diff --git a/packages/powermem-langchain/README.md b/packages/powermem-langchain/README.md index 8cd1f28fd1..0de2e4adda 100644 --- a/packages/powermem-langchain/README.md +++ b/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. diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..83a8ed973d 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/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) diff --git a/packages/powermem-langchain/tests/test_middleware.py b/packages/powermem-langchain/tests/test_middleware.py index 4f2fbbcfd9..9c902e74d3 100644 --- a/packages/powermem-langchain/tests/test_middleware.py +++ b/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