Task1: PowerMem #四个人一队 || Task1: PowerMem #A team of four people

packages/powermem-langchain/README.md

@@ -1,9 +1,8 @@
 # powermem-langchain
 
 This package is the LangChain middleware exercise for the VLDB 2026 summer
-school branch. It provides a package skeleton, a no-op middleware scaffold,
-contract tests, and a runnable OpenAI example. It does not provide a complete
-PowerMem middleware implementation.
+school branch. It provides a PowerMem middleware implementation, contract
+tests, and a runnable OpenAI example.
 
 The goal is to integrate PowerMem as a long-term memory layer for LangChain v1
 agents. The provided scaffold is intentionally small; you may adjust it as your
@@ -112,6 +111,5 @@ uv run --no-project \
     --user-id summer-school-demo
 ```
 
-With the placeholder middleware, the command can run but will not show the expected
-memory injection or write-back behavior. After implementation, the output should
-show seeded memories before the agent call and updated memories afterward.
+The output should show seeded memories before the agent call and updated memories
+afterward.

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

@@ -1,15 +1,22 @@
-"""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 entry point for PowerMem."""
 
 from __future__ import annotations
 
-from typing import Any, NotRequired, TypedDict
+import asyncio
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Any, NotRequired, TypedDict, cast
+
+from langchain.agents.middleware import (
+    AgentMiddleware,
+    AgentState,
+    ModelRequest,
+    ModelResponse,
+)
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
 
-from langchain.agents.middleware import AgentMiddleware, AgentState
+
+logger = logging.getLogger(__name__)
 
 
 class PowerMemState(AgentState):
@@ -25,7 +32,7 @@ class PowerMemStateUpdate(TypedDict):
 
 
 class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]):
-    """Placeholder for the summer school implementation."""
+    """Use PowerMem as a long-term memory layer for a LangChain agent."""
 
     state_schema = PowerMemState
 
@@ -38,17 +45,129 @@ 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
+
+    @staticmethod
+    def _latest_message_text(
+        messages: list[BaseMessage],
+        message_type: type[BaseMessage],
+    ) -> str | None:
+        for message in reversed(messages):
+            if isinstance(message, message_type) and message.text:
+                return message.text
+        return None
+
+    @staticmethod
+    def _format_search_results(result: Any) -> str | None:
+        if not isinstance(result, dict):
+            return None
+
+        memories: list[str] = []
+        for item in result.get("results", []):
+            if isinstance(item, str):
+                content = item
+            elif isinstance(item, dict):
+                content = item.get("memory") or item.get("content")
+            else:
+                continue
+
+            if content:
+                memories.append(str(content))
+
+        if not memories:
+            return None
+        return "\n".join(f"- {memory}" for memory in memories)
+
+    def _load_memory(self, state: PowerMemState) -> PowerMemStateUpdate | None:
+        query = self._latest_message_text(state["messages"], HumanMessage)
+        if not query:
+            return None
+
+        try:
+            result = self.memory.search(
+                query,
+                user_id=self.user_id,
+                limit=self.search_limit,
+            )
+        except Exception:
+            logger.warning(
+                "PowerMem search failed; continuing without memory",
+                exc_info=True,
+            )
+            return None
+
+        context = self._format_search_results(result)
+        if context is None:
+            return None
+        return {"powermem_context": context}
+
+    @staticmethod
+    def _inject_memory(request: ModelRequest[Any]) -> ModelRequest[Any]:
+        state = request.state or {}
+        context = state.get("powermem_context")
+        if not context:
+            return request
+
+        memory_prompt = f"Relevant long-term memories:\n{context}"
+        if request.system_message is None:
+            system_message = SystemMessage(content=memory_prompt)
+        else:
+            content = [
+                *request.system_message.content_blocks,
+                {"type": "text", "text": f"\n\n{memory_prompt}"},
+            ]
+            system_message = SystemMessage(
+                content=cast("list[str | dict[str, str]]", content)
+            )
+
+        return request.override(system_message=system_message)
+
+    def _save_interaction(self, state: PowerMemState) -> None:
+        messages = state["messages"]
+        user_message = self._latest_message_text(messages, HumanMessage)
+        assistant_message = self._latest_message_text(messages, AIMessage)
+        if not user_message or not assistant_message:
+            return
+
+        self.memory.add(
+            [
+                {"role": "user", "content": user_message},
+                {"role": "assistant", "content": assistant_message},
+            ],
+            user_id=self.user_id,
+        )
 
     def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None:
-        pass
+        return self._load_memory(state)
 
     async def abefore_agent(
         self,
         state: PowerMemState,
         runtime,
     ) -> PowerMemStateUpdate | None:
-        pass
+        return await asyncio.to_thread(self._load_memory, state)
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest[Any],
+        handler: Callable[[ModelRequest[Any]], ModelResponse[Any]],
+    ) -> ModelResponse[Any]:
+        return handler(self._inject_memory(request))
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest[Any],
+        handler: Callable[[ModelRequest[Any]], Awaitable[ModelResponse[Any]]],
+    ) -> ModelResponse[Any]:
+        return await handler(self._inject_memory(request))
 
     def after_agent(self, state: PowerMemState, runtime) -> None:
-        pass
+        if self.save_interactions:
+            self._save_interaction(state)
+
+    async def aafter_agent(self, state: PowerMemState, runtime) -> None:
+        if self.save_interactions:
+            await asyncio.to_thread(self._save_interaction, state)

packages/powermem-langchain/tests/test_middleware.py

@@ -183,3 +183,28 @@ 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])
+
+
+@pytest.mark.asyncio
+async def test_async_agent_persists_interaction(tmp_path: Path):
+    memory = _sqlite_memory(tmp_path)
+    model = CapturingChatModel(responses=["Async stored response"])
+    agent = create_agent(
+        model=model,
+        tools=[],
+        middleware=[
+            PowerMemMiddleware(
+                memory=memory,
+                user_id="async-user",
+                save_interactions=True,
+            )
+        ],
+    )
+
+    await agent.ainvoke(
+        {"messages": [HumanMessage(content="Remember this async note.")]}
+    )
+
+    memory_text = _stored_memory_text(memory, "async-user")
+    assert "Remember this async note." in memory_text
+    assert "Async stored response" in memory_text