Vldb 2026

packages/powermem-langchain/examples/openai_agent.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import argparse
+import traceback
 from typing import Any
 
 from langchain.agents import create_agent
@@ -123,6 +124,7 @@ def main() -> None:
     try:
         memory = create_memory()
     except Exception as exc:
+        traceback.print_exc()
         raise SystemExit(
             "Failed to create PowerMem memory. Configure PowerMem first, for "
             "example by setting LLM_PROVIDER, LLM_API_KEY, DATABASE_PROVIDER, "

packages/powermem-langchain/src/powermem_langchain/middleware.py

@@ -7,25 +7,26 @@
 
 from __future__ import annotations
 
-from typing import Any, NotRequired, TypedDict
+import asyncio
+import logging
+from typing import Any, NotRequired
 
 from langchain.agents.middleware import AgentMiddleware, AgentState
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
 
-class PowerMemState(AgentState):
-    """State schema reserved for the PowerMem middleware implementation."""
-
-    powermem_context: NotRequired[str]
+logger = logging.getLogger(__name__)
 
 
-class PowerMemStateUpdate(TypedDict):
-    """State update returned by memory-loading middleware hooks."""
+class PowerMemState(AgentState):
+    """State schema reserved for the PowerMem middleware implementation."""
 
-    powermem_context: str
+    powermem_memories: NotRequired[list[str]]
 
 
 class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]):
-    """Placeholder for the summer school implementation."""
+    """LangChain v1 middleware that injects PowerMem retrieval into the model
+    context and persists the interaction after the agent run completes."""
 
     state_schema = PowerMemState
 
@@ -36,19 +37,168 @@ def __init__(
         user_id: str | None = None,
         search_limit: int = 5,
         save_interactions: bool = True,
+        system_template: str = (
+            "<system>\n"
+            "The following memories were retrieved from PowerMem and may be "
+            "relevant to the user's request. Use them when answering, but do "
+            "not mention them explicitly unless asked.\n"
+            "{memories}\n"
+            "</system>"
+        ),
+        infer: bool = False,
         **kwargs: Any,
     ) -> None:
-        pass
-
-    def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None:
-        pass
+        if user_id is None:
+            raise ValueError(
+                "PowerMemMiddleware requires an explicit user_id; "
+                "do not rely on global or runtime-inferred user identity."
+            )
+        self.memory = memory
+        self.user_id = user_id
+        self.search_limit = search_limit
+        self.save_interactions = save_interactions
+        self.system_template = system_template
+        self.infer = infer
+
+    def _retrieve(self, query: str) -> list[str]:
+        try:
+            result = self.memory.search(
+                query,
+                user_id=self.user_id,
+                limit=self.search_limit,
+            )
+        except Exception as exc:
+            logger.warning(
+                "PowerMem search failed for user_id=%s: %s", self.user_id, exc
+            )
+            return []
+        memories = result.get("results", []) if isinstance(result, dict) else []
+        return [
+            item["memory"]
+            for item in memories
+            if isinstance(item, dict) and item.get("memory")
+        ]
+
+    async def _aretrieve(self, query: str) -> list[str]:
+        return await asyncio.to_thread(self._retrieve, query)
+
+    @staticmethod
+    def _latest_user_text(messages: list[Any] | None) -> str | None:
+        if not messages:
+            return None
+        for message in reversed(messages):
+            if isinstance(message, HumanMessage):
+                content = message.content
+                if isinstance(content, str) and content.strip():
+                    return content
+                if isinstance(content, list):
+                    parts = [
+                        part.get("text", "")
+                        for part in content
+                        if isinstance(part, dict) and part.get("text")
+                    ]
+                    joined = "".join(parts).strip()
+                    if joined:
+                        return joined
+        return None
+
+    @staticmethod
+    def _last_of_type(messages: list[Any] | None, message_type: type) -> Any | None:
+        if not messages:
+            return None
+        for message in reversed(messages):
+            if isinstance(message, message_type):
+                return message
+        return None
+
+    @staticmethod
+    def _message_text(message: Any) -> str:
+        content = getattr(message, "content", message)
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts = [
+                part.get("text", "")
+                for part in content
+                if isinstance(part, dict) and part.get("text")
+            ]
+            return "".join(parts)
+        return str(content)
+
+    def before_agent(self, state: PowerMemState, runtime) -> dict[str, Any] | None:
+        query = self._latest_user_text(state.get("messages"))
+        if not query:
+            return None
+        memories = self._retrieve(query)
+        return {"powermem_memories": memories}
 
     async def abefore_agent(
         self,
         state: PowerMemState,
         runtime,
-    ) -> PowerMemStateUpdate | None:
-        pass
+    ) -> dict[str, Any] | None:
+        query = self._latest_user_text(state.get("messages"))
+        if not query:
+            return None
+        memories = await self._aretrieve(query)
+        return {"powermem_memories": memories}
+
+    def _format_system_message(self, memories: list[str]) -> SystemMessage:
+        body = "\n".join(f"- {memory}" for memory in memories)
+        return SystemMessage(content=self.system_template.format(memories=body))
+
+    def wrap_model_call(self, request, handler):
+        memories = request.state.get("powermem_memories") or []
+        if not memories:
+            return handler(request)
+        system_message = self._format_system_message(memories)
+        request = request.override(messages=[system_message, *request.messages])
+        return handler(request)
+
+    async def awrap_model_call(self, request, handler):
+        memories = request.state.get("powermem_memories") or []
+        if not memories:
+            return await handler(request)
+        system_message = self._format_system_message(memories)
+        request = request.override(messages=[system_message, *request.messages])
+        return await handler(request)
+
+    def _persist_interaction(self, user_message: Any, assistant_message: Any) -> None:
+        try:
+            self.memory.add(
+                messages=[
+                    {"role": "user", "content": self._message_text(user_message)},
+                    {
+                        "role": "assistant",
+                        "content": self._message_text(assistant_message),
+                    },
+                ],
+                user_id=self.user_id,
+                infer=self.infer,
+            )
+        except Exception as exc:
+            logger.warning(
+                "PowerMem add failed for user_id=%s: %s", self.user_id, exc
+            )
 
     def after_agent(self, state: PowerMemState, runtime) -> None:
-        pass
+        if not self.save_interactions:
+            return
+        messages = state.get("messages") or []
+        user_message = self._last_of_type(messages, HumanMessage)
+        assistant_message = self._last_of_type(messages, AIMessage)
+        if user_message is None or assistant_message is None:
+            return
+        self._persist_interaction(user_message, assistant_message)
+
+    async def aafter_agent(self, state: PowerMemState, runtime) -> None:
+        if not self.save_interactions:
+            return
+        messages = state.get("messages") or []
+        user_message = self._last_of_type(messages, HumanMessage)
+        assistant_message = self._last_of_type(messages, AIMessage)
+        if user_message is None or assistant_message is None:
+            return
+        await asyncio.to_thread(
+            self._persist_interaction, user_message, assistant_message
+        )