planet/backend/app/services/ai_client.py

from __future__ import annotations

import asyncio
import json

import httpx
from fastapi import Depends, HTTPException, status
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import settings
from app.db.session import get_db
from app.schemas.ai import (
    AIProviderStatusResponse,
    SituationalAnalysisRequest,
    SituationalAnalysisResponse,
)


class AIProviderClient:
    def __init__(
        self,
        *,
        service_url: str | None = None,
        service_token: str | None = None,
        timeout: int | None = None,
        retry_attempts: int | None = None,
        llm_config: dict | None = None,
    ) -> None:
        self.service_url = (
            service_url if service_url is not None else settings.AI_PROVIDER_SERVICE_URL
        ).rstrip("/")
        self.service_token = (
            service_token if service_token is not None else settings.AI_PROVIDER_SERVICE_TOKEN
        )
        self.timeout = timeout if timeout is not None else settings.AI_PROVIDER_TIMEOUT_SECONDS
        self.retry_attempts = max(
            retry_attempts if retry_attempts is not None else settings.AI_PROVIDER_RETRY_ATTEMPTS,
            1,
        )
        self.llm_config = llm_config or {}

    def _headers(self, request_id: str | None = None) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        if self.service_token:
            headers["X-Provider-Token"] = self.service_token
        if request_id:
            headers["X-Request-ID"] = request_id
        llm_header_map = {
            "provider": "X-AI-Provider",
            "provider_api": "X-AI-Provider-API",
            "base_url": "X-AI-Base-URL",
            "api_key": "X-AI-API-Key",
            "model": "X-AI-Model",
            "max_tokens": "X-AI-Max-Tokens",
            "anthropic_version": "X-AI-Anthropic-Version",
        }
        for key, header_name in llm_header_map.items():
            value = self.llm_config.get(key)
            if value not in (None, ""):
                headers[header_name] = str(value)
        model_provider_apis = self.llm_config.get("model_provider_apis")
        if isinstance(model_provider_apis, dict) and model_provider_apis:
            headers["X-AI-Model-Provider-APIs"] = json.dumps(model_provider_apis)
        return headers

    async def get_status(self, request_id: str | None = None) -> AIProviderStatusResponse:
        if not self.service_url:
            return AIProviderStatusResponse(
                provider="unconfigured",
                enabled=False,
                configured=False,
                model=None,
                base_url=None,
            )

        data = await self._request("GET", "/v1/provider/status", request_id=request_id)
        return AIProviderStatusResponse.model_validate(data)

    async def analyze(
        self,
        payload: SituationalAnalysisRequest,
        request_id: str | None = None,
    ) -> SituationalAnalysisResponse:
        if not self.service_url:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="AI provider service URL is not configured.",
            )

        data = await self._request(
            "POST",
            "/v1/analyze",
            json=payload.model_dump(),
            request_id=request_id,
        )
        return SituationalAnalysisResponse.model_validate(data)

    async def _request(
        self,
        method: str,
        path: str,
        json: dict | None = None,
        request_id: str | None = None,
    ) -> dict:
        last_error: Exception | None = None
        for attempt in range(1, self.retry_attempts + 1):
            try:
                async with httpx.AsyncClient(timeout=self.timeout) as client:
                    response = await client.request(
                        method,
                        f"{self.service_url}{path}",
                        headers=self._headers(request_id),
                        json=json,
                    )
                    response.raise_for_status()
                    return response.json()
            except httpx.HTTPStatusError as exc:
                last_error = exc
                if attempt < self.retry_attempts and exc.response.status_code >= 500:
                    await asyncio.sleep(0.3 * attempt)
                    continue
                detail = exc.response.text or "AI provider service returned an error"
                raise HTTPException(
                    status_code=status.HTTP_502_BAD_GATEWAY,
                    detail=f"AI provider service request failed: {detail}",
                ) from exc
            except httpx.HTTPError as exc:
                last_error = exc
                if attempt < self.retry_attempts:
                    await asyncio.sleep(0.3 * attempt)
                    continue
                raise HTTPException(
                    status_code=status.HTTP_502_BAD_GATEWAY,
                    detail=f"Failed to reach AI provider service: {exc}",
                ) from exc

        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail=f"AI provider service request failed: {last_error}",
        )


async def get_ai_provider_client(db: AsyncSession = Depends(get_db)) -> AIProviderClient:
    from app.api.v1.settings import get_runtime_ai_provider_config

    runtime_config = await get_runtime_ai_provider_config(db)
    return AIProviderClient(
        service_url=runtime_config["service_url"],
        service_token=runtime_config["service_token"],
        timeout=runtime_config["timeout_seconds"],
        retry_attempts=runtime_config["retry_attempts"],
        llm_config=runtime_config.get("llm_config") or {},
    )