feat: implement PowerMem LangChain middleware

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

@@ -1,19 +1,17 @@
-"""LangChain middleware entry point for PowerMem.
-
-The VLDB 2026 summer school branch intentionally provides only the public entry
-point. Students are expected to replace this placeholder with a LangChain
-middleware implementation that satisfies the package contract tests.
-"""
+"""LangChain middleware for loading and storing PowerMem memories."""
 
 from __future__ import annotations
 
+import asyncio
+import inspect
 from typing import Any, NotRequired, TypedDict
 
 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."""
+    """State schema used by the PowerMem middleware."""
 
     powermem_context: NotRequired[str]
 
@@ -22,10 +20,11 @@ class PowerMemStateUpdate(TypedDict):
     """State update returned by memory-loading middleware hooks."""
 
     powermem_context: str
+    messages: NotRequired[list]
 
 
 class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]):
-    """Placeholder for the summer school implementation."""
+    """Add relevant user memories to an agent prompt and persist its reply."""
 
     state_schema = PowerMemState
 
@@ -38,17 +37,134 @@ def __init__(
         save_interactions: bool = True,
         **kwargs: Any,
     ) -> None:
-        pass
+        self.memory = memory
+        self.user_id = user_id
+        self.search_limit = search_limit
+        self.save_interactions = save_interactions
+
+        if not hasattr(memory, "search"):
+            raise TypeError("memory must provide a search method")
+        if search_limit < 1:
+            raise ValueError("search_limit must be positive")
 
     def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None:
-        pass
+        query = self._latest_human_text(state)
+        if self.user_id is None or not query:
+            return None
+        try:
+            result = self.memory.search(query, user_id=self.user_id, limit=self.search_limit)
+        except Exception:
+            return None
+        return self._memory_update(result)
 
     async def abefore_agent(
         self,
         state: PowerMemState,
         runtime,
     ) -> PowerMemStateUpdate | None:
-        pass
+        query = self._latest_human_text(state)
+        if self.user_id is None or not query:
+            return None
+        try:
+            search = getattr(self.memory, "asearch", None)
+            if callable(search):
+                result = search(query, user_id=self.user_id, limit=self.search_limit)
+            elif inspect.iscoroutinefunction(self.memory.search):
+                result = self.memory.search(
+                    query, user_id=self.user_id, limit=self.search_limit
+                )
+            else:
+                result = await asyncio.to_thread(
+                    self.memory.search,
+                    query,
+                    user_id=self.user_id,
+                    limit=self.search_limit,
+                )
+            if inspect.isawaitable(result):
+                result = await result
+        except Exception:
+            return None
+        return self._memory_update(result)
 
     def after_agent(self, state: PowerMemState, runtime) -> None:
-        pass
+        if not self.save_interactions or self.user_id is None:
+            return
+        messages = self._interaction_messages(state)
+        if not messages:
+            return
+        try:
+            self.memory.add(messages=messages, user_id=self.user_id)
+        except Exception:
+            return
+
+    async def aafter_agent(self, state: PowerMemState, runtime) -> None:
+        if not self.save_interactions or self.user_id is None:
+            return
+        messages = self._interaction_messages(state)
+        if not messages:
+            return
+        try:
+            if inspect.iscoroutinefunction(self.memory.add):
+                result = self.memory.add(messages=messages, user_id=self.user_id)
+            else:
+                result = await asyncio.to_thread(
+                    self.memory.add, messages=messages, user_id=self.user_id
+                )
+            if inspect.isawaitable(result):
+                await result
+        except Exception:
+            return
+
+    @classmethod
+    def _memory_update(cls, result: Any) -> PowerMemStateUpdate | None:
+        context = cls._context_from_search(result)
+        if not context:
+            return None
+        return {
+            "powermem_context": context,
+            "messages": [SystemMessage(content=context)],
+        }
+
+    @staticmethod
+    def _latest_human_text(state: PowerMemState) -> str:
+        for message in reversed(state.get("messages", [])):
+            if isinstance(message, HumanMessage):
+                return str(message.content)
+        return ""
+
+    @staticmethod
+    def _context_from_search(result: Any) -> str:
+        if not isinstance(result, dict):
+            return ""
+        memories = result.get("results", [])
+        if not isinstance(memories, list):
+            return ""
+        texts = [
+            str(item["memory"])
+            for item in memories
+            if isinstance(item, dict) and item.get("memory")
+        ]
+        if not texts:
+            return ""
+        return "Relevant memories from PowerMem:\n" + "\n".join(f"- {text}" for text in texts)
+
+    @staticmethod
+    def _interaction_messages(state: PowerMemState) -> list[dict[str, str]]:
+        human = next(
+            (
+                message
+                for message in reversed(state.get("messages", []))
+                if isinstance(message, HumanMessage)
+            ),
+            None,
+        )
+        assistant = next(
+            (message for message in reversed(state.get("messages", [])) if isinstance(message, AIMessage)),
+            None,
+        )
+        if human is None or assistant is None:
+            return []
+        return [
+            {"role": "user", "content": str(human.content)},
+            {"role": "assistant", "content": str(assistant.content)},
+        ]

packages/powermem-langchain/tests/test_middleware.py

@@ -1,4 +1,4 @@
-from __future__ import annotations
+from __future__ import annotations
 
 import uuid
 from pathlib import Path
@@ -7,7 +7,7 @@
 import pytest
 from langchain.agents import create_agent
 from langchain_core.language_models.chat_models import SimpleChatModel
-from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
 from powermem import Memory
 from powermem_langchain import PowerMemMiddleware
 from pydantic import Field
@@ -38,6 +38,26 @@ def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
         raise RuntimeError("search failed")
 
 
+class CapturingMemory:
+    def __init__(self) -> None:
+        self.add_calls: list[tuple[list[dict[str, str]], str]] = []
+
+    def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        return {"results": []}
+
+    def add(self, *, messages: list[dict[str, str]], user_id: str) -> None:
+        self.add_calls.append((messages, user_id))
+class AsyncCapturingMemory:
+    def __init__(self) -> None:
+        self.search_calls: list[tuple[str, str, int]] = []
+        self.add_calls: list[tuple[list[dict[str, str]], str]] = []
+
+    async def search(self, query: str, *, user_id: str, limit: int) -> dict[str, Any]:
+        self.search_calls.append((query, user_id, limit))
+        return {"results": [{"memory": "Async user preference."}]}
+
+    async def add(self, *, messages: list[dict[str, str]], user_id: str) -> None:
+        self.add_calls.append((messages, user_id))
 def _message_text(messages: list[BaseMessage]) -> str:
     return "\n".join(str(message.content) for message in messages)
 
@@ -183,3 +203,70 @@ async def test_async_agent_uses_powermem_memory(tmp_path: Path):
     await agent.ainvoke({"messages": [HumanMessage(content="Use my async profile.")]})
 
     assert "User prefers async examples." in _message_text(model.calls[0])
+
+
+def test_no_user_id_skips_memory_io(tmp_path: Path):
+    memory = _sqlite_memory(tmp_path)
+    model = CapturingChatModel(responses=["ok"])
+    agent = create_agent(
+        model=model,
+        tools=[],
+        middleware=[PowerMemMiddleware(memory=memory, save_interactions=True)],
+    )
+
+    agent.invoke({"messages": [HumanMessage(content="Do not use memory.")]})
+
+    assert memory.get_all()["results"] == []
+
+
+@pytest.mark.asyncio
+async def test_async_memory_search_and_write_back_are_awaited():
+    memory = AsyncCapturingMemory()
+    middleware = PowerMemMiddleware(memory=memory, user_id="async-user")
+    state = {
+        "messages": [
+            HumanMessage(content="Use my async profile."),
+            AIMessage(content="Using it now."),
+        ]
+    }
+
+    update = await middleware.abefore_agent(state, runtime=None)
+    await middleware.aafter_agent(state, runtime=None)
+
+    assert memory.search_calls == [("Use my async profile.", "async-user", 5)]
+    assert update is not None
+    assert "Async user preference." in update["powermem_context"]
+    assert memory.add_calls == [
+        (
+            [
+                {"role": "user", "content": "Use my async profile."},
+                {"role": "assistant", "content": "Using it now."},
+            ],
+            "async-user",
+        )
+    ]
+
+
+def test_persists_latest_interaction_from_multi_turn_state():
+    memory = CapturingMemory()
+    middleware = PowerMemMiddleware(memory=memory, user_id="alice")
+    state = {
+        "messages": [
+            HumanMessage(content="Old question."),
+            AIMessage(content="Old answer."),
+            HumanMessage(content="Current question."),
+            AIMessage(content="Current answer."),
+        ]
+    }
+
+    middleware.after_agent(state, runtime=None)
+
+    assert memory.add_calls == [
+        (
+            [
+                {"role": "user", "content": "Current question."},
+                {"role": "assistant", "content": "Current answer."},
+            ],
+            "alice",
+        )
+    ]