diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..b549c91b80 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -13,19 +13,27 @@ class PowerMemState(AgentState): - """State schema reserved for the PowerMem middleware implementation.""" + """可选的状态字典类型,包含可选的 powermem_context 字段。""" powermem_context: NotRequired[str] class PowerMemStateUpdate(TypedDict): - """State update returned by memory-loading middleware hooks.""" + """可更新状态的字典类型,包含可选的 powermem_context 字段。""" powermem_context: str +import logging + +from langchain.agents.middleware import AgentMiddleware +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage + +logger = logging.getLogger(__name__) + + + class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" state_schema = PowerMemState @@ -38,17 +46,185 @@ def __init__( save_interactions: bool = True, **kwargs: Any, ) -> None: - pass - def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + super().__init__(**kwargs) + self.memory = memory + self.user_id = user_id + self.search_limit = search_limit + self.save_interactions = save_interactions - async def abefore_agent( - self, - state: PowerMemState, - runtime, - ) -> PowerMemStateUpdate | None: - pass + # ------------------------------------------------------------------ + # Hooks – before the agent runs (retrieve & inject memories) + # ------------------------------------------------------------------ + + def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: + """在调用agent之前,从PowerMem中检索相关记忆,并将其作为系统消息插入到消息列表的最前面。""" + context = self._retrieve_memories_sync(state) # 从memory获取 + if context is None: + return None + # 把memory作为系统消息插入到消息列表最前面 + state["messages"].insert(0, SystemMessage(content=context)) + return {"powermem_context": context} + + async def abefore_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: + context = await self._retrieve_memories_async(state) + if context is None: + return None + state["messages"].insert(0, SystemMessage(content=context)) + return {"powermem_context": context} + + # ------------------------------------------------------------------ + # Hooks – after the agent finishes (save interaction) + # ------------------------------------------------------------------ def after_agent(self, state: PowerMemState, runtime) -> None: - pass + """在调用agent之后,将用户和AI的对话保存到PowerMem中。""" + self._save_interaction_sync(state) + + async def aafter_agent(self, state: PowerMemState, runtime) -> None: + """Async version of after_agent.""" + await self._save_interaction_async(state) + + # ------------------------------------------------------------------ + # Internal – memory retrieval + # ------------------------------------------------------------------ + + def _retrieve_memories_sync(self, state: PowerMemState) -> str | None: + """同步记忆检索。返回格式化的上下文或 None。""" + user_msg = self._get_last_user_message(state) # 获取最近的用户信息 + if user_msg is None: + return None + + # 使用用户消息内容作为查询词,在记忆库中搜索与当前用户相关的历史记忆,最多返回 search_limit 条 + try: + result = self.memory.search( + query=user_msg.content, + user_id=self.user_id, + limit=self.search_limit, + ) + except Exception: + logger.warning("PowerMem search failed", exc_info=True) + return None + # 提取记忆列表 + memories = result.get("results", []) if isinstance(result, dict) else [] + if not memories: + return None + # 将记忆列表格式化为文本块,返回给调用者 + return self._format_memories(memories) + + async def _retrieve_memories_async(self, state: PowerMemState) -> str | None: + """异步记忆检索。返回格式化的上下文或 None。""" + user_msg = self._get_last_user_message(state) + if user_msg is None: + return None + + try: + # Use the async search if available; fall back to sync wrapped in a thread. + if hasattr(self.memory, "asearch"): + result = await self.memory.asearch( + query=user_msg.content, + user_id=self.user_id, + limit=self.search_limit, + ) + else: + import asyncio + + result = await asyncio.to_thread( + self.memory.search, + query=user_msg.content, + user_id=self.user_id, + limit=self.search_limit, + ) + except Exception: + logger.warning("PowerMem async search failed", exc_info=True) + return None + + memories = result.get("results", []) if isinstance(result, dict) else [] + if not memories: + return None + + return self._format_memories(memories) + + + @staticmethod + def _format_memories(memories) -> str: + """将记忆列表格式化为文本块,返回给调用者。""" + parts = [] + for m in memories: + if isinstance(m, dict): + # 记忆是字典,取 "memory" 字段 + parts.append(f"- {m.get('memory', str(m))}") + elif hasattr(m, "content"): + # 记忆是对象,有 .content 属性 + parts.append(f"- {m.content}") + elif isinstance(m, str): + # 记忆直接是字符串 + parts.append(f"- {m}") + else: + # 后备方案 + parts.append(f"- {str(m)}") + lines = "\n".join(parts) + return f"相关记忆:\n{lines}" + + # ------------------------------------------------------------------ + # Internal – interaction saving + # ------------------------------------------------------------------ + + def _save_interaction_sync(self, state: PowerMemState) -> None: + """保存用户和AI的对话到PowerMem中。""" + if not self.save_interactions: + return + user_msg, ai_msg = self._get_last_user_ai_pair(state) + if user_msg is None or ai_msg is None: + return + text = f"用户:{user_msg.content}\n助手:{ai_msg.content}" + try: + self.memory.add(messages=text, user_id=self.user_id) + except Exception: + logger.warning("PowerMem save failed", exc_info=True) + + async def _save_interaction_async(self, state: PowerMemState) -> None: + """异步保存用户和AI的对话到PowerMem中。""" + if not self.save_interactions: + return + user_msg, ai_msg = self._get_last_user_ai_pair(state) + if user_msg is None or ai_msg is None: + return + text = f"用户:{user_msg.content}\n助手:{ai_msg.content}" + try: + if hasattr(self.memory, "aadd"): + await self.memory.aadd(messages=text, user_id=self.user_id) + else: + import asyncio + await asyncio.to_thread( + self.memory.add, messages=text, user_id=self.user_id + ) + except Exception: + logger.warning("PowerMem async save failed", exc_info=True) + + # ------------------------------------------------------------------ + # Helpers – message extraction + # ------------------------------------------------------------------ + + @staticmethod + def _get_last_user_message(state: PowerMemState) -> HumanMessage | None: + """Return the most recent HumanMessage, or None.""" + for msg in reversed(state["messages"]): + if isinstance(msg, HumanMessage): + return msg + return None + + @staticmethod + def _get_last_ai_message(state: PowerMemState) -> AIMessage | None: + """Return the most recent AIMessage, or None.""" + for msg in reversed(state["messages"]): + if isinstance(msg, AIMessage): + return msg + return None + + @classmethod + def _get_last_user_ai_pair( + cls, state: PowerMemState + ) -> tuple[HumanMessage | None, AIMessage | None]: + """Return the last user message and the last AI message as a pair.""" + return cls._get_last_user_message(state), cls._get_last_ai_message(state) \ No newline at end of file