feat: implement powermem langchain middleware

packages/powermem-langchain/src/powermem_langchain/middleware.py

@@ -7,9 +7,15 @@
 
 from __future__ import annotations
 
+import logging
+from collections.abc import Sequence
 from typing import Any, NotRequired, TypedDict
 
 from langchain.agents.middleware import AgentMiddleware, AgentState
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+
+
+logger = logging.getLogger(__name__)
 
 
 class PowerMemState(AgentState):
@@ -18,14 +24,15 @@ class PowerMemState(AgentState):
     powermem_context: NotRequired[str]
 
 
-class PowerMemStateUpdate(TypedDict):
+class PowerMemStateUpdate(TypedDict, total=False):
     """State update returned by memory-loading middleware hooks."""
 
     powermem_context: str
+    messages: list[BaseMessage]
 
 
 class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]):
-    """Placeholder for the summer school implementation."""
+    """LangChain middleware that gives an agent access to PowerMem memories."""
 
     state_schema = PowerMemState
 
@@ -38,17 +45,124 @@ 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
+        self.kwargs = kwargs
 
     def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None:
-        pass
+        query = self._latest_message_text(state.get("messages", ()), HumanMessage)
+        if not query:
+            return None
+
+        context = self._search_context(query)
+        if not context:
+            return None
+
+        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
+
+        messages = state.get("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
+
+        memory_text = f"User: {user_message}\nAssistant: {assistant_message}"
+        self.memory.add(memory_text, user_id=self.user_id, infer=False)
+
+    def _search_context(self, query: str) -> str | None:
+        try:
+            result = self.memory.search(
+                query,
+                user_id=self.user_id,
+                limit=self.search_limit,
+            )
+        except Exception:
+            logger.exception("PowerMem search failed; continuing without memories.")
+            return None
+
+        memories = self._memory_items(result)
+        if not memories:
+            return None
+
+        lines = ["Relevant memories from PowerMem:"]
+        lines.extend(f"- {memory}" for memory in memories)
+        return "\n".join(lines)
+
+    def _memory_items(self, result: Any) -> list[str]:
+        if not isinstance(result, dict):
+            return []
+
+        items = result.get("results", [])
+        if not isinstance(items, Sequence) or isinstance(items, str):
+            return []
+
+        memories: list[str] = []
+        for item in items[: self.search_limit]:
+            if isinstance(item, dict):
+                content = item.get("memory") or item.get("content")
+            else:
+                content = str(item)
+
+            text = self._content_to_text(content)
+            if text:
+                memories.append(text)
+
+        return memories
+
+    def _latest_message_text(
+        self,
+        messages: Sequence[Any],
+        message_type: type[BaseMessage],
+    ) -> str | None:
+        for message in reversed(messages):
+            if isinstance(message, message_type):
+                text = self._content_to_text(message.content)
+                if text:
+                    return text
+            elif isinstance(message, dict):
+                role = str(message.get("role") or "").lower()
+                if (
+                    message_type is HumanMessage
+                    and role in {"human", "user"}
+                    or message_type is AIMessage
+                    and role in {"ai", "assistant"}
+                ):
+                    text = self._content_to_text(message.get("content"))
+                    if text:
+                        return text
+
+        return None
+
+    def _content_to_text(self, content: Any) -> str:
+        if content is None:
+            return ""
+        if isinstance(content, str):
+            return content.strip()
+        if isinstance(content, Sequence) and not isinstance(content, (bytes, bytearray)):
+            parts: list[str] = []
+            for item in content:
+                if isinstance(item, dict):
+                    value = item.get("text") or item.get("content")
+                    if value:
+                        parts.append(str(value))
+                elif item is not None:
+                    parts.append(str(item))
+            return "\n".join(parts).strip()
+
+        return str(content).strip()