""" Memori Agent 中间件 实现记忆召回和存储的 AgentMiddleware """ import asyncio import logging from typing import Any, Dict, List, Optional from langchain.agents.middleware import AgentMiddleware, AgentState from langgraph.runtime import Runtime from .memori_config import MemoriConfig from .memori_manager import MemoriManager, get_memori_manager logger = logging.getLogger("app") class MemoriMiddleware(AgentMiddleware): """ Memori 记忆中间件 功能： 1. before_agent: 召回相关记忆并注入到上下文 2. after_agent: 后台异步提取和存储新记忆 """ def __init__( self, memori_manager: MemoriManager, config: MemoriConfig, ): """初始化 MemoriMiddleware Args: memori_manager: MemoriManager 实例 config: MemoriConfig 配置 """ self.memori_manager = memori_manager self.config = config def _extract_user_query(self, state: AgentState) -> str: """从状态中提取用户查询 Args: state: Agent 状态 Returns: 用户查询文本 """ messages = state.get("messages", []) if not messages: return "" # 获取最后一条消息 last_message = messages[-1] # 尝试获取内容 content = getattr(last_message, "content", None) if content is None: content = last_message.get("content", "") if isinstance(last_message, dict) else "" return str(content) if content else "" def _format_memories(self, memories: List[Dict[str, Any]]) -> str: """格式化记忆列表为文本 Args: memories: 记忆列表 Returns: 格式化的记忆文本 """ if not memories: return "" lines = [] for i, memory in enumerate(memories, 1): content = memory.get("content", "") similarity = memory.get("similarity", 0.0) fact_type = memory.get("fact_type", "fact") # 添加相似度分数（调试用） lines.append(f"{i}. [{fact_type}] {content}") return "\n".join(lines) def _inject_memory_context(self, state: AgentState, memory_text: str) -> AgentState: """将记忆上下文注入到状态中 Args: state: 原始状态 memory_text: 记忆文本 Returns: 更新后的状态 """ if not memory_text or not self.config.inject_memory_to_system_prompt: return state # 生成记忆提示 memory_prompt = self.config.get_memory_prompt([memory_text]) # 检查是否有系统消息 messages = state.get("messages", []) if not messages: return state # 在系统消息后添加记忆上下文 from langchain_core.messages import SystemMessage # 查找系统消息 system_message = None for msg in messages: if hasattr(msg, "type") and msg.type == "system": system_message = msg break elif isinstance(msg, dict) and msg.get("role") == "system": system_message = msg break if system_message: # 修改现有系统消息 if hasattr(system_message, "content"): original_content = system_message.content system_message.content = original_content + memory_prompt elif isinstance(system_message, dict): original_content = system_message.get("content", "") system_message["content"] = original_content + memory_prompt else: # 添加新的系统消息 new_messages = list(messages) new_messages.insert(0, SystemMessage(content=memory_prompt)) state = {**state, "messages": new_messages} return state def before_agent(self, state: AgentState, runtime: Runtime) -> Dict[str, Any] | None: """Agent 执行前：召回相关记忆（同步版本） Args: state: Agent 状态 runtime: 运行时上下文 Returns: 更新后的状态或 None """ if not self.config.is_enabled(): return None try: # 提取用户查询 query = self._extract_user_query(state) if not query: return None # 获取 attribution 参数 entity_id, process_id = self.config.get_attribution_tuple() session_id = self.config.session_id or runtime.config.get("configurable", {}).get("thread_id", "default") # 召回记忆（同步方式 - 在后台任务中执行） memories = asyncio.run(self._recall_memories_async(query, entity_id, process_id, session_id)) if memories: # 格式化记忆 memory_text = self._format_memories(memories) # 注入到状态 updated_state = self._inject_memory_context(state, memory_text) logger.info(f"Injected {len(memories)} memories into context") return updated_state return None except Exception as e: logger.error(f"Error in MemoriMiddleware.before_agent: {e}") return None async def abefore_agent(self, state: AgentState, runtime: Runtime) -> Dict[str, Any] | None: """Agent 执行前：召回相关记忆（异步版本） Args: state: Agent 状态 runtime: 运行时上下文 Returns: 更新后的状态或 None """ if not self.config.is_enabled(): return None try: # 提取用户查询 query = self._extract_user_query(state) if not query: logger.debug("No user query found, skipping memory recall") return None # 获取 attribution 参数 entity_id, process_id = self.config.get_attribution_tuple() session_id = self.config.session_id or runtime.config.get("configurable", {}).get("thread_id", "default") # 召回记忆 memories = await self._recall_memories_async(query, entity_id, process_id, session_id) if memories: # 格式化记忆 memory_text = self._format_memories(memories) # 注入到状态 updated_state = self._inject_memory_context(state, memory_text) logger.info(f"Injected {len(memories)} memories into context (similarity > {self.config.semantic_search_threshold})") return updated_state return None except Exception as e: logger.error(f"Error in MemoriMiddleware.abefore_agent: {e}") return None async def _recall_memories_async( self, query: str, entity_id: str, process_id: str, session_id: str ) -> List[Dict[str, Any]]: """异步召回记忆 Args: query: 查询文本 entity_id: 实体 ID process_id: 进程 ID session_id: 会话 ID Returns: 记忆列表 """ return await self.memori_manager.recall_memories( query=query, entity_id=entity_id, process_id=process_id, session_id=session_id, config=self.config, ) def after_agent(self, state: AgentState, runtime: Runtime) -> None: """Agent 执行后：触发记忆增强（同步版本） Args: state: Agent 状态 runtime: 运行时上下文 """ if not self.config.is_enabled() or not self.config.augmentation_enabled: return try: # 触发后台增强任务 asyncio.create_task(self._trigger_augmentation_async(state, runtime)) except Exception as e: logger.error(f"Error in MemoriMiddleware.after_agent: {e}") async def aafter_agent(self, state: AgentState, runtime: Runtime) -> None: """Agent 执行后：触发记忆增强（异步版本） 注意：Memori 的增强会自动在后台执行，这里主要是记录日志 Args: state: Agent 状态 runtime: 运行时上下文 """ if not self.config.is_enabled() or not self.config.augmentation_enabled: return try: # 如果配置了等待超时，则等待增强完成 if self.config.augmentation_wait_timeout is not None: entity_id, process_id = self.config.get_attribution_tuple() session_id = self.config.session_id or runtime.config.get("configurable", {}).get("thread_id", "default") await self.memori_manager.wait_for_augmentation( entity_id=entity_id, process_id=process_id, session_id=session_id, timeout=self.config.augmentation_wait_timeout, ) except Exception as e: logger.error(f"Error in MemoriMiddleware.aafter_agent: {e}") async def _trigger_augmentation_async(self, state: AgentState, runtime: Runtime) -> None: """触发记忆增强任务 注意：Memori 的 LLM 客户端注册后会自动捕获对话并进行增强， 这里不需要手动触发，只是确保会话正确设置 Args: state: Agent 状态 runtime: 运行时上下文 """ # Memori 的增强是自动的，这里主要是确保配置正确 # 如果需要手动触发，可以在这里实现 pass def create_memori_middleware( bot_id: str, user_identifier: str, session_id: str, enabled: bool = True, semantic_search_top_k: int = 5, semantic_search_threshold: float = 0.7, memori_manager: Optional[MemoriManager] = None, ) -> Optional[MemoriMiddleware]: """创建 MemoriMiddleware 的工厂函数 Args: bot_id: Bot ID user_identifier: 用户标识 session_id: 会话 ID enabled: 是否启用 semantic_search_top_k: 语义搜索返回数量 semantic_search_threshold: 语义搜索相似度阈值 memori_manager: MemoriManager 实例（如果为 None，使用全局实例） Returns: MemoriMiddleware 实例或 None """ if not enabled: return None # 获取或使用提供的 manager manager = memori_manager or get_memori_manager() # 创建配置 config = MemoriConfig( enabled=True, entity_id=user_identifier, process_id=bot_id, session_id=session_id, semantic_search_top_k=semantic_search_top_k, semantic_search_threshold=semantic_search_threshold, ) return MemoriMiddleware(memori_manager=manager, config=config)