ZNJJ-api-server/src/core/fastgpt_client.py

"""FastGPT client dependency injection."""
from contextlib import AsyncExitStack, asynccontextmanager
from dataclasses import dataclass
from fastapi import FastAPI
from fastgpt_client import AsyncChatClient
from .config import Config


@dataclass(frozen=True)
class FastGPTProfile:
    """FastGPT client and app metadata for one client mode."""

    client: AsyncChatClient
    app_id: str


_fastgpt_profiles: dict[str, FastGPTProfile] = {}


def _profile_settings() -> dict[str, dict[str, str | None]]:
    return {
        "direct": {
            "api_key": Config.FASTGPT_DIRECT_API_KEY,
            "base_url": Config.FASTGPT_DIRECT_BASE_URL,
            "app_id": Config.FASTGPT_DIRECT_APP_ID,
        },
        "browser_addon": {
            "api_key": Config.FASTGPT_BROWSER_ADDON_API_KEY,
            "base_url": Config.FASTGPT_BROWSER_ADDON_BASE_URL,
            "app_id": Config.FASTGPT_BROWSER_ADDON_APP_ID,
        },
    }


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage all mode-specific FastGPT clients."""
    Config.validate()

    async with AsyncExitStack() as stack:
        for client_mode, settings in _profile_settings().items():
            client = AsyncChatClient(
                api_key=settings["api_key"],
                base_url=settings["base_url"],
                timeout=60.0,
                max_retries=3,
                retry_delay=1.0,
                enable_logging=Config.DEBUG,
            )
            await stack.enter_async_context(client)
            _fastgpt_profiles[client_mode] = FastGPTProfile(
                client=client,
                app_id=settings["app_id"],
            )

        try:
            yield
        finally:
            _fastgpt_profiles.clear()


def get_fastgpt_profile(client_mode: str) -> FastGPTProfile:
    """Get the initialized FastGPT profile for a client mode."""
    try:
        return _fastgpt_profiles[client_mode]
    except KeyError as exc:
        if client_mode not in _profile_settings():
            raise ValueError(f"Unsupported clientMode: {client_mode}") from exc
        raise RuntimeError(
            f"FastGPT profile is not initialized for clientMode={client_mode}"
        ) from exc


def get_fastgpt_client(client_mode: str = "direct") -> AsyncChatClient:
    """Get a FastGPT client, defaulting to the legacy direct mode."""
    return get_fastgpt_profile(client_mode).client