MemoryOS 2.0.3__py3-none-any.whl

memos/mem_agent/deepsearch_agent.py

@@ -0,0 +1,391 @@
+"""
+Deep Search Agent implementation for MemOS.
+
+This module implements a sophisticated deep search agent that performs iterative
+query refinement and memory retrieval to provide comprehensive answers.
+"""
+
+import json
+import re
+
+from typing import TYPE_CHECKING, Any
+
+from memos.configs.mem_agent import DeepSearchAgentConfig
+from memos.llms.base import BaseLLM
+from memos.log import get_logger
+from memos.mem_agent.base import BaseMemAgent
+from memos.memories.textual.item import TextualMemoryItem
+from memos.memories.textual.tree import TreeTextMemory
+from memos.templates.mem_agent_prompts import (
+    FINAL_GENERATION_PROMPT,
+    QUERY_REWRITE_PROMPT,
+    REFLECTION_PROMPT,
+)
+
+
+if TYPE_CHECKING:
+    from memos.types import MessageList
+
+logger = get_logger(__name__)
+
+
+class JSONResponseParser:
+    """Elegant JSON response parser for LLM outputs"""
+
+    @staticmethod
+    def parse(response: str) -> dict[str, Any]:
+        """Parse JSON response from LLM output with fallback strategies"""
+        # Clean response text by removing code block markers
+        cleaned = re.sub(r"^```(?:json)?\s*\n?|```\s*$", "", response.strip(), flags=re.IGNORECASE)
+
+        # Try parsing with multiple strategies
+        for text in [cleaned, re.search(r"\{.*\}", cleaned, re.DOTALL)]:
+            if not text:
+                continue
+            try:
+                return json.loads(text if isinstance(text, str) else text.group())
+            except json.JSONDecodeError:
+                continue
+
+        raise ValueError(f"Cannot parse JSON response: {response[:100]}...")
+
+
+class QueryRewriter(BaseMemAgent):
+    """Specialized agent for rewriting queries based on conversation history"""
+
+    def __init__(self, llm: BaseLLM, name: str = "QueryRewriter"):
+        self.llm = llm
+        self.name = name
+
+    def run(self, query: str, history: list[str] | None = None) -> str:
+        """Rewrite query to be standalone and more searchable"""
+        history = history or []
+        history_context = self._format_history(history)
+
+        prompt = QUERY_REWRITE_PROMPT.format(history=history_context, query=query)
+        messages = [{"role": "user", "content": prompt}]
+        try:
+            response = self.llm.generate(messages)
+            logger.info(f"[{self.name}] Rewritten query: {response.strip()}")
+            return response.strip()
+        except Exception as e:
+            logger.error(f"[{self.name}] Query rewrite failed: {e}")
+            return query
+
+    def _format_history(self, history: list[str]) -> str:
+        """Format conversation history for prompt context"""
+        if not history:
+            return "No previous conversation"
+        return "\n".join(f"- {msg}" for msg in history[-5:])
+
+
+class ReflectionAgent:
+    """Specialized agent for analyzing information sufficiency"""
+
+    def __init__(self, llm: BaseLLM, name: str = "Reflector"):
+        self.llm = llm
+        self.name = name
+
+    def run(self, query: str, context: list[str]) -> dict[str, Any]:
+        """Analyze whether retrieved context is sufficient to answer the query"""
+        context_summary = self._format_context(context)
+        prompt = REFLECTION_PROMPT.format(query=query, context=context_summary)
+
+        try:
+            response = self.llm.generate([{"role": "user", "content": prompt}])
+            logger.info(f"[{self.name}] Reflection response: {response}")
+
+            result = JSONResponseParser.parse(response.strip())
+            logger.info(f"[{self.name}] Reflection result: {result}")
+            return result
+
+        except Exception as e:
+            logger.error(f"[{self.name}] Reflection analysis failed: {e}")
+            return self._fallback_response()
+
+    def _format_context(self, context: list[str]) -> str:
+        """Format context strings for analysis with length limits"""
+        return "\n".join(
+            f"- {ctx[:200]}..." if len(ctx) > 200 else f"- {ctx}" for ctx in context[:10]
+        )
+
+    def _fallback_response(self) -> dict[str, Any]:
+        """Return safe fallback when reflection fails"""
+        return {
+            "status": "sufficient",
+            "reasoning": "Unable to analyze, proceeding with available information",
+            "missing_entities": [],
+        }
+
+
+class DeepSearchMemAgent(BaseMemAgent):
+    """
+    Main orchestrator agent implementing the deep search pipeline.
+
+    This agent coordinates multiple sub-agents to perform iterative query refinement,
+    memory retrieval, and information synthesis as shown in the architecture diagram.
+    """
+
+    def __init__(
+        self,
+        llm: BaseLLM,
+        memory_retriever: TreeTextMemory | None = None,
+        config: DeepSearchAgentConfig | None = None,
+    ):
+        """
+        Initialize DeepSearchMemAgent.
+
+        Args:
+            llm: Language model for query rewriting and response generation
+            memory_retriever: Memory retrieval interface (e.g., naive_mem_cube.text_mem)
+            config: Configuration for deep search behavior
+        """
+        self.config = config or DeepSearchAgentConfig(agent_name="DeepSearchMemAgent")
+        self.max_iterations = self.config.max_iterations
+        self.timeout = self.config.timeout
+        self.llm: BaseLLM = llm
+        self.query_rewriter: QueryRewriter = QueryRewriter(llm, "QueryRewriter")
+        self.reflector: ReflectionAgent = ReflectionAgent(llm, "Reflector")
+        self.memory_retriever = memory_retriever
+
+    def run(self, query: str, **kwargs) -> str | list[TextualMemoryItem]:
+        """
+        Main execution method implementing the deep search pipeline.
+
+        Args:
+            query: User query string
+            **kwargs: Additional arguments (history, user_id, etc.)
+        Returns:
+            Comprehensive response string
+        """
+        if not self.llm:
+            raise RuntimeError("LLM not initialized.")
+
+        history = kwargs.get("history", [])
+        user_id = kwargs.get("user_id")
+        generated_answer = kwargs.get("generated_answer")
+
+        # Step 1: Query Rewriting
+        current_query = self.query_rewriter.run(query, history)
+
+        accumulated_context = []
+        accumulated_memories = []
+        search_keywords = []  # Can be extended with keyword extraction
+
+        # Step 2: Iterative Search and Reflection Loop
+        for iteration in range(self.max_iterations):
+            logger.info(f"Starting iteration {iteration + 1}/{self.max_iterations}")
+
+            search_results = self._perform_memory_search(
+                current_query, keywords=search_keywords, user_id=user_id, history=history
+            )
+
+            if search_results:
+                context_batch = [self._extract_context_from_memory(mem) for mem in search_results]
+                accumulated_context.extend(context_batch)
+                reflection_result = self.reflector.run(current_query, context_batch)
+                status = reflection_result.get("status", "sufficient")
+                reasoning = reflection_result.get("reasoning", "")
+
+                logger.info(f"Reflection status: {status} - {reasoning}")
+
+                if status == "sufficient":
+                    logger.info("Sufficient information collected")
+                    accumulated_memories.extend(search_results)
+                    break
+                elif status == "needs_raw":
+                    logger.info("Need original sources, retrieving raw content")
+                    accumulated_memories.extend(self._set_source_from_memory(search_results))
+                    break
+                elif status == "missing_info":
+                    accumulated_memories.extend(search_results)
+                    missing_entities = reflection_result.get("missing_entities", [])
+                    logger.info(f"Missing information: {missing_entities}")
+                    current_query = reflection_result.get("new_search_query")
+                    if not current_query:
+                        refined_query = self._refine_query_for_missing_info(
+                            current_query, missing_entities
+                        )
+                        current_query = refined_query
+                        logger.info(f"Refined query: {current_query}")
+            else:
+                logger.warning(f"No search results for iteration {iteration + 1}")
+                if iteration == 0:
+                    current_query = query
+                else:
+                    break
+
+        if not generated_answer:
+            return self._remove_duplicate_memories(accumulated_memories)
+        else:
+            return self._generate_final_answer(
+                query, accumulated_memories, accumulated_context, history
+            )
+
+    def _remove_duplicate_memories(
+        self, memories: list[TextualMemoryItem]
+    ) -> list[TextualMemoryItem]:
+        """
+        Remove duplicate memories based on memory content.
+
+        Args:
+            memories: List of TextualMemoryItem objects to deduplicate
+
+        Returns:
+            List of unique TextualMemoryItem objects (first occurrence kept)
+        """
+        seen = set()
+        return [
+            memory
+            for memory in memories
+            if (content := getattr(memory, "memory", "").strip())
+            and content not in seen
+            and not seen.add(content)
+        ]
+
+    def _generate_final_answer(
+        self,
+        original_query: str,
+        search_results: list[TextualMemoryItem],
+        context: list[str],
+        history: list[str] | None = None,
+        sources: list[str] | None = None,
+        missing_info: str | None = None,
+    ) -> str:
+        """
+        Generate the final answer.
+        """
+        context_str = "\n".join([f"- {ctx}" for ctx in context[:20]])
+        prompt = FINAL_GENERATION_PROMPT.format(
+            query=original_query,
+            sources=sources,
+            context=context_str if context_str else "No specific context retrieved",
+            missing_info=missing_info if missing_info else "None identified",
+        )
+        messages: MessageList = [{"role": "user", "content": prompt}]
+        response = self.llm.generate(messages)
+        return response.strip()
+
+    def _perform_memory_search(
+        self,
+        query: str,
+        keywords: list[str] | None = None,
+        user_id: str | None = None,
+        history: list[str] | None = None,
+        top_k: int = 10,
+    ) -> list[TextualMemoryItem]:
+        """
+        Perform memory search using the configured retriever.
+
+        Args:
+            query: Search query
+            keywords: Additional keywords for search
+            user_id: User identifier
+            top_k: Number of results to retrieve
+
+        Returns:
+            List of retrieved memory items
+        """
+        if not self.memory_retriever:
+            logger.warning("Memory retriever not configured, returning empty results")
+            return []
+
+        try:
+            # Use the memory retriever interface
+            # This is a placeholder - actual implementation depends on the retriever interface
+            search_query = query
+            if keywords and len(keywords) > 1:
+                search_query = f"{query} {' '.join(keywords[:3])}"  # Combine with top keywords
+
+            # Assuming the retriever has a search method similar to TreeTextMemory
+            results = self.memory_retriever.search(
+                query=search_query,
+                top_k=top_k,
+                mode="fast",
+                user_name=user_id,
+                info={"history": history},
+            )
+
+            return results if isinstance(results, list) else []
+
+        except Exception as e:
+            logger.error(f"Error performing memory search: {e}")
+            return []
+
+    def _extract_context_from_memory(self, memory_item: TextualMemoryItem) -> str:
+        """Extract readable context from a memory item."""
+        if hasattr(memory_item, "memory"):
+            return str(memory_item.memory)
+        elif hasattr(memory_item, "content"):
+            return str(memory_item.content)
+        else:
+            return str(memory_item)
+
+    def _refine_query_for_missing_info(self, query: str, missing_entities: list[str]) -> str:
+        """Refine the query to search for missing information."""
+        if not missing_entities:
+            return query
+
+        # Simple refinement strategy - append missing entities
+        entities_str = " ".join(missing_entities[:3])  # Limit to top 3 entities
+        refined_query = f"{query} {entities_str}"
+
+        return refined_query
+
+    def _set_source_from_memory(
+        self, memory_items: list[TextualMemoryItem]
+    ) -> list[TextualMemoryItem]:
+        """set source from memory item"""
+        for memory_item in memory_items:
+            if not hasattr(memory_item.metadata, "sources"):
+                continue
+            chat_sources = [
+                f"{source.chat_time} {source.role}: {source.content}"
+                for source in memory_item.metadata.sources
+                if hasattr(source, "type") and source.type == "chat"
+            ]
+            if chat_sources:
+                memory_item.memory = "\n".join(chat_sources) + "\n"
+        return memory_items
+
+    def _generate_final_answer(
+        self,
+        original_query: str,
+        search_results: list[TextualMemoryItem],
+        context: list[str],
+        missing_info: str = "",
+    ) -> str:
+        """
+        Generate the final comprehensive answer.
+
+        Args:
+            original_query: Original user query
+            search_results: All retrieved memory items
+            context: Extracted context strings
+            missing_info: Information about missing data
+
+        Returns:
+            Final answer string
+        """
+        # Prepare context for the prompt
+        context_str = "\n".join([f"- {ctx}" for ctx in context[:20]])  # Limit context
+        sources = (
+            f"Retrieved {len(search_results)} memory items"
+            if search_results
+            else "No specific sources"
+        )
+
+        prompt = FINAL_GENERATION_PROMPT.format(
+            query=original_query,
+            sources=sources,
+            context=context_str if context_str else "No specific context retrieved",
+            missing_info=missing_info if missing_info else "None identified",
+        )
+        messages: MessageList = [{"role": "user", "content": prompt}]
+
+        try:
+            response = self.llm.generate(messages)
+            return response.strip()
+        except Exception as e:
+            logger.error(f"Error generating final answer: {e}")
+            return f"I apologize, but I encountered an error while processing your query: {original_query}. Please try again."

