74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
"""持久化装配：按配置选择 MongoDB 或进程内存，统一返回会话存储、检查点和关闭函数"""
|
||
|
||
from collections.abc import Awaitable, Callable
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||
from loguru import logger
|
||
|
||
from app.config import settings
|
||
from app.services.conversation import ConversationStore, InMemoryConversationStore
|
||
|
||
|
||
@dataclass
|
||
class Persistence:
|
||
store: ConversationStore
|
||
# 为空时由图使用默认的内存检查点
|
||
checkpointer: BaseCheckpointSaver | None
|
||
close: Callable[[], Awaitable[None]]
|
||
# 启用 MongoDB 时的数据库，供物料等其他集合复用同一连接
|
||
db: Any = None
|
||
|
||
|
||
async def _close_nothing() -> None:
|
||
return None
|
||
|
||
|
||
async def open_persistence() -> Persistence:
|
||
if not settings.mongodb_uri:
|
||
logger.warning(
|
||
"未配置 MONGODB_URI，会话记录和检查点仅保存在进程内存，重启后丢失"
|
||
)
|
||
return Persistence(InMemoryConversationStore(), None, _close_nothing)
|
||
|
||
# 依赖较重且只在启用 MongoDB 时使用，放到函数内导入
|
||
from langgraph.checkpoint.mongodb import MongoDBSaver
|
||
from pymongo import AsyncMongoClient, MongoClient
|
||
|
||
from app.processor.graph import build_checkpoint_serde
|
||
from app.services.conversation_mongo import MongoConversationStore
|
||
|
||
# 会话记录使用异步客户端；官方检查点基于同步客户端，异步方法由其放入线程池执行
|
||
async_client = AsyncMongoClient(settings.mongodb_uri, serverSelectionTimeoutMS=5000)
|
||
sync_client = MongoClient(settings.mongodb_uri, serverSelectionTimeoutMS=5000)
|
||
|
||
async def close() -> None:
|
||
try:
|
||
await async_client.close()
|
||
finally:
|
||
sync_client.close()
|
||
|
||
try:
|
||
# 启动时探测连接，配置错误尽早暴露；异常信息不包含连接串，避免泄露账号密码
|
||
try:
|
||
await async_client.admin.command("ping")
|
||
except Exception as exc:
|
||
raise RuntimeError(f"无法连接 MongoDB：{type(exc).__name__}") from None
|
||
|
||
store = MongoConversationStore(async_client[settings.mongodb_db])
|
||
await store.ensure_indexes()
|
||
# 与内存检查点共用同一份类型白名单，保证读回后仍是业务模型对象
|
||
checkpointer = MongoDBSaver(
|
||
sync_client,
|
||
settings.mongodb_db,
|
||
ttl=settings.mongodb_checkpoint_ttl or None,
|
||
serde=build_checkpoint_serde(),
|
||
)
|
||
except BaseException:
|
||
await close()
|
||
raise
|
||
|
||
logger.info("已启用 MongoDB 持久化，库名 {}", settings.mongodb_db)
|
||
return Persistence(store, checkpointer, close, async_client[settings.mongodb_db])
|