""" 对话记忆管理器 """ import asyncio from datetime import datetime from typing import Dict, List, Optional from langchain_core.messages import BaseMessage, messages_from_dict, messages_to_dict from pydantic import BaseModel, ConfigDict, Field, field_serializer from app.agent.history.message import legacy_identity from app.agent.tools.result import messages_for_persistence from app.application.messaging.chat import ( AgentChatPersistenceService, AgentChatRecord, AgentChatService, get_configured_agent_chat_persistence, get_configured_agent_chat_service, ) from app.runtime.log import logger from app.runtime.settings import get_runtime_setting class ConversationMemory(BaseModel): """ 对话记忆模型 消息字段依赖 langchain_core,因此模型归属智能体记忆模块;普通 schema 导入不加载 LangChain。 """ session_id: str = Field(description="会话ID") user_id: Optional[str] = Field(default=None, description="用户ID") owner_user_id: Optional[str] = Field(default=None, description="持久化会话原始归属,缓存复用时再次校验") owner_username: Optional[str] = Field(default=None, description="无 ID 旧记录的用户名归属") messages: List[BaseMessage] = Field(default_factory=list, description="消息列表") updated_at: datetime = Field(default_factory=datetime.now, description="更新时间") model_config = ConfigDict() @field_serializer("updated_at", when_used="json") def serialize_datetime(self, value: datetime) -> str: """ 把记忆更新时间序列化为标准日期字符串。 """ return value.isoformat() class MemoryManager: """ 对话记忆管理器 """ def __init__( self, chat: AgentChatService | None = None, persistence: AgentChatPersistenceService | None = None, ) -> None: """ 创建记忆缓存,并保存组合根显式注入的会话能力。 """ self._chat = chat self._persistence = persistence # 内存中的会话记忆缓存 self.memory_cache: Dict[str, ConversationMemory] = {} # 内存缓存清理任务 self.cleanup_task: Optional[asyncio.Task[None]] = None def configure( self, chat: AgentChatService, persistence: AgentChatPersistenceService, ) -> None: """ 在 Agent 启动前绑定唯一会话查询与写入服务。 """ self._chat = chat self._persistence = persistence def _chat_service(self) -> AgentChatService: """ 返回显式注入服务,兼容测试未装配时使用既有应用服务。 """ return self._chat or get_configured_agent_chat_service() def _chat_persistence(self) -> AgentChatPersistenceService: """ 返回显式注入写服务,兼容测试未装配时使用既有应用服务。 """ return self._persistence or get_configured_agent_chat_persistence() def initialize(self) -> None: """ 初始化记忆管理器 """ try: if self.cleanup_task and not self.cleanup_task.done(): return # 启动内存缓存清理任务(Redis通过TTL自动过期) self.cleanup_task = asyncio.create_task( self._cleanup_expired_memories() ) logger.info("对话记忆管理器初始化完成") except Exception as e: logger.warning(f"Redis连接失败,将使用内存存储: {e}") async def close(self) -> None: """ 关闭记忆管理器 """ if self.cleanup_task: self.cleanup_task.cancel() try: await self.cleanup_task except asyncio.CancelledError: pass self.cleanup_task = None logger.info("对话记忆管理器已关闭") @staticmethod def _get_memory_key(session_id: str, user_id: Optional[str]) -> str: """ 计算内存Key """ return f"{user_id}:{session_id}" if user_id else session_id def get_memory(self, session_id: str, user_id: str) -> Optional[ConversationMemory]: """ 获取内存中的记忆 """ cache_key = self._get_memory_key(session_id, user_id) return self.memory_cache.get(cache_key) @staticmethod def _chat_lookup_params( session_id: str, user_id: str, ) -> tuple[dict[str, str], ...]: """ 返回用户专属会话优先、旧无用户会话兜底的查询顺序。 """ return ( {"session_id": session_id, "user_id": user_id}, {"session_id": session_id}, ) def _restore_agent_messages( self, *, session_id: str, user_id: str, chat: Optional[AgentChatRecord], allow_cross_owner: bool = False, legacy_username: Optional[str] = None, ) -> List[BaseMessage]: """ 仅恢复当前用户或无归属旧记录,再清理工具图像并回填缓存。 """ if not chat or not chat.agent_messages: return [] # 无用户条件的兼容查询只用于无归属旧记录,不能扩大到其他明确归属账号。 owner_id = getattr(chat, "user_id", None) owner_username = getattr(chat, "username", None) foreign_owner = ( str(owner_id) != str(user_id) if owner_id else bool(owner_username and owner_username != legacy_username) ) if foreign_owner and not allow_cross_owner: return [] try: messages = messages_for_persistence(messages_from_dict(chat.agent_messages)) for position, message in enumerate(messages): if not message.id: message.id = legacy_identity(session_id, position, message) except Exception as e: logger.debug(f"恢复持久化Agent消息失败: {e}") return [] memory = ConversationMemory( session_id=session_id, user_id=user_id, owner_user_id=owner_id, owner_username=owner_username, messages=messages, ) self.save_memory(memory) return memory.messages def get_agent_messages( self, session_id: str, user_id: str, *, allow_cross_owner: bool = False, legacy_username: Optional[str] = None, ) -> List[BaseMessage]: """ 为Agent获取最近的消息。 优先使用内存缓存,缓存不存在时从数据库恢复上一轮持久化的原始 messages。 """ memory = self.get_memory(session_id, user_id) if memory: return self._authorized_cached_messages(memory, user_id, allow_cross_owner, legacy_username) try: service = self._chat_service() chat = None for lookup_params in self._chat_lookup_params(session_id, user_id): chat = service.get_sync(**lookup_params) if chat: break except Exception as e: logger.debug(f"读取持久化Agent会话失败: {e}") return [] return self._restore_agent_messages( session_id=session_id, user_id=user_id, chat=chat, allow_cross_owner=allow_cross_owner, legacy_username=legacy_username, ) @staticmethod def _authorized_cached_messages( memory: ConversationMemory, user_id: str, allow_cross_owner: bool, legacy_username: Optional[str], ) -> List[BaseMessage]: """ 恢复授权不会永久附着于缓存,每次跨归属读取都需要显式许可。 """ foreign_owner = ( memory.owner_user_id != str(user_id) if memory.owner_user_id else bool(memory.owner_username and memory.owner_username != legacy_username) ) if foreign_owner and not allow_cross_owner: return [] return memory.messages async def async_get_agent_messages( self, session_id: str, user_id: str, *, allow_cross_owner: bool = False, legacy_username: Optional[str] = None, ) -> List[BaseMessage]: """ 异步恢复 Agent 消息,查询与会话应用服务保持同一异步端口。 """ memory = self.get_memory(session_id, user_id) if memory: return self._authorized_cached_messages(memory, user_id, allow_cross_owner, legacy_username) try: service = self._chat_service() chat = None for lookup_params in self._chat_lookup_params(session_id, user_id): chat = await service.get(**lookup_params) if chat: break except Exception as e: logger.debug(f"读取持久化Agent会话失败: {e}") return [] return self._restore_agent_messages( session_id=session_id, user_id=user_id, chat=chat, allow_cross_owner=allow_cross_owner, legacy_username=legacy_username, ) def _update_agent_messages( self, *, session_id: str, user_id: str, messages: List[BaseMessage], ) -> None: """ 统一更新同步和异步写路径共享的内存消息状态。 """ memory = self.get_memory(session_id, user_id) if not memory: memory = ConversationMemory(session_id=session_id, user_id=user_id) memory.messages = messages memory.updated_at = datetime.now() self.save_memory(memory) def save_agent_messages( self, session_id: str, user_id: str, messages: List[BaseMessage] ) -> None: """ 当前会话在缓存保留原图,持久化副本只保存工具图片的失效说明。 """ self._update_agent_messages( session_id=session_id, user_id=user_id, messages=messages, ) try: self._chat_service().save_agent_messages( session_id=session_id, user_id=user_id, messages=messages_to_dict(messages_for_persistence(messages)), ) except Exception as e: logger.debug(f"持久化Agent消息失败: {e}") async def async_save_agent_messages( self, session_id: str, user_id: str, messages: List[BaseMessage] ) -> None: """ 缓存保留原图,去除工具图片的副本经有界数据库 worker 持久化。 """ self._update_agent_messages( session_id=session_id, user_id=user_id, messages=messages, ) try: persistence = self._chat_persistence() await persistence.async_save_agent_messages( session_id=session_id, user_id=user_id, messages=messages_to_dict(messages_for_persistence(messages)), ) except Exception as e: logger.debug(f"持久化Agent消息失败: {e}") def save_memory(self, memory: ConversationMemory) -> None: """ 保存记忆到内存缓存 注意:Redis中的记忆通过TTL机制自动过期,这里只更新内存缓存,Redis会在下次访问时自动过期 """ cache_key = self._get_memory_key(memory.session_id, memory.user_id) self.memory_cache[cache_key] = memory def clear_memory(self, session_id: str, user_id: str) -> None: """ 清空会话记忆 """ cache_key = self._get_memory_key(session_id, user_id) if cache_key in self.memory_cache: del self.memory_cache[cache_key] logger.info(f"会话记忆已清空: session_id={session_id}, user_id={user_id}") async def _cleanup_expired_memories(self) -> None: """ 清理内存中过期记忆的后台任务 注意:Redis中的记忆通过TTL机制自动过期,这里只清理内存缓存 """ while True: try: # 每小时清理一次 await asyncio.sleep(3600) current_time = datetime.now() expired_sessions = [] # 只检查内存缓存中的过期记忆 # Redis中的记忆会通过TTL自动过期,无需手动处理 for cache_key, memory in self.memory_cache.items(): if ( current_time - memory.updated_at ).days > get_runtime_setting('LLM_MEMORY_RETENTION_DAYS'): expired_sessions.append(cache_key) # 只清理内存缓存,不删除Redis中的键(Redis会自动过期) for cache_key in expired_sessions: if cache_key in self.memory_cache: del self.memory_cache[cache_key] if expired_sessions: logger.info(f"清理了{len(expired_sessions)}个过期内存会话记忆") except asyncio.CancelledError: break except Exception as e: logger.error(f"清理记忆时发生错误: {e}") memory_manager = MemoryManager()