55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import asyncio
|
||
from typing import Optional
|
||
|
||
from config.env import RAGFlowConfig
|
||
from utils.ragflow_asy_util import AsyncRAGFlowClient
|
||
|
||
|
||
class RAGFlowClientManager:
|
||
"""管理RAGFlow异步客户端的生命周期，复用底层HTTP会话"""
|
||
|
||
_client: Optional[AsyncRAGFlowClient] = None
|
||
_lock: asyncio.Lock = asyncio.Lock()
|
||
|
||
@classmethod
|
||
async def init_client(cls) -> AsyncRAGFlowClient:
|
||
async with cls._lock:
|
||
if cls._client is None:
|
||
cls._client = AsyncRAGFlowClient(
|
||
RAGFlowConfig.RAGFLOW_BASE_URL,
|
||
RAGFlowConfig.RAGFLOW_API_KEY,
|
||
timeout=60,
|
||
connector_kwargs={
|
||
"limit": 128,
|
||
"ssl": False,
|
||
"keepalive_timeout": 45,
|
||
},
|
||
)
|
||
await cls._client.create_session()
|
||
return cls._client
|
||
|
||
@classmethod
|
||
async def get_client(cls) -> AsyncRAGFlowClient:
|
||
if cls._client is None:
|
||
return await cls.init_client()
|
||
return cls._client
|
||
|
||
@classmethod
|
||
async def shutdown_client(cls) -> None:
|
||
async with cls._lock:
|
||
if cls._client is not None:
|
||
await cls._client.close_session()
|
||
cls._client = None
|
||
|
||
|
||
async def init_ragflow_client() -> AsyncRAGFlowClient:
|
||
return await RAGFlowClientManager.init_client()
|
||
|
||
|
||
async def get_ragflow_client() -> AsyncRAGFlowClient:
|
||
return await RAGFlowClientManager.get_client()
|
||
|
||
|
||
async def shutdown_ragflow_client() -> None:
|
||
await RAGFlowClientManager.shutdown_client()
|