From 1c452a8994ccfb4b2caaa9c83398fd9a1687158d Mon Sep 17 00:00:00 2001 From: yanngfrye Date: Fri, 17 Jul 2026 09:59:36 +0800 Subject: [PATCH] problem1 --- .../src/powermem_langchain/middleware.py | 173 +++++++++++++++--- 1 file changed, 150 insertions(+), 23 deletions(-) diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f..266c9a308 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -1,31 +1,56 @@ """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. +Provides PowerMemMiddleware, a LangChain v1 AgentMiddleware that integrates +PowerMem as a long-term memory layer for agents. """ from __future__ import annotations -from typing import Any, NotRequired, TypedDict +import logging +from typing import Any from langchain.agents.middleware import AgentMiddleware, AgentState +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +logger = logging.getLogger(__name__) -class PowerMemState(AgentState): - """State schema reserved for the PowerMem middleware implementation.""" - - powermem_context: NotRequired[str] +class PowerMemState(AgentState): + """State schema for the PowerMem middleware. -class PowerMemStateUpdate(TypedDict): - """State update returned by memory-loading middleware hooks.""" + Carries the retrieved memory context through the agent graph so that + downstream hooks (e.g. ``before_model``) can inject it into the prompt. + """ powermem_context: str class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" + """LangChain v1 middleware that connects an agent to PowerMem. + + On every agent invocation the middleware: + + * Searches PowerMem for memories relevant to the latest user message + **before** the model is called. + * Injects the retrieved memories into the conversation as a system + message so the model can see them. + * After the agent finishes, persists the user message and assistant + response back to PowerMem (when ``save_interactions=True``). + + Parameters + ---------- + memory: + A PowerMem ``Memory`` instance (or any object exposing a compatible + ``search(query, user_id, limit)`` and ``add(messages, user_id, infer)`` + API). + user_id: + Explicit user identity used for both search and persistence. + search_limit: + Maximum number of memories to retrieve per search (default 5). + save_interactions: + When ``True`` (the default), the user message and assistant response + are written back to PowerMem after the agent completes. + """ state_schema = PowerMemState @@ -38,17 +63,119 @@ def __init__( save_interactions: bool = True, **kwargs: Any, ) -> None: - pass - - def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + self._memory = memory + self._user_id = user_id + self._search_limit = search_limit + self._save_interactions = save_interactions + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _get_latest_user_message(state: PowerMemState) -> str: + """Return the content of the most recent ``HumanMessage`` in state.""" + messages = state.get("messages", []) + for msg in reversed(messages): + if isinstance(msg, HumanMessage): + return str(msg.content) + return "" + + def _search_memories(self, query: str) -> str: + """Search PowerMem and return formatted memory lines. + + Fail-open: returns an empty string when the search fails so the + agent can continue without interruption. + """ + try: + result = self._memory.search( + query=query, + user_id=self._user_id, + limit=self._search_limit, + ) + except Exception: + logger.warning("PowerMem search failed; continuing without memories.", exc_info=True) + return "" + + memories = result.get("results", []) if isinstance(result, dict) else [] + if not memories: + return "" + + lines: list[str] = [] + for mem in memories: + memory_text = mem.get("memory", "") if isinstance(mem, dict) else "" + if memory_text: + lines.append(f"- {memory_text}") + return "\n".join(lines) + + def _persist_interaction(self, state: PowerMemState) -> None: + """Persist user and assistant messages to PowerMem when enabled.""" + if not self._save_interactions: + return + + messages = state.get("messages", []) + user_parts: list[str] = [] + assistant_parts: list[str] = [] + + for msg in messages: + if isinstance(msg, HumanMessage): + user_parts.append(str(msg.content)) + elif isinstance(msg, AIMessage): + assistant_parts.append(str(msg.content)) + + if not user_parts and not assistant_parts: + return + + combined = ( + "User: " + + " | ".join(user_parts) + + "\nAssistant: " + + " | ".join(assistant_parts) + ) + try: + self._memory.add(combined, user_id=self._user_id, infer=False) + except Exception: + logger.warning( + "PowerMem add failed; interaction was not persisted.", exc_info=True + ) + + # ------------------------------------------------------------------ + # Middleware hooks + # ------------------------------------------------------------------ + + def before_agent( + self, state: PowerMemState, runtime: Any + ) -> dict[str, Any] | None: + """Search PowerMem and inject relevant memories before the agent runs. + + Returns a state update that appends a ``SystemMessage`` with the + retrieved memories so the model can reference them. + """ + query = self._get_latest_user_message(state) + if not query: + return {"powermem_context": ""} + + context = self._search_memories(query) + updates: dict[str, Any] = {"powermem_context": context} + + if context: + system_msg = SystemMessage( + content=f"Relevant information from previous conversations:\n{context}" + ) + updates["messages"] = [system_msg] + + return updates async def abefore_agent( - self, - state: PowerMemState, - runtime, - ) -> PowerMemStateUpdate | None: - pass - - def after_agent(self, state: PowerMemState, runtime) -> None: - pass + self, state: PowerMemState, runtime: Any + ) -> dict[str, Any] | None: + """Async variant of :meth:`before_agent`.""" + return self.before_agent(state, runtime) + + def after_agent(self, state: PowerMemState, runtime: Any) -> None: + """Persist the completed interaction to PowerMem.""" + self._persist_interaction(state) + + async def aafter_agent(self, state: PowerMemState, runtime: Any) -> None: + """Async variant of :meth:`after_agent`.""" + self.after_agent(state, runtime)