feat：接入mongodb

2026-09-30 04:27:35 -07:00
parent f92909ccda
commit 16212d7138
11 changed files with 663 additions and 43 deletions
@@ -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,
+        )