aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/middleware/memory.py

@@ -0,0 +1,532 @@
+"""Memory middleware for automatic recall and persistence.
+
+This module provides MemoryMiddleware, which implements the AgentMiddleware protocol
+to handle memory operations (recall and persistence) through the middleware lifecycle
+hooks instead of special graph nodes.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+from concurrent.futures import Future
+from typing import TYPE_CHECKING, Any
+
+from langchain_core.messages import HumanMessage
+
+from aip_agents.constants import TEXT_PREVIEW_LENGTH
+from aip_agents.middleware.base import AgentMiddleware
+from aip_agents.tools.memory_search import Mem0DeleteTool, Mem0SearchTool
+from aip_agents.utils.logger import get_logger
+
+if TYPE_CHECKING:  # pragma: no cover
+    from aip_agents.memory import BaseMemory
+
+logger = get_logger(__name__)
+
+MAX_QUERY_LENGTH: int = 128
+DEFAULT_MEMORY_LIMIT: int = 10
+FORGET_WORD_PATTERN: str = r"\bforget\b"
+
+
+class MemoryMiddleware(AgentMiddleware):
+    """Middleware for automatic memory recall and persistence.
+
+    This middleware replaces the dedicated memory graph entry node by using
+    middleware lifecycle hooks to:
+    1. Recall relevant memories before the agent runs (abefore_run hook)
+    2. Persist interactions after the agent completes (aafter_run hook)
+
+    The middleware handles memory retrieval logic directly, maintaining backward
+    compatibility with existing memory configuration.
+
+    Attributes:
+        tools: Empty list - memory middleware doesn't contribute tools.
+        system_prompt_additions: None - no prompt additions needed.
+    """
+
+    tools: list[Any] = []
+    system_prompt_additions: str | None = None
+
+    def __init__(
+        self,
+        memory: BaseMemory | None,
+        memory_agent_id: str,
+        save_interaction_to_memory: bool = True,
+        memory_retrieval_limit: int = DEFAULT_MEMORY_LIMIT,
+    ) -> None:
+        """Initialize the MemoryMiddleware.
+
+        Args:
+            memory: The memory backend instance (Mem0Memory or compatible).
+            memory_agent_id: The agent ID used for memory scoping.
+            save_interaction_to_memory: Whether to save interactions to memory.
+            memory_retrieval_limit: Maximum memories to retrieve per recall.
+        """
+        self._memory = memory
+        self._memory_agent_id = memory_agent_id
+        self._save_interaction_to_memory = save_interaction_to_memory
+        self._memory_retrieval_limit = max(memory_retrieval_limit, 1)
+
+        if memory is not None:
+            self._memory_search_tool = Mem0SearchTool(memory=memory, default_user_id=memory_agent_id)
+            self._memory_delete_tool = Mem0DeleteTool(memory=memory, default_user_id=memory_agent_id)
+        else:
+            self._memory_search_tool = None
+            self._memory_delete_tool = None
+
+    @property
+    def memory_enabled(self) -> bool:
+        """Check whether memory is enabled.
+
+        Returns:
+            True when a memory adapter is set.
+        """
+        return self._memory is not None
+
+    async def abefore_run(
+        self,
+        state: dict[str, Any],
+        config: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        """Execute memory recall before the agent run.
+
+        Extracts the user query from state messages, retrieves relevant memories
+        directly using the memory search tool, and enhances the query with
+        formatted memory context.
+
+        Args:
+            state: Initial agent state containing messages.
+            config: Optional run configuration.
+
+        Returns:
+            State updates with enhanced message if memories were found.
+        """
+        if not self.memory_enabled:
+            return {}
+
+        user_query = self._extract_user_query_from_messages(state.get("messages", []))
+        if not user_query:
+            return {}
+
+        self._store_original_query(config, user_query)
+
+        try:
+            memory_user_id = self._extract_memory_user_id(config)
+            effective_user_id = memory_user_id or self._memory_agent_id
+
+            if self._is_forget_request(user_query):
+                await self._delete_memories(
+                    user_query=user_query, user_id=effective_user_id, state=state, config=config
+                )
+                return {"memory_action": "delete"}
+
+            memories = await self._retrieve_memories(user_query, effective_user_id, state, config)
+            if not memories:
+                return {}
+
+            tagged_memory = self._format_memories(memories)
+            enhanced_query = f"{user_query}\n\n{tagged_memory}".strip()
+            logger.info("MemoryMiddleware: enhanced query with %d memories", len(memories))
+
+            return {"messages": [HumanMessage(content=enhanced_query)]}
+
+        except Exception as e:
+            logger.warning("MemoryMiddleware: enhancement failed: %s", e)
+            return {}
+
+    async def _retrieve_memories(
+        self,
+        query: str,
+        user_id: str,
+        state: dict[str, Any],
+        config: dict[str, Any] | None,
+    ) -> list[dict[str, Any]]:
+        """Retrieve memories using the memory search tool."""
+        if self._memory_search_tool is None:
+            return []  # pragma: no cover
+
+        metadata = self._build_metadata(state, user_id)
+        search_args = {
+            "query": query[:MAX_QUERY_LENGTH] if query else None,
+            "limit": self._memory_retrieval_limit,
+            "metadata": metadata,
+        }
+        tool_config = self._create_tool_config(config, user_id)
+
+        try:
+            result = await self._memory_search_tool.ainvoke(search_args, config=tool_config)
+            memories = self._parse_memories_from_result(result)
+            if memories:
+                logger.info("MemoryMiddleware: retrieved %d memories", len(memories))
+            return memories
+        except Exception as e:  # pragma: no cover
+            logger.warning("MemoryMiddleware: retrieval failed: %s", e)  # pragma: no cover
+            return []  # pragma: no cover
+
+    def _create_tool_config(
+        self,
+        config: dict[str, Any] | None,
+        user_id: str,
+    ) -> dict[str, Any]:
+        """Create tool configuration with user_id."""
+        tool_config = dict(config) if config else {}
+        tool_config.setdefault("configurable", {})
+        tool_config["configurable"]["user_id"] = user_id
+        return tool_config
+
+    def _build_metadata(self, state: dict[str, Any], user_id: str) -> dict[str, Any]:
+        """Build metadata dict for memory tools, filtering internal keys."""
+        state_metadata = state.get("metadata", {})
+        metadata = {k: v for k, v in state_metadata.items() if k not in {"tool_configs", "configurable", "thread_id"}}
+        metadata["user_id"] = user_id
+        return metadata
+
+    def _parse_memories_from_result(self, result: Any) -> list[dict[str, Any]]:
+        """Parse memory list from tool result."""
+        if isinstance(result, list):
+            return [m for m in result if isinstance(m, dict)]
+        if isinstance(result, str):
+            try:
+                parsed = json.loads(result)
+                if isinstance(parsed, list):
+                    return [m for m in parsed if isinstance(m, dict)]
+            except json.JSONDecodeError:
+                pass
+        return []
+
+    def _format_memories(self, memories: list[dict[str, Any]]) -> str:
+        """Format memory hits using the underlying tool formatter."""
+        if not memories:
+            return ""
+        if self._memory_search_tool is not None:
+            return self._memory_search_tool.format_hits(memories, with_tag=True)
+        lines = ["<MEMORY_CONTEXT>"]
+        for i, memory in enumerate(memories, 1):
+            content = memory.get("memory", str(memory))
+            lines.append(f"{i}. {content}")
+        lines.append("</MEMORY_CONTEXT>")
+        return "\n".join(lines)
+
+    def _is_forget_request(self, user_query: str) -> bool:
+        """Return True when the user request should delete/clear memories."""
+        normalized = user_query.strip().lower()
+        if not normalized:
+            return False
+
+        # Guard against reminder phrasing such as "don't forget to ...".
+        if self._is_non_destructive_forget_phrase(normalized):
+            return False
+
+        # Explicit forget intent (imperative/request form).
+        if re.search(r"^(please\s+)?((can|could|would)\s+you\s+)?forget\b", normalized):
+            return True
+
+        # Accept intent phrasing where "forget" is not sentence-initial.
+        if re.search(r"\bi\s+(want|need|would\s+like)\s+to\s+forget\b", normalized):
+            return True
+
+        # Allow explicit memory management phrasing.
+        if re.search(r"\b(clear|delete|remove)\b", normalized) and self._has_memory_context(normalized):
+            return True
+
+        return False
+
+    def _extract_delete_query(self, user_query: str) -> str:
+        """Extract the semantic target for deletion from a forget request."""
+        normalized = user_query.strip()
+
+        # Strip common prefixes so we delete the actual target content.
+        normalized = re.sub(r"(?i)^\s*(please\s+)?(forget|delete|remove|clear)\b\s*", "", normalized)
+        normalized = re.sub(r"(?i)^\s*(that\s+)?", "", normalized)
+        return normalized.strip() or user_query.strip()
+
+    @staticmethod
+    def _is_non_destructive_forget_phrase(normalized_query: str) -> bool:
+        """Return True for reminder phrasing that should not delete memory."""
+        non_destructive_patterns = (
+            r"\bdon['’]?t\s+forget\b",
+            r"\bdo\s+not\s+forget\b",
+            r"\bnever\s+forget\b",
+            r"\bforget\s+to\b",
+        )
+        return any(re.search(pattern, normalized_query) for pattern in non_destructive_patterns)
+
+    @staticmethod
+    def _has_memory_context(normalized_query: str) -> bool:
+        """Return True when the query refers to memory explicitly or implicitly."""
+        return bool(
+            re.search(
+                r"\bmemory\b|\bmemories\b|\byou\s+remember\b|\bremember\s+about\s+me\b|\bwhat\s+you\s+know\s+about\s+me\b",
+                normalized_query,
+            )
+        )
+
+    @staticmethod
+    def _is_scoped_forget_query(normalized_query: str) -> bool:
+        """Return True for scoped forget requests that should not map to delete_all."""
+        if not re.search(FORGET_WORD_PATTERN, normalized_query):
+            return False
+        if not re.search(r"\b(all|everything)\b", normalized_query):
+            return False
+        if re.search(r"\babout\s+me\b", normalized_query):
+            return False
+        return bool(re.search(r"\b(all|everything)\b(?:\s+\w+){0,5}\s+\b(about|regarding|on)\b", normalized_query))
+
+    def _is_delete_all_request(self, user_query: str) -> bool:
+        normalized = user_query.strip().lower()
+        if self._is_non_destructive_forget_phrase(normalized):
+            return False
+
+        has_delete_verb = bool(re.search(r"\b(forget|delete|clear|remove)\b", normalized))
+        has_all_scope = bool(re.search(r"\b(all|everything)\b", normalized))
+        if not (has_delete_verb and has_all_scope):
+            return False
+
+        if self._is_scoped_forget_query(normalized):
+            return False
+
+        if self._has_memory_context(normalized):
+            return True
+
+        # Keep concise global intent working (e.g. "forget everything").
+        return bool(re.search(FORGET_WORD_PATTERN, normalized))
+
+    async def _delete_memories(
+        self,
+        *,
+        user_query: str,
+        user_id: str,
+        state: dict[str, Any],
+        config: dict[str, Any] | None,
+    ) -> None:
+        """Best-effort deletion for forget/clear requests."""
+        if self._memory_delete_tool is None:
+            return  # pragma: no cover
+
+        await self._flush_pending_memory_writes(timeout=30.0)
+
+        tool_config = self._create_tool_config(config, user_id)
+        metadata = self._build_metadata(state, user_id)
+
+        if self._is_delete_all_request(user_query):
+            args: dict[str, Any] = {"delete_all": True, "metadata": metadata}
+        else:
+            delete_query = self._extract_delete_query(user_query)[:MAX_QUERY_LENGTH]
+            args = {"query": delete_query, "metadata": metadata}
+
+        try:
+            await self._memory_delete_tool.ainvoke(args, config=tool_config)
+        except Exception as exc:  # noqa: BLE001  # pragma: no cover
+            logger.warning("MemoryMiddleware: delete ignored error: %s", exc)  # pragma: no cover
+
+    async def _flush_pending_memory_writes(self, timeout: float | None = None) -> None:
+        flush = getattr(self._memory, "flush_pending_writes", None)
+        if not callable(flush):
+            return
+        try:
+            await asyncio.to_thread(flush, timeout=timeout)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("MemoryMiddleware: pending memory flush ignored error: %s", exc)
+
+    async def aafter_run(
+        self,
+        *,
+        final_state: dict[str, Any],
+        output: str | None,
+        config: dict[str, Any] | None,
+        error: Exception | None,
+    ) -> None:
+        """Execute memory persistence after the agent run."""
+        if error is not None:
+            return
+        if not self._should_save_interaction(final_state):
+            return
+        if not output:
+            return  # pragma: no cover
+
+        user_query = self._extract_original_query(config=config, final_state=final_state)
+        if not user_query:
+            return
+
+        memory_user_id = self._extract_memory_user_id(config)
+        self._memory_save_interaction(
+            user_text=user_query,
+            ai_text=output,
+            memory_user_id=memory_user_id,
+        )
+
+    async def aon_final_response(self, content: str, context: dict[str, Any]) -> None:
+        """Persist memory early during streaming final response (best-effort)."""
+        if not self.memory_enabled or not self._save_interaction_to_memory:
+            return
+        if not isinstance(content, str) or not content:
+            return
+        if context.get("memory_persisted"):
+            return
+        final_state = context.get("final_state")
+        if isinstance(final_state, dict) and final_state.get("memory_action") == "delete":
+            context["memory_persisted"] = True
+            return
+        user_text = context.get("original_query")
+        if not isinstance(user_text, str) or not user_text:
+            return
+        config = context.get("config")
+        memory_user_id = self._extract_memory_user_id(config) if isinstance(config, dict) else None
+        if not memory_user_id:
+            memory_user_id = context.get("memory_user_id") if isinstance(context.get("memory_user_id"), str) else None
+        try:
+            self._memory_save_interaction(
+                user_text=user_text,
+                ai_text=content,
+                memory_user_id=memory_user_id,
+            )
+        finally:
+            context["memory_persisted"] = True
+
+    def _should_save_interaction(self, final_state: dict[str, Any] | None) -> bool:
+        """Return True when interaction should be saved to memory."""
+        if not self._save_interaction_to_memory or not self.memory_enabled:
+            return False
+        if isinstance(final_state, dict) and final_state.get("memory_persisted"):
+            return False  # pragma: no cover
+        if isinstance(final_state, dict) and final_state.get("memory_action") == "delete":
+            return False  # pragma: no cover
+        if self._contains_memory_delete_action(final_state):
+            return False
+        return True
+
+    @staticmethod
+    def _contains_memory_delete_action(final_state: dict[str, Any] | None) -> bool:
+        """Return True when final state includes a delete memory action block."""
+        if not isinstance(final_state, dict):
+            return False  # pragma: no cover
+        messages = final_state.get("messages")
+        if not isinstance(messages, list):
+            return False  # pragma: no cover
+        for message in messages:
+            content = getattr(message, "content", None)
+            if not isinstance(content, str):
+                continue
+            if "<MEMORY_ACTION>" in content and "action=delete" in content:
+                return True
+        return False
+
+    @staticmethod
+    def _extract_user_query_from_messages(messages: list[Any]) -> str | None:
+        """Get latest user query string from messages."""
+        if not messages:
+            return None
+        for msg in reversed(messages):
+            if isinstance(msg, HumanMessage) and hasattr(msg, "content"):
+                content = msg.content
+                if isinstance(content, str) and content.strip():
+                    return content
+        return None
+
+    def _extract_memory_user_id(self, config: dict[str, Any] | None) -> str | None:
+        """Extract memory_user_id from config if present."""
+        if not isinstance(config, dict):
+            return None  # pragma: no cover
+        configurable = config.get("configurable", {})
+        if isinstance(configurable, dict):
+            user_id = configurable.get("memory_user_id")
+            if user_id:
+                return user_id
+        metadata = config.get("metadata", {})
+        if isinstance(metadata, dict):
+            user_id = metadata.get("memory_user_id")
+            if user_id:
+                return user_id
+        return None
+
+    @staticmethod
+    def _store_original_query(config: dict[str, Any] | None, user_query: str) -> None:
+        """Store original user query in config metadata for persistence."""
+        if not isinstance(config, dict) or not isinstance(user_query, str) or not user_query.strip():
+            return
+        metadata = config.get("metadata")
+        if metadata is None:
+            config["metadata"] = {"original_query": user_query}
+            return
+        if not isinstance(metadata, dict):
+            return
+        existing = metadata.get("original_query")
+        if isinstance(existing, str) and existing.strip():
+            return
+        metadata["original_query"] = user_query
+
+    def _extract_original_query(
+        self,
+        *,
+        config: dict[str, Any] | None,
+        final_state: dict[str, Any] | None,
+    ) -> str | None:
+        """Extract original user query, preferring config metadata over state messages."""
+        if isinstance(config, dict):
+            metadata = config.get("metadata")
+            if isinstance(metadata, dict):
+                original_query = metadata.get("original_query")
+                if isinstance(original_query, str) and original_query.strip():
+                    return original_query
+        if isinstance(final_state, dict):
+            state_original_query = final_state.get("original_query")
+            if isinstance(state_original_query, str) and state_original_query.strip():
+                return state_original_query
+            return self._extract_user_query_from_messages(final_state.get("messages", []))
+        return None
+
+    def _memory_save_interaction(
+        self,
+        user_text: str,
+        ai_text: str,
+        memory_user_id: str | None = None,
+    ) -> None:
+        """Persist the user/assistant pair (best-effort)."""
+        if not (self._save_interaction_to_memory and self.memory_enabled and user_text and ai_text):
+            return  # pragma: no cover
+
+        try:
+            user_id = memory_user_id or self._memory_agent_id
+
+            logger.info("MemoryMiddleware: saving interaction for user_id='%s'", user_id)
+            logger.debug(
+                "MemoryMiddleware: interaction preview - User: '%s...' AI: '%s...'",
+                user_text[:TEXT_PREVIEW_LENGTH],
+                ai_text[:TEXT_PREVIEW_LENGTH],
+            )
+
+            save_async = getattr(self._memory, "save_interaction_async", None)
+            if callable(save_async):
+                future = save_async(user_text=str(user_text), ai_text=str(ai_text), user_id=user_id)
+                self._watch_memory_future(future, user_id)
+            elif hasattr(self._memory, "save_interaction"):
+                self._memory.save_interaction(  # pragma: no cover
+                    user_text=str(user_text),
+                    ai_text=str(ai_text),
+                    user_id=user_id,
+                )
+            else:
+                logger.warning(  # pragma: no cover
+                    "MemoryMiddleware: save_interaction not available on adapter type=%s",
+                    type(self._memory).__name__,
+                )
+        except Exception as e:
+            logger.debug("MemoryMiddleware: save_interaction ignored error: %s", e)
+
+    @staticmethod
+    def _watch_memory_future(future: Any, user_id: str) -> None:
+        """Attach logging to async memory writes."""
+        if not isinstance(future, Future):
+            return
+
+        def _log_completion(done: Future) -> None:
+            exc = done.exception()
+            if exc:  # pragma: no cover
+                logger.warning("MemoryMiddleware: async save failed for user_id='%s': %s", user_id, exc)
+
+        future.add_done_callback(_log_completion)