From 5d3617dca692acfc4fdf663aa7194e1ada6167ea Mon Sep 17 00:00:00 2001 From: saber1231-lab Date: Wed, 15 Jul 2026 14:26:00 +0800 Subject: [PATCH] feat: implement PowerMem LangChain middleware --- .../src/powermem_langchain/middleware.py | 140 ++++++++++++++++-- .../tests/test_middleware.py | 91 +++++++++++- 2 files changed, 217 insertions(+), 14 deletions(-) diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..568dab1e9b 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/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)}, + ] diff --git a/packages/powermem-langchain/tests/test_middleware.py b/packages/powermem-langchain/tests/test_middleware.py index 4f2fbbcfd9..4f496aa6cc 100644 --- a/packages/powermem-langchain/tests/test_middleware.py +++ b/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", + ) + ]