199 lines
7.8 KiB
Python
199 lines
7.8 KiB
Python
"""会话存储的 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 ConversationSummary, Turn
|
|
from app.services.conversation import make_title
|
|
|
|
# 单次清理的空会话上限，避免积压时一次请求扫描过多
|
|
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")
|
|
# 按用户查询会话
|
|
await self._conversations.create_index("user_id")
|
|
# 按大屏过滤会话
|
|
await self._conversations.create_index([("user_id", ASCENDING), ("screen_id", ASCENDING)])
|
|
|
|
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, user_id: int, screen_id: str | None = None) -> str:
|
|
await self._purge_expired()
|
|
conversation_id = str(uuid4())
|
|
# user_id 来自用户中心，会话只对创建者可见
|
|
await self._conversations.insert_one(
|
|
{
|
|
"_id": conversation_id,
|
|
"created_at": _now(),
|
|
"user_id": user_id,
|
|
"screen_id": screen_id,
|
|
}
|
|
)
|
|
return conversation_id
|
|
|
|
async def exists(self, conversation_id: str, user_id: int) -> bool:
|
|
# 不属于该用户的会话与不存在同样处理，不泄露会话是否存在
|
|
document = await self._conversations.find_one(
|
|
{"_id": conversation_id, "user_id": user_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,
|
|
)
|
|
|
|
def _with_stats(self, user_id: int, screen_id: str | None) -> list[dict]:
|
|
match: dict = {"user_id": user_id}
|
|
if screen_id is not None:
|
|
match["screen_id"] = screen_id
|
|
# $unwind 会丢弃没有任何轮次的空会话
|
|
return [
|
|
{"$match": match},
|
|
{
|
|
"$lookup": {
|
|
"from": self._turns.name,
|
|
"localField": "_id",
|
|
"foreignField": "conversation_id",
|
|
"pipeline": [
|
|
{"$sort": {"created_at": 1, "_id": 1}},
|
|
{
|
|
"$group": {
|
|
"_id": None,
|
|
"first_prompt": {"$first": "$prompt"},
|
|
"updated_at": {"$last": "$created_at"},
|
|
}
|
|
},
|
|
],
|
|
"as": "stats",
|
|
}
|
|
},
|
|
{"$unwind": "$stats"},
|
|
]
|
|
|
|
async def list_by_user(
|
|
self, user_id: int, limit: int, offset: int, screen_id: str | None = None
|
|
) -> list[ConversationSummary]:
|
|
pipeline = [
|
|
*self._with_stats(user_id, screen_id),
|
|
{"$sort": {"stats.updated_at": -1, "_id": -1}},
|
|
{"$skip": offset},
|
|
{"$limit": limit},
|
|
]
|
|
cursor = await self._conversations.aggregate(pipeline)
|
|
return [
|
|
ConversationSummary(
|
|
conversation_id=item["_id"],
|
|
title=item.get("title") or make_title(item["stats"]["first_prompt"]),
|
|
screen_id=item.get("screen_id"),
|
|
updated_at=item["stats"]["updated_at"],
|
|
)
|
|
for item in await cursor.to_list()
|
|
]
|
|
|
|
async def count_by_user(self, user_id: int, screen_id: str | None = None) -> int:
|
|
pipeline = [*self._with_stats(user_id, screen_id), {"$count": "total"}]
|
|
cursor = await self._conversations.aggregate(pipeline)
|
|
counted = await cursor.to_list()
|
|
return counted[0]["total"] if counted else 0
|
|
|
|
async def rename(self, conversation_id: str, user_id: int, title: str) -> bool:
|
|
if not await self.exists(conversation_id, user_id):
|
|
return False
|
|
await self._conversations.update_one(
|
|
{"_id": conversation_id, "user_id": user_id}, {"$set": {"title": title}}
|
|
)
|
|
return True
|
|
|
|
async def delete(self, conversation_id: str, user_id: int) -> bool:
|
|
if not await self.exists(conversation_id, user_id):
|
|
return False
|
|
await self._turns.delete_many({"conversation_id": conversation_id})
|
|
await self._conversations.delete_one({"_id": conversation_id, "user_id": user_id})
|
|
return True
|