119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
"""会话存储：创建时即占位，用于校验 ID 由后端签发；此处为协议和进程内存实现，MongoDB 实现见 conversation_mongo"""
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Protocol
|
|
from uuid import uuid4
|
|
|
|
from app.config import settings
|
|
from app.schemas.conversation import ConversationSummary, Turn
|
|
|
|
TITLE_MAX_LENGTH = 30
|
|
|
|
|
|
def make_title(prompt: str) -> str:
|
|
"""会话标题取首轮用户输入，合并空白后截断"""
|
|
title = " ".join(prompt.split())
|
|
if len(title) <= TITLE_MAX_LENGTH:
|
|
return title
|
|
return title[:TITLE_MAX_LENGTH] + "…"
|
|
|
|
|
|
class ConversationStore(Protocol):
|
|
"""内存和 MongoDB 实现同一组方法，业务代码不需要关心具体存储"""
|
|
|
|
async def create(self, user_id: int) -> str: ...
|
|
|
|
async def exists(self, conversation_id: str, user_id: int) -> bool:
|
|
"""会话已签发、未过期且属于该用户；不属于时与不存在同样返回 False"""
|
|
...
|
|
|
|
async def load_turns(
|
|
self, conversation_id: str, limit: int | None
|
|
) -> list[Turn]: ...
|
|
|
|
async def save_turn(self, turn: Turn) -> None: ...
|
|
|
|
async def list_by_user(
|
|
self, user_id: int, limit: int, offset: int
|
|
) -> list[ConversationSummary]:
|
|
"""该用户已产生对话的会话，按最近更新时间倒序；没有对话的空会话不返回"""
|
|
...
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
class InMemoryConversationStore:
|
|
def __init__(self, empty_ttl: int | None = None):
|
|
ttl = empty_ttl if empty_ttl is not None else settings.conversation_empty_ttl
|
|
self._empty_ttl = timedelta(seconds=ttl)
|
|
# 会话 ID -> 创建时间；轮次按 thread_id 去重并保持首次写入顺序
|
|
self._created: dict[str, datetime] = {}
|
|
self._owner: dict[str, int] = {}
|
|
self._turns: dict[str, dict[str, Turn]] = {}
|
|
|
|
def _purge_expired(self) -> None:
|
|
# 只清理没有任何轮次且超过保留时长的空会话，有对话的会话不受影响
|
|
deadline = _now() - self._empty_ttl
|
|
for conversation_id in [
|
|
key
|
|
for key, created in self._created.items()
|
|
if created < deadline and not self._turns.get(key)
|
|
]:
|
|
del self._created[conversation_id]
|
|
self._owner.pop(conversation_id, None)
|
|
self._turns.pop(conversation_id, None)
|
|
|
|
async def create(self, user_id: int) -> str:
|
|
self._purge_expired()
|
|
conversation_id = str(uuid4())
|
|
self._created[conversation_id] = _now()
|
|
self._owner[conversation_id] = user_id
|
|
return conversation_id
|
|
|
|
async def exists(self, conversation_id: str, user_id: int) -> bool:
|
|
self._purge_expired()
|
|
return (
|
|
conversation_id in self._created
|
|
and self._owner.get(conversation_id) == user_id
|
|
)
|
|
|
|
async def load_turns(self, conversation_id: str, limit: int | None) -> list[Turn]:
|
|
turns = list(self._turns.get(conversation_id, {}).values())
|
|
turns.sort(key=lambda turn: turn.created_at)
|
|
if limit is None:
|
|
return turns
|
|
# 只取最近 limit 轮，limit 为 0 时不带历史
|
|
return turns[-limit:] if limit else []
|
|
|
|
async def save_turn(self, turn: Turn) -> None:
|
|
# 同一任务再次写入（澄清恢复后）保留首次创建时间，只更新内容
|
|
existing = self._turns.setdefault(turn.conversation_id, {}).get(turn.thread_id)
|
|
if existing is not None:
|
|
turn = turn.model_copy(update={"created_at": existing.created_at})
|
|
self._turns[turn.conversation_id][turn.thread_id] = turn
|
|
|
|
async def list_by_user(
|
|
self, user_id: int, limit: int, offset: int
|
|
) -> list[ConversationSummary]:
|
|
summaries = []
|
|
for conversation_id, owner in self._owner.items():
|
|
turns = sorted(
|
|
self._turns.get(conversation_id, {}).values(),
|
|
key=lambda turn: turn.created_at,
|
|
)
|
|
if owner != user_id or not turns:
|
|
continue
|
|
summaries.append(
|
|
ConversationSummary(
|
|
conversation_id=conversation_id,
|
|
title=make_title(turns[0].prompt),
|
|
updated_at=turns[-1].created_at,
|
|
)
|
|
)
|
|
summaries.sort(
|
|
key=lambda item: (item.updated_at, item.conversation_id), reverse=True
|
|
)
|
|
return summaries[offset : offset + limit]
|