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

"""后端与单一会话存储的生命周期、依赖注入。"""
from contextlib import AsyncExitStack, asynccontextmanager

from fastapi import FastAPI
from fastgpt_client import AsyncChatClient

from ..agent.graph import build_graph
from ..backends.chat import ChatBackend
from ..backends.fastgpt import FastGPTBackend
from ..backends.langgraph import LangGraphBackend
from ..services.app_service import AccidentAppService
from ..storage.checkpoints import delete_checkpoint
from .config import Config

_fastgpt_client: AsyncChatClient | None = None
_chat_backend: ChatBackend | None = None
_langgraph_graph = None
_app_service: AccidentAppService | None = None

FASTGPT_TIMEOUT_SECONDS = 60.0
FASTGPT_MAX_RETRIES = 3


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _chat_backend, _fastgpt_client, _langgraph_graph, _app_service
    Config.validate()
    llm = None
    async with AsyncExitStack() as resources:
        try:
            if Config.AGENT_BACKEND == 'langgraph':
                from ..agent.llm import build_llm_service
                from ..storage.runtime import StorageSettings, open_storage

                storage = await resources.enter_async_context(open_storage(StorageSettings.from_config(Config)))
                _langgraph_graph = build_graph(storage.checkpointer)
                llm = build_llm_service()
                _chat_backend = LangGraphBackend(_langgraph_graph, llm=llm)
                _app_service = AccidentAppService(_chat_backend, storage.repository, locks=storage.locks)
            else:
                _fastgpt_client = AsyncChatClient(
                    api_key=Config.FASTGPT_API_KEY, base_url=Config.FASTGPT_BASE_URL,
                    timeout=FASTGPT_TIMEOUT_SECONDS, max_retries=FASTGPT_MAX_RETRIES + 1,
                    retry_delay=1.0, enable_logging=Config.DEBUG,
                )
                await _fastgpt_client.__aenter__()
                _chat_backend = FastGPTBackend(_fastgpt_client)
            yield
        finally:
            _chat_backend = None
            _langgraph_graph = None
            _app_service = None
            if llm is not None and hasattr(llm, 'aclose'):
                await llm.aclose()
            if _fastgpt_client:
                await _fastgpt_client.__aexit__(None, None, None)
            _fastgpt_client = None


async def delete_chat_session(thread_id: str) -> None:
    await delete_checkpoint(_langgraph_graph, thread_id)


def get_fastgpt_client() -> AsyncChatClient:
    if _fastgpt_client is None:
        raise RuntimeError('FastGPT client not initialized')
    return _fastgpt_client


def get_app_service() -> AccidentAppService:
    if _app_service is None:
        raise RuntimeError('Application service not initialized')
    return _app_service


def get_chat_backend() -> ChatBackend:
    if _chat_backend is None:
        raise RuntimeError('Chat backend not initialized')
    return _chat_backend