diff --git a/packages/powermem-langchain/README.md b/packages/powermem-langchain/README.md index 8cd1f28fd1..e26e346f4e 100644 --- a/packages/powermem-langchain/README.md +++ b/packages/powermem-langchain/README.md @@ -1,9 +1,9 @@ # 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. +This package provides PowerMem long-term memory middleware for LangChain v1 +agents. It retrieves relevant memories before an agent run, injects them into +model-visible context, and can save the final user/assistant interaction after +the run. 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 @@ -13,7 +13,7 @@ This is useful beyond a single `create_agent` example: LangChain middleware is the extension point for controlling agent execution, and related projects such as Deep Agents also compose capabilities through middleware. -## Task +## Usage Implement: @@ -44,7 +44,7 @@ agent = create_agent( ) ``` -At minimum, the implementation should: +The middleware: - Search PowerMem with the latest user message before the model call. - Add relevant memories to the model-visible context. @@ -52,6 +52,13 @@ At minimum, the implementation should: - Skip write-back when `save_interactions=False`. - Use the explicit `user_id` constructor argument. - Keep the agent usable when PowerMem search fails. +- Support both `invoke()` and `ainvoke()` agent execution. + +`user_id` is required and must be supplied explicitly. This keeps application +identity separate from LangChain runtime or checkpoint state. Retrieved memory +is added to the model request's system context without modifying the persisted +conversation messages. Interaction write-back uses `infer=False`, preserving +the final user and assistant messages as explicit PowerMem records. ## Tests @@ -112,6 +119,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, an answer informed +by those memories, and updated memories after the call. diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..14807d0460 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -1,31 +1,38 @@ -"""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 that provides PowerMem-backed long-term memory.""" from __future__ import annotations +import inspect +import logging +from asyncio import to_thread +from collections.abc import Awaitable, Callable, Mapping, Sequence 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, BaseMessage, HumanMessage, SystemMessage + +logger = logging.getLogger(__name__) class PowerMemState(AgentState): - """State schema reserved for the PowerMem middleware implementation.""" + """Agent state used to carry retrieved memory into model-call hooks.""" powermem_context: NotRequired[str] class PowerMemStateUpdate(TypedDict): - """State update returned by memory-loading middleware hooks.""" + """State update returned after loading relevant memories.""" powermem_context: str class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" + """Use a PowerMem client as a LangChain agent's long-term memory layer.""" state_schema = PowerMemState @@ -38,17 +45,225 @@ def __init__( save_interactions: bool = True, **kwargs: Any, ) -> None: - pass + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected PowerMemMiddleware arguments: {unexpected}") + if memory is None: + raise ValueError("memory must be a PowerMem client") + if not isinstance(user_id, str) or not user_id.strip(): + raise ValueError("user_id must be a non-empty explicit user identity") + if isinstance(search_limit, bool) or not isinstance(search_limit, int): + raise TypeError("search_limit must be an integer") + if search_limit < 1: + raise ValueError("search_limit must be at least 1") + + self.memory = memory + self.user_id = user_id + self.search_limit = search_limit + self.save_interactions = save_interactions def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + del runtime + query = _latest_message_text(state.get("messages", []), HumanMessage) + if not query: + return None + + try: + result = self.memory.search( + query=query, + user_id=self.user_id, + limit=self.search_limit, + ) + if inspect.isawaitable(result): + _close_awaitable(result) + raise TypeError( + "The configured PowerMem client is asynchronous; use agent.ainvoke()" + ) + except Exception: + logger.warning("PowerMem search failed; continuing without memory", exc_info=True) + return None + + context = _format_search_result(result) + return {"powermem_context": context} if context else None async def abefore_agent( self, state: PowerMemState, runtime, ) -> PowerMemStateUpdate | None: - pass + del runtime + query = _latest_message_text(state.get("messages", []), HumanMessage) + if not query: + return None + + try: + result = await _call_async_compatible( + self.memory.search, + query=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 = _format_search_result(result) + return {"powermem_context": context} if context else None + + def wrap_model_call( + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], ModelResponse[Any]], + ) -> ModelResponse[Any]: + return handler(_inject_memory_context(request)) + + async def awrap_model_call( + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], Awaitable[ModelResponse[Any]]], + ) -> ModelResponse[Any]: + return await handler(_inject_memory_context(request)) def after_agent(self, state: PowerMemState, runtime) -> None: - pass + del runtime + if not self.save_interactions: + return + + interaction = _final_interaction(state.get("messages", [])) + if interaction is None: + return + + result = self.memory.add( + messages=interaction, + user_id=self.user_id, + infer=False, + ) + if inspect.isawaitable(result): + _close_awaitable(result) + raise TypeError( + "The configured PowerMem client is asynchronous; use agent.ainvoke()" + ) + + async def aafter_agent(self, state: PowerMemState, runtime) -> None: + del runtime + if not self.save_interactions: + return + + interaction = _final_interaction(state.get("messages", [])) + if interaction is None: + return + + await _call_async_compatible( + self.memory.add, + messages=interaction, + user_id=self.user_id, + infer=False, + ) + + +async def _call_async_compatible(function: Callable[..., Any], **kwargs: Any) -> Any: + """Call native async methods directly and move sync SDK calls off the event loop.""" + + if inspect.iscoroutinefunction(function): + return await function(**kwargs) + + result = await to_thread(function, **kwargs) + if inspect.isawaitable(result): + return await result + return result + + +def _close_awaitable(awaitable: Any) -> None: + close = getattr(awaitable, "close", None) + if callable(close): + close() + + +def _latest_message_text( + messages: Sequence[Any], + message_type: type[BaseMessage], +) -> str | None: + for message in reversed(messages): + if isinstance(message, message_type): + text = _message_text(message) + if text: + return text + return None + + +def _message_text(message: BaseMessage) -> str: + text = message.text + if isinstance(text, str): + return text.strip() + return str(text).strip() + + +def _format_search_result(result: Any) -> str: + if not isinstance(result, Mapping): + return "" + + raw_memories = result.get("results", []) + if not isinstance(raw_memories, Sequence) or isinstance(raw_memories, (str, bytes)): + return "" + + memories: list[str] = [] + for item in raw_memories: + if isinstance(item, Mapping): + value = item.get("memory") or item.get("content") + else: + value = item + if value is not None and (text := str(value).strip()): + memories.append(text) + + if not memories: + return "" + + lines = [ + "Relevant long-term memories from PowerMem:", + "Use these only when they are relevant to the current request.", + ] + lines.extend(f"- {memory}" for memory in memories) + return "\n".join(lines) + + +def _inject_memory_context(request: ModelRequest[Any]) -> ModelRequest[Any]: + context = request.state.get("powermem_context") + if not context: + return request + + existing = request.system_message + if existing is None: + system_message = SystemMessage(content=context) + elif isinstance(existing.content, str): + separator = "\n\n" if existing.content else "" + system_message = existing.model_copy( + update={"content": f"{existing.content}{separator}{context}"} + ) + else: + content = [*existing.content, {"type": "text", "text": context}] + system_message = existing.model_copy(update={"content": content}) + + return request.override(system_message=system_message) + + +def _final_interaction(messages: Sequence[Any]) -> list[dict[str, str]] | None: + assistant_index: int | None = None + assistant_text: str | None = None + for index in range(len(messages) - 1, -1, -1): + message = messages[index] + if isinstance(message, AIMessage) and (text := _message_text(message)): + assistant_index = index + assistant_text = text + break + + if assistant_index is None or assistant_text is None: + return None + + user_text = _latest_message_text(messages[:assistant_index], HumanMessage) + if not user_text: + return None + + return [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": assistant_text}, + ] diff --git a/packages/powermem-langchain/tests/test_middleware.py b/packages/powermem-langchain/tests/test_middleware.py index 4f2fbbcfd9..07bbf0b6d2 100644 --- a/packages/powermem-langchain/tests/test_middleware.py +++ b/packages/powermem-langchain/tests/test_middleware.py @@ -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 BaseMessage, HumanMessage, SystemMessage from powermem import Memory from powermem_langchain import PowerMemMiddleware from pydantic import Field @@ -38,6 +38,21 @@ def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: raise RuntimeError("search failed") +class RecordingMemory: + def __init__(self, results: list[dict[str, Any]] | None = None) -> None: + self.results = results or [] + self.search_calls: list[dict[str, Any]] = [] + self.add_calls: list[dict[str, Any]] = [] + + def search(self, **kwargs: Any) -> dict[str, Any]: + self.search_calls.append(kwargs) + return {"results": self.results} + + def add(self, **kwargs: Any) -> dict[str, Any]: + self.add_calls.append(kwargs) + return {"results": []} + + def _message_text(messages: list[BaseMessage]) -> str: return "\n".join(str(message.content) for message in messages) @@ -162,6 +177,67 @@ def test_search_failure_is_fail_open_by_default(): assert result["messages"][-1].content == "Agent still runs" +def test_uses_latest_user_message_and_explicit_user_id(): + memory = RecordingMemory() + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="explicit-user", + search_limit=3, + ) + ], + ) + + agent.invoke( + { + "messages": [ + HumanMessage(content="Older question"), + HumanMessage(content="Newest question"), + ] + } + ) + + assert memory.search_calls == [ + {"query": "Newest question", "user_id": "explicit-user", "limit": 3} + ] + assert memory.add_calls[0]["user_id"] == "explicit-user" + + +def test_preserves_existing_system_prompt_when_injecting_memory(): + memory = RecordingMemory([{"memory": "PowerMem context"}]) + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + system_prompt="Original system prompt", + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=False, + ) + ], + ) + + agent.invoke({"messages": [HumanMessage(content="Question")]}) + + system_messages = [ + message for message in model.calls[0] if isinstance(message, SystemMessage) + ] + assert len(system_messages) == 1 + assert "Original system prompt" in str(system_messages[0].content) + assert "PowerMem context" in str(system_messages[0].content) + + +def test_rejects_missing_user_id(): + with pytest.raises(ValueError, match="user_id"): + PowerMemMiddleware(memory=RecordingMemory()) + + @pytest.mark.asyncio async def test_async_agent_uses_powermem_memory(tmp_path: Path): memory = _sqlite_memory(tmp_path) @@ -183,3 +259,27 @@ 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(): + memory = RecordingMemory() + model = CapturingChatModel(responses=["Async response"]) + agent = create_agent( + model=model, + tools=[], + middleware=[PowerMemMiddleware(memory=memory, user_id="async-user")], + ) + + await agent.ainvoke({"messages": [HumanMessage(content="Async request")]}) + + assert memory.add_calls == [ + { + "messages": [ + {"role": "user", "content": "Async request"}, + {"role": "assistant", "content": "Async response"}, + ], + "user_id": "async-user", + "infer": False, + } + ]