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

"""FastGPT client and chat backend dependency injection."""
from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastgpt_client import AsyncChatClient
from ..backends.chat import ChatBackend
from ..backends.fastgpt import FastGPTBackend
from .config import Config


_fastgpt_client: AsyncChatClient | None = None
_chat_backend: ChatBackend | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage FastGPT client lifecycle."""
    global _chat_backend, _fastgpt_client
    Config.validate()

    _fastgpt_client = AsyncChatClient(
        api_key=Config.FASTGPT_API_KEY,
        base_url=Config.FASTGPT_BASE_URL,
        timeout=60.0,
        max_retries=3,
        retry_delay=1.0,
        enable_logging=Config.DEBUG,
    )
    await _fastgpt_client.__aenter__()
    _chat_backend = FastGPTBackend(_fastgpt_client)

    try:
        yield
    finally:
        _chat_backend = None
        if _fastgpt_client:
            await _fastgpt_client.__aexit__(None, None, None)
        _fastgpt_client = None


def get_fastgpt_client() -> AsyncChatClient:
    """Get the initialized FastGPT client."""
    if _fastgpt_client is None:
        raise RuntimeError("FastGPT client not initialized")
    return _fastgpt_client


def get_chat_backend() -> ChatBackend:
    """Get the backend-neutral chat service."""
    if _chat_backend is None:
        raise RuntimeError("Chat backend not initialized")
    return _chat_backend