From 102b31cf7da04cc2989fd78bd7a09edd9bf86c43 Mon Sep 17 00:00:00 2001 From: vladaliii <2773158174@qq.com> Date: Tue, 14 Jul 2026 15:18:14 +0800 Subject: [PATCH] impl Powermem LangChain Middleware --- .../examples/openai_agent.py | 27 ++- .../src/powermem_langchain/middleware.py | 225 +++++++++++++++++- 2 files changed, 239 insertions(+), 13 deletions(-) diff --git a/packages/powermem-langchain/examples/openai_agent.py b/packages/powermem-langchain/examples/openai_agent.py index f06972a0bb..0ddb0cd530 100644 --- a/packages/powermem-langchain/examples/openai_agent.py +++ b/packages/powermem-langchain/examples/openai_agent.py @@ -1,9 +1,8 @@ -"""Run a PowerMem-backed LangChain agent with an OpenAI chat model. +"""使用 OpenAI 聊天模型运行一个由 PowerMem 提供长期记忆的 LangChain agent。 -This example is intentionally written as an end-to-end CLI check. Once -PowerMemMiddleware is implemented and the environment is configured, it should -seed PowerMem, invoke a LangChain agent backed by OpenAI, and print enough -information to verify memory retrieval and write-back behavior. +这个示例是一个端到端的命令行检查程序。实现 PowerMemMiddleware 并配置好环境后， +它会先向 PowerMem 写入一条种子记忆，再调用由 OpenAI 驱动的 LangChain agent， +最后打印足够的信息，用来确认记忆检索、上下文注入和交互写回是否生效。 """ from __future__ import annotations @@ -28,6 +27,10 @@ class DemoSettings(BaseSettings): + """从环境变量和可选的 .env 文件读取示例配置。""" + + # 专用于本示例的变量可以使用 POWERMEM_LANGCHAIN_ 前缀；extra="ignore" + # 让 PowerMem 自己使用的其他环境变量不会触发 Pydantic 校验错误。 model_config = SettingsConfigDict( env_file=".env", env_file_encoding="utf-8", @@ -37,6 +40,7 @@ class DemoSettings(BaseSettings): openai_api_key: str | None = Field( default=None, + # 同时兼容 OpenAI 的标准变量和 PowerMem 示例中使用的 LLM_API_KEY。 validation_alias=AliasChoices("OPENAI_API_KEY", "LLM_API_KEY"), ) openai_base_url: str | None = Field( @@ -54,6 +58,8 @@ class DemoSettings(BaseSettings): def parse_args(settings: DemoSettings) -> argparse.Namespace: + """解析命令行参数，并以环境配置作为默认值。""" + parser = argparse.ArgumentParser( description="Run the PowerMem LangChain middleware OpenAI demo." ) @@ -77,6 +83,9 @@ def parse_args(settings: DemoSettings) -> argparse.Namespace: def create_openai_model(settings: DemoSettings, model: str, temperature: float): + """按需导入并创建 ChatOpenAI，缺少依赖或密钥时给出明确提示。""" + + # 延迟导入使包的基础功能和测试不必安装示例专用的 langchain-openai。 try: from langchain_openai import ChatOpenAI except ImportError as exc: @@ -101,6 +110,8 @@ def create_openai_model(settings: DemoSettings, model: str, temperature: float): def print_search_results(title: str, result: dict[str, Any]) -> None: + """以便于人工核对的格式打印 PowerMem 检索结果。""" + memories = result.get("results", []) print(title) if not memories: @@ -113,6 +124,7 @@ def print_search_results(title: str, result: dict[str, Any]) -> None: def main() -> None: + # 配置优先来自命令行；未传入的参数使用环境变量或类中定义的默认值。 settings = DemoSettings() args = parse_args(settings) @@ -121,6 +133,7 @@ def main() -> None: print(f"model: {args.model}") try: + # create_memory() 会读取 PowerMem 自己的数据库、LLM 和嵌入模型配置。 memory = create_memory() except Exception as exc: raise SystemExit( @@ -130,9 +143,11 @@ def main() -> None: ) from exc if not args.skip_seed: + # infer=False 表示直接保存演示文本，不让 PowerMem 的 LLM 再做事实提取。 memory.add(args.seed_memory, user_id=args.user_id, infer=False) print(f"seed_memory: {args.seed_memory}") + # 调用 agent 前先检索一次，用于确认种子记忆已经存在。 before = memory.search( args.prompt, user_id=args.user_id, @@ -147,6 +162,7 @@ def main() -> None: save_interactions=not args.no_save, ) + # middleware 应在模型调用前注入相关记忆，并在 agent 结束后按需写回交互。 agent = create_agent( model=create_openai_model(settings, args.model, args.temperature), tools=[], @@ -159,6 +175,7 @@ def main() -> None: print("assistant:") print(answer) + # 再次检索同一主题，用于观察本轮用户消息和助手回复是否已被写回。 after = memory.search( args.prompt, user_id=args.user_id, diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..31070a5f1e 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -7,9 +7,17 @@ from __future__ import annotations +import asyncio +from collections.abc import Awaitable, Callable from typing import Any, NotRequired, TypedDict -from langchain.agents.middleware import AgentMiddleware, AgentState +from langchain.agents.middleware import ( + AgentMiddleware, + AgentState, + ModelRequest, + ModelResponse, +) +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage class PowerMemState(AgentState): @@ -25,7 +33,7 @@ class PowerMemStateUpdate(TypedDict): class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" + """Connect a LangChain agent to a PowerMem memory instance.""" state_schema = PowerMemState @@ -38,17 +46,218 @@ def __init__( save_interactions: bool = True, **kwargs: Any, ) -> None: - pass + """Store the configuration used by the middleware hooks. - def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + Args: + memory: PowerMem-compatible memory instance. + user_id: Explicit business user identity used for search and write-back. + search_limit: Maximum number of memories retrieved for one request. + save_interactions: Whether completed interactions should be written back. + **kwargs: Reserved for future middleware options. + """ + if memory is None: + raise ValueError("memory must not be None") + if not isinstance(user_id, str) or not user_id.strip(): + raise ValueError("user_id must be a non-empty string") + if not isinstance(search_limit, int) or isinstance(search_limit, bool): + raise TypeError("search_limit must be an integer") + if search_limit <= 0: + raise ValueError("search_limit must be greater than zero") + if not isinstance(save_interactions, bool): + raise TypeError("save_interactions must be a boolean") + + self.memory = memory + self.user_id = user_id + self.search_limit = search_limit + self.save_interactions = save_interactions + self.options = dict(kwargs) + + @staticmethod + def _latest_user_text(state: PowerMemState) -> str | None: + """Return the text of the latest user message.""" + for message in reversed(state.get("messages", [])): + if not isinstance(message, HumanMessage): + continue + + if isinstance(message.content, str): + text = message.content.strip() + return text or None + + return None + + @staticmethod + def _latest_assistant_text(state: PowerMemState) -> str | None: + """Return the text of the final non-empty assistant message.""" + for message in reversed(state.get("messages", [])): + if not isinstance(message, AIMessage): + continue + + if isinstance(message.content, str): + text = message.content.strip() + + if text: + return text + + return None + + def _search_context(self, query: str) -> str: + """Search PowerMem and format retrieved memories for the model.""" + try: + result = self.memory.search( + query, + user_id=self.user_id, + limit=self.search_limit, + ) + except Exception: + # PowerMem is an enhancement. Search failures must not stop the agent. + return "" + + if not isinstance(result, dict): + return "" + + memories: list[str] = [] + + for item in result.get("results", []): + if not isinstance(item, dict): + continue + + memory_text = item.get("memory") + + if isinstance(memory_text, str) and memory_text.strip(): + memories.append(memory_text.strip()) + + if not memories: + return "" + + formatted_memories = "\n".join( + f"- {memory_text}" for memory_text in memories + ) + + return ( + "Relevant long-term memories from PowerMem:\n" + f"{formatted_memories}" + ) + + @staticmethod + def _request_with_context(request: ModelRequest) -> ModelRequest: + """Return a model request containing the retrieved PowerMem context.""" + context = request.state.get("powermem_context", "") + + if not context: + return request + + current_system_message = request.system_message + + if current_system_message is None: + new_system_message = SystemMessage(content=context) + + elif isinstance(current_system_message.content, str): + original_content = current_system_message.content.strip() + + if original_content: + new_content = f"{original_content}\n\n{context}" + else: + new_content = context + + new_system_message = SystemMessage(content=new_content) + + else: + new_content = list(current_system_message.content) + new_content.append( + { + "type": "text", + "text": context, + } + ) + + new_system_message = SystemMessage(content=new_content) + + return request.override( + system_message=new_system_message, + ) + + def before_agent( + self, + state: PowerMemState, + runtime, + ) -> PowerMemStateUpdate: + """Load memories related to the latest user message.""" + query = self._latest_user_text(state) + + if query is None: + return {"powermem_context": ""} + + context = self._search_context(query) + + return { + "powermem_context": context, + } + + def wrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], ModelResponse], + ) -> ModelResponse: + """Inject PowerMem context and execute the synchronous model call.""" + new_request = self._request_with_context(request) + return handler(new_request) + + async def awrap_model_call( + self, + request: ModelRequest, + handler: Callable[ + [ModelRequest], + Awaitable[ModelResponse], + ], + ) -> ModelResponse: + """Inject PowerMem context and execute the asynchronous model call.""" + new_request = self._request_with_context(request) + return await handler(new_request) async def abefore_agent( self, state: PowerMemState, runtime, - ) -> PowerMemStateUpdate | None: - pass + ) -> PowerMemStateUpdate: + """Asynchronously load memories related to the latest user message.""" + query = self._latest_user_text(state) + + if query is None: + return {"powermem_context": ""} + + context = await asyncio.to_thread( + self._search_context, + query, + ) + + return { + "powermem_context": context, + } def after_agent(self, state: PowerMemState, runtime) -> None: - pass + """Write the completed interaction back to PowerMem.""" + if not self.save_interactions: + return None + + user_text = self._latest_user_text(state) + assistant_text = self._latest_assistant_text(state) + + if user_text is None or assistant_text is None: + return None + + interaction = ( + f"User: {user_text}\n" + f"Assistant: {assistant_text}" + ) + + try: + self.memory.add( + interaction, + user_id=self.user_id, + infer=False, + ) + except Exception: + # 写回失败不能破坏已经成功生成的 Agent 回复。 + return None + + return None