memos/mem_agent/factory.py

@@ -0,0 +1,36 @@
+from typing import Any, ClassVar
+
+from memos.configs.mem_agent import MemAgentConfigFactory
+from memos.mem_agent.base import BaseMemAgent
+from memos.mem_agent.deepsearch_agent import DeepSearchMemAgent
+
+
+class MemAgentFactory:
+    """Factory class for creating MemAgent instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "deep_search": DeepSearchMemAgent,
+    }
+
+    @classmethod
+    def from_config(
+        cls, config_factory: MemAgentConfigFactory, llm: Any, memory_retriever: Any | None = None
+    ) -> BaseMemAgent:
+        """
+        Create a MemAgent instance from configuration.
+
+        Args:
+            config_factory: Configuration factory for the agent
+            llm: Language model instance
+            memory_retriever: Memory retrieval interface (e.g., naive_mem_cube.text_mem)
+
+        Returns:
+            Initialized MemAgent instance
+        """
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        mem_agent_class = cls.backend_to_class[backend]
+        return mem_agent_class(
+            llm=llm, memory_retriever=memory_retriever, config=config_factory.config
+        )

memos/mem_chat/base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+
+from memos.configs.mem_chat import BaseMemChatConfig
+from memos.mem_cube.base import BaseMemCube
+
+
+class BaseMemChat(ABC):
+    """Base class for all MemChat."""
+
+    @abstractmethod
+    def __init__(self, config: BaseMemChatConfig):
+        """Initialize the MemChat with the given configuration."""
+
+    @property
+    @abstractmethod
+    def mem_cube(self) -> BaseMemCube:
+        """The memory cube associated with this MemChat."""
+
+    @mem_cube.setter
+    @abstractmethod
+    def mem_cube(self, value: BaseMemCube) -> None:
+        """The memory cube associated with this MemChat."""
+
+    @abstractmethod
+    def run(self) -> None:
+        """Run the MemChat.
+
+        This `run` method can represent the core logic of a MemChat.
+        It could be an iterative chat process.
+        """

memos/mem_chat/factory.py

@@ -0,0 +1,21 @@
+from typing import Any, ClassVar
+
+from memos.configs.mem_chat import MemChatConfigFactory
+from memos.mem_chat.base import BaseMemChat
+from memos.mem_chat.simple import SimpleMemChat
+
+
+class MemChatFactory(BaseMemChat):
+    """Factory class for creating MemChat instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "simple": SimpleMemChat,
+    }
+
+    @classmethod
+    def from_config(cls, config_factory: MemChatConfigFactory) -> BaseMemChat:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        mem_chat_class = cls.backend_to_class[backend]
+        return mem_chat_class(config_factory.config)

