diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..b668fefa39 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -1,31 +1,37 @@ -"""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 asyncio +import logging 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 extended with formatted long-term-memory context.""" powermem_context: NotRequired[str] class PowerMemStateUpdate(TypedDict): - """State update returned by memory-loading middleware hooks.""" + """State update returned after loading memories.""" powermem_context: str class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" + """Retrieve PowerMem context and persist completed agent interactions.""" state_schema = PowerMemState @@ -38,17 +44,109 @@ def __init__( save_interactions: bool = True, **kwargs: Any, ) -> None: - pass - - def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + if memory is None: + raise ValueError("memory must not be None") + if not isinstance(user_id, str) or not user_id.strip(): + raise ValueError("user_id must be a non-empty string") + if not isinstance(search_limit, int) or isinstance(search_limit, bool) or search_limit <= 0: + raise ValueError("search_limit must be a positive integer") + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"unexpected middleware arguments: {names}") + + self.memory = memory + self.user_id = user_id + self.search_limit = search_limit + self.save_interactions = save_interactions + + @staticmethod + def _message_text(message: BaseMessage) -> str: + return message.text.strip() + + @classmethod + def _latest_message(cls, state: PowerMemState, message_type: type[BaseMessage]) -> str | None: + for message in reversed(state.get("messages", [])): + if isinstance(message, message_type): + text = cls._message_text(message) + if text: + return text + return None + + @staticmethod + def _format_memories(result: Any) -> str: + if not isinstance(result, dict): + return "" + memories: list[str] = [] + for item in result.get("results", []): + if isinstance(item, dict): + value = item.get("memory") or item.get("content") + if value is not None and str(value).strip(): + memories.append(str(value).strip()) + if not memories: + return "" + body = "\n".join(f"- {memory}" for memory in memories) + return f"Relevant long-term memories for this user:\n{body}" + + def _load_context(self, state: PowerMemState) -> str: + query = self._latest_message(state, HumanMessage) + if not query: + return "" + 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 "" + return self._format_memories(result) + + def before_agent(self, state: PowerMemState, runtime: Any) -> PowerMemStateUpdate: + return {"powermem_context": self._load_context(state)} async def abefore_agent( self, state: PowerMemState, - runtime, - ) -> PowerMemStateUpdate | None: - pass - - def after_agent(self, state: PowerMemState, runtime) -> None: - pass + runtime: Any, + ) -> PowerMemStateUpdate: + context = await asyncio.to_thread(self._load_context, state) + return {"powermem_context": context} + + @staticmethod + def _request_with_context(request: ModelRequest[Any]) -> ModelRequest[Any]: + context = request.state.get("powermem_context", "") + if not context: + return request + return request.override(messages=[SystemMessage(content=context), *request.messages]) + + def wrap_model_call(self, request: ModelRequest[Any], handler: Any) -> ModelResponse: + return handler(self._request_with_context(request)) + + async def awrap_model_call(self, request: ModelRequest[Any], handler: Any) -> ModelResponse: + return await handler(self._request_with_context(request)) + + def _save_interaction(self, state: PowerMemState) -> None: + if not self.save_interactions: + return + user_text = self._latest_message(state, HumanMessage) + assistant_text = self._latest_message(state, AIMessage) + if not user_text or not assistant_text: + return + try: + self.memory.add( + [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": assistant_text}, + ], + user_id=self.user_id, + infer=False, + ) + except Exception: + logger.warning("PowerMem write-back failed; returning agent result", exc_info=True) + + def after_agent(self, state: PowerMemState, runtime: Any) -> None: + self._save_interaction(state) + + async def aafter_agent(self, state: PowerMemState, runtime: Any) -> None: + await asyncio.to_thread(self._save_interaction, state)