diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..218265eef2 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -1,32 +1,23 @@ -"""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 from langchain.agents.middleware import AgentMiddleware, AgentState +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage class PowerMemState(AgentState): - """State schema reserved for the PowerMem middleware implementation.""" - powermem_context: NotRequired[str] class PowerMemStateUpdate(TypedDict): - """State update returned by memory-loading middleware hooks.""" - - powermem_context: str + powermem_context: NotRequired[str] + messages: NotRequired[list[BaseMessage]] class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" - state_schema = PowerMemState def __init__( @@ -38,17 +29,109 @@ def __init__( save_interactions: bool = True, **kwargs: Any, ) -> None: - pass + super().__init__(**kwargs) + self.memory = memory + self.user_id = user_id or "default" + self.search_limit = search_limit + self.save_interactions = save_interactions def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + latest_user_message = self._latest_message_text(state, HumanMessage) + if not latest_user_message: + return None + + try: + result = self.memory.search( + latest_user_message, + user_id=self.user_id, + limit=self.search_limit, + ) + except Exception: + return None + + memories = self._extract_memory_texts(result) + if not memories: + return None + + context = "Relevant long-term memories:\n" + "\n".join( + f"- {memory}" for memory in memories + ) + + return { + "powermem_context": context, + "messages": [SystemMessage(content=context)], + } async def abefore_agent( self, state: PowerMemState, runtime, ) -> PowerMemStateUpdate | None: - pass + return self.before_agent(state, runtime) def after_agent(self, state: PowerMemState, runtime) -> None: - pass + if not self.save_interactions: + return + + user_text = self._latest_message_text(state, HumanMessage) + assistant_text = self._latest_message_text(state, AIMessage) + + if not user_text or not assistant_text: + return + + interaction = f"User: {user_text}\nAssistant: {assistant_text}" + + try: + self.memory.add(interaction, user_id=self.user_id, infer=False) + except Exception: + return + + @classmethod + def _latest_message_text( + cls, + state: PowerMemState, + message_type: type[BaseMessage], + ) -> str: + for message in reversed(state.get("messages", [])): + if isinstance(message, message_type): + return cls._content_to_text(message.content) + return "" + + @staticmethod + def _content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + + if isinstance(content, list): + parts = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict): + parts.append(str(item.get("text") or item.get("content") or item)) + else: + parts.append(str(item)) + return "\n".join(parts) + + return str(content) + + @classmethod + def _extract_memory_texts(cls, result: Any) -> list[str]: + if isinstance(result, dict): + items = result.get("results", []) + elif isinstance(result, list): + items = result + else: + items = [] + + memories: list[str] = [] + for item in items: + if isinstance(item, dict): + text = item.get("memory") or item.get("content") or item.get("text") + else: + text = getattr(item, "memory", None) or getattr(item, "content", None) + + if text: + memories.append(cls._content_to_text(text)) + + return memories