memos/mem_chat/simple.py

@@ -0,0 +1,200 @@
+import os
+
+from typing import Literal
+
+from memos.configs.mem_chat import SimpleMemChatConfig
+from memos.llms.factory import LLMFactory
+from memos.log import get_logger
+from memos.mem_chat.base import BaseMemChat
+from memos.mem_cube.base import BaseMemCube
+from memos.memories.activation.kv import move_dynamic_cache_htod
+from memos.memories.textual.item import TextualMemoryItem
+from memos.types import ChatHistory, MessageList
+
+
+logger = get_logger(__name__)
+
+
+class SimpleMemChat(BaseMemChat):
+    """Simple MemChat class."""
+
+    def __init__(self, config: SimpleMemChatConfig):
+        """Initialize the MemChat with the given configuration."""
+        self.config = config
+        self.chat_llm = LLMFactory.from_config(config.chat_llm)
+        self._mem_cube = None
+
+    @property
+    def mem_cube(self) -> BaseMemCube:
+        """The memory cube associated with this MemChat."""
+        return self._mem_cube
+
+    @mem_cube.setter
+    def mem_cube(self, value: BaseMemCube) -> None:
+        """The memory cube associated with this MemChat."""
+        self._mem_cube = value
+
+    def run(self) -> None:
+        """Run the MemChat."""
+
+        # Start MemChat
+
+        print(
+            "\n📢 [System] " + "Simple MemChat is running.\n"
+            "Commands: 'bye' to quit, 'clear' to clear chat history, 'mem' to show all memories, 'export' to export chat history\n",
+        )
+
+        messages = []
+        while True:
+            # Get user input
+
+            user_input = input("👤 [You] ").strip()
+            print()
+
+            if user_input.lower() == "bye":
+                break
+            elif user_input.lower() == "clear":
+                messages = []
+                print("📢 [System] Chat history cleared.")
+                continue
+            elif user_input.lower() == "mem":
+                if self.config.enable_textual_memory:
+                    all_memories = self.mem_cube.text_mem.get_all()
+                    print(f"🧠 [Memory] \n{self._str_memories(all_memories)}\n")
+                else:
+                    print("📢 [System] Textual memory is not enabled.\n")
+                continue
+            elif user_input.lower() == "export":
+                if messages:
+                    filepath = self._export_chat_history(messages)
+                    print(f"📢 [System] Chat history exported to: {filepath}\n")
+                else:
+                    print("📢 [System] No chat history to export.\n")
+                continue
+            elif user_input == "":
+                continue
+
+            # Get memories
+
+            if self.config.enable_textual_memory:
+                memories = self.mem_cube.text_mem.search(user_input, top_k=self.config.top_k)
+                print(
+                    f"🧠 [Memory] Searched memories:\n{self._str_memories(memories, mode='concise')}\n"
+                )
+                system_prompt = self._build_system_prompt(memories)
+            else:
+                system_prompt = self._build_system_prompt()
+            current_messages = [
+                {"role": "system", "content": system_prompt},
+                *messages,
+                {"role": "user", "content": user_input},
+            ]
+
+            if self.config.enable_activation_memory:
+                past_key_values = None
+                loaded_kv_cache_item = next(
+                    iter(self.mem_cube.act_mem.kv_cache_memories.values()), None
+                )
+                if loaded_kv_cache_item is not None:
+                    # If has loaded kv cache, we move it to device before inferring.
+                    # Currently, we move only single kv cache item
+                    past_key_values = loaded_kv_cache_item
+                    past_key_values.kv_cache = move_dynamic_cache_htod(
+                        past_key_values.kv_cache, self.chat_llm.model.device
+                    )
+
+                # Generate response
+                response = self.chat_llm.generate(
+                    current_messages,
+                    past_key_values=past_key_values.kv_cache if past_key_values else None,
+                )
+            else:
+                # Generate response without activation memory
+                response = self.chat_llm.generate(current_messages)
+
+            print(f"🤖 [Assistant] {response}\n")
+            messages.append({"role": "user", "content": user_input})
+            messages.append({"role": "assistant", "content": response})
+            messages = messages[
+                -self.config.max_turns_window :
+            ]  # Keep only recent messages to avoid context overflow
+
+            # Extract memories
+
+            if self.config.enable_textual_memory:
+                new_memories = self.mem_cube.text_mem.extract(messages[-2:])
+                for memory in new_memories:
+                    memory.metadata.user_id = self.config.user_id
+                    memory.metadata.session_id = self.config.session_id
+                    memory.metadata.status = "activated"
+                self.mem_cube.text_mem.add(new_memories)
+                print(
+                    f"🧠 [Memory] Stored {len(new_memories)} new memory(ies):\n"
+                    f"{self._str_memories(new_memories, 'concise')}\n"
+                )
+
+        # Stop MemChat
+
+        print("📢 [System] MemChat has stopped.")
+
+    def _build_system_prompt(self, memories: list | None = None) -> str:
+        """Build system prompt with optional memories context."""
+        base_prompt = (
+            "You are a knowledgeable and helpful AI assistant. "
+            "You have access to conversation memories that help you provide more personalized responses. "
+            "Use the memories to understand the user's context, preferences, and past interactions. "
+            "If memories are provided, reference them naturally when relevant, but don't explicitly mention having memories."
+        )
+
+        if memories:
+            memory_context = "\n\n## Memories:\n"
+            for i, memory in enumerate(memories, 1):
+                memory_context += f"{i}. ({memory.metadata.memory_time}) {memory.memory}\n"
+            return base_prompt + memory_context
+
+        return base_prompt
+
+    def _str_memories(
+        self, memories: list[TextualMemoryItem], mode: Literal["concise", "full"] = "full"
+    ) -> str:
+        """Format memories for display."""
+        if not memories:
+            return "No memories."
+        if mode == "concise":
+            return "\n".join(f"{i + 1}. {memory.memory}" for i, memory in enumerate(memories))
+        elif mode == "full":
+            return "\n".join(f"{i + 1}. {memory}" for i, memory in enumerate(memories))
+
+    def _export_chat_history(self, messages: MessageList, output_dir: str = "chat_exports") -> str:
+        """Export chat history to JSON file.
+
+        Args:
+            messages: List of chat messages
+            output_dir: Directory to save the export file
+
+        Returns:
+            Path to the exported JSON file
+        """
+        # Create output directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Generate filename with user_id and timestamp
+        timestamp = self.config.created_at.strftime("%Y%m%d_%H%M%S")
+        filename = f"{self.config.user_id}_{timestamp}_chat_history.json"
+        filepath = os.path.join(output_dir, filename)
+
+        # Prepare export data
+        export_data = ChatHistory(
+            user_id=self.config.user_id,
+            session_id=self.config.session_id,
+            created_at=self.config.created_at,
+            total_messages=len(messages),
+            chat_history=messages,
+        )
+
+        # Write to JSON file
+        with open(filepath, "w", encoding="utf-8") as f:
+            f.write(export_data.model_dump_json(indent=4, exclude_none=True, warnings="none"))
+
+        logger.info(f"Chat history exported to {filepath}")
+        return filepath