impl Powermem LangChain Middleware

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,

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