feat：接入mongodb

2026-09-30 04:27:35 -07:00
parent f92909ccda
commit 16212d7138
11 changed files with 663 additions and 43 deletions
@@ -1,4 +1,4 @@
-"""会话存储：创建时即占位，用于校验 ID 由后端签发；当前为进程内存实现"""
+"""会话存储：创建时即占位，用于校验 ID 由后端签发；此处为协议和进程内存实现，MongoDB 实现见 conversation_mongo"""

 from datetime import UTC, datetime, timedelta
 from typing import Protocol
@@ -9,7 +9,7 @@ from app.schemas.conversation import Turn


 class ConversationStore(Protocol):
-    """后续接入 MongoDB 时实现同一组方法，业务代码不需要改动"""
+    """内存和 MongoDB 实现同一组方法，业务代码不需要关心具体存储"""

    async def create(self) -> str: ...

@@ -0,0 +1,116 @@
+"""会话存储的 MongoDB 实现，方法语义与内存实现一致"""
+
+from datetime import UTC, datetime, timedelta
+from uuid import uuid4
+
+from bson.codec_options import CodecOptions
+from pymongo import ASCENDING, DESCENDING
+from pymongo.asynchronous.database import AsyncDatabase
+
+from app.config import settings
+from app.schemas.conversation import Turn
+
+# 单次清理的空会话上限，避免积压时一次请求扫描过多
+PURGE_BATCH = 200
+
+
+def _now() -> datetime:
+    return datetime.now(UTC)
+
+
+class MongoConversationStore:
+    def __init__(
+        self, db: AsyncDatabase, empty_ttl: int | None = None, prefix: str = ""
+    ):
+        ttl = empty_ttl if empty_ttl is not None else settings.conversation_empty_ttl
+        self._empty_ttl = timedelta(seconds=ttl)
+        # 读回的时间统一带 UTC 时区，与 Turn 的时间保持一致；集合名前缀用于测试隔离
+        options = CodecOptions(tz_aware=True, tzinfo=UTC)
+        self._conversations = db.get_collection(
+            f"{prefix}conversations", codec_options=options
+        )
+        self._turns = db.get_collection(f"{prefix}turns", codec_options=options)
+
+    async def ensure_indexes(self) -> None:
+        # 创建索引是幂等的，启动时调用一次即可
+        # 同一任务只有一轮，恢复后的重复写入靠唯一索引落到同一条记录
+        await self._turns.create_index(
+            [("conversation_id", ASCENDING), ("thread_id", ASCENDING)], unique=True
+        )
+        await self._turns.create_index(
+            [("conversation_id", ASCENDING), ("created_at", ASCENDING)]
+        )
+        # 空会话清理按创建时间扫描
+        await self._conversations.create_index("created_at")
+
+    async def _purge_expired(self) -> None:
+        # 只清理没有任何轮次且超过保留时长的空会话，有对话的会话不受影响
+        deadline = _now() - self._empty_ttl
+        expired = await (
+            self._conversations.find({"created_at": {"$lt": deadline}}, {"_id": 1})
+            .limit(PURGE_BATCH)
+            .to_list()
+        )
+        ids = [item["_id"] for item in expired]
+        if not ids:
+            return
+        used = set(
+            await self._turns.distinct(
+                "conversation_id", {"conversation_id": {"$in": ids}}
+            )
+        )
+        empty = [item for item in ids if item not in used]
+        if empty:
+            await self._conversations.delete_many({"_id": {"$in": empty}})
+
+    async def create(self) -> str:
+        await self._purge_expired()
+        conversation_id = str(uuid4())
+        # user_id 为后续长期记忆预留，当前恒为空
+        await self._conversations.insert_one(
+            {"_id": conversation_id, "created_at": _now(), "user_id": None}
+        )
+        return conversation_id
+
+    async def exists(self, conversation_id: str) -> bool:
+        document = await self._conversations.find_one({"_id": conversation_id})
+        if document is None:
+            return False
+        if document["created_at"] >= _now() - self._empty_ttl:
+            return True
+        # 超过空会话保留时长后，只有已产生对话的会话仍视为存在
+        return (
+            await self._turns.find_one({"conversation_id": conversation_id}, {"_id": 1})
+            is not None
+        )
+
+    async def load_turns(self, conversation_id: str, limit: int | None) -> list[Turn]:
+        if limit == 0:
+            return []
+        cursor = self._turns.find({"conversation_id": conversation_id})
+        if limit is None:
+            documents = await cursor.sort(
+                [("created_at", ASCENDING), ("_id", ASCENDING)]
+            ).to_list()
+        else:
+            # 倒序取最近 limit 轮，再反转为时间正序
+            documents = await (
+                cursor.sort([("created_at", DESCENDING), ("_id", DESCENDING)])
+                .limit(limit)
+                .to_list()
+            )
+            documents.reverse()
+        return [
+            Turn.model_validate({k: v for k, v in item.items() if k != "_id"})
+            for item in documents
+        ]
+
+    async def save_turn(self, turn: Turn) -> None:
+        data = turn.model_dump()
+        created_at = data.pop("created_at")
+        # 同一任务再次写入（澄清恢复后）保留首次创建时间，只更新内容
+        await self._turns.update_one(
+            {"conversation_id": turn.conversation_id, "thread_id": turn.thread_id},
+            {"$set": data, "$setOnInsert": {"created_at": created_at}},
+            upsert=True,
+        )
@@ -0,0 +1,70 @@
+"""持久化装配：按配置选择 MongoDB 或进程内存，统一返回会话存储、检查点和关闭函数"""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+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]]
+
+
+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)