MemoryOS 2.0.3__py3-none-any.whl

memos/mem_os/main.py

@@ -0,0 +1,582 @@
+import concurrent.futures
+import json
+import os
+
+from typing import Any
+
+from memos.configs.mem_os import MOSConfig
+from memos.context.context import ContextThreadPoolExecutor
+from memos.llms.factory import LLMFactory
+from memos.log import get_logger
+from memos.mem_os.core import MOSCore
+from memos.mem_os.utils.default_config import get_default
+from memos.memories.textual.base import BaseTextMemory
+from memos.templates.mos_prompts import (
+    COT_DECOMPOSE_PROMPT,
+    PRO_MODE_WELCOME_MESSAGE,
+    SYNTHESIS_PROMPT,
+)
+
+
+logger = get_logger(__name__)
+
+
+class MOS(MOSCore):
+    """
+    The MOS (Memory Operating System) class inherits from MOSCore.
+    This class maintains backward compatibility with the original MOS interface.
+    """
+
+    def __init__(self, config: MOSConfig | None = None):
+        """
+        Initialize MOS with optional automatic configuration.
+
+        Args:
+            config (MOSConfig, optional): MOS configuration. If None, will use automatic configuration from environment variables.
+        """
+        if config is None:
+            # Auto-configure if no config provided
+            config, default_cube = self._auto_configure()
+            self._auto_registered_cube = default_cube
+        else:
+            self._auto_registered_cube = None
+
+        self.enable_cot = config.PRO_MODE
+        if config.PRO_MODE:
+            print(PRO_MODE_WELCOME_MESSAGE)
+            logger.info(PRO_MODE_WELCOME_MESSAGE)
+        super().__init__(config)
+
+        # Auto-register cube if one was created
+        if self._auto_registered_cube is not None:
+            self.register_mem_cube(self._auto_registered_cube)
+            logger.info(
+                f"Auto-registered default cube: {self._auto_registered_cube.config.cube_id}"
+            )
+
+    def _auto_configure(self, **kwargs) -> tuple[MOSConfig, Any]:
+        """
+        Automatically configure MOS with default settings.
+
+        Returns:
+            tuple[MOSConfig, Any]: MOS configuration and default MemCube
+        """
+        # Get configuration from environment variables
+        openai_api_key = os.getenv("OPENAI_API_KEY")
+        openai_api_base = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1")
+        text_mem_type = os.getenv("MOS_TEXT_MEM_TYPE", "general_text")
+
+        if not openai_api_key:
+            raise ValueError("OPENAI_API_KEY environment variable is required")
+
+        logger.info(f"Auto-configuring MOS with text_mem_type: {text_mem_type}")
+        return get_default(
+            openai_api_key=openai_api_key,
+            openai_api_base=openai_api_base,
+            text_mem_type=text_mem_type,
+        )
+
+    @classmethod
+    def simple(cls) -> "MOS":
+        """
+        Create a MOS instance with automatic configuration from environment variables.
+
+        This is the simplest way to get started with MemOS.
+
+        Environment variables needed:
+        - OPENAI_API_KEY: Your OpenAI API key
+        - OPENAI_API_BASE: OpenAI API base URL (optional, defaults to "https://api.openai.com/v1")
+        - MOS_TEXT_MEM_TYPE: Text memory type (optional, defaults to "general_text")
+
+        Returns:
+            MOS: Configured MOS instance with auto-registered default cube
+
+        Example:
+            ```python
+            # Set environment variables
+            export OPENAI_API_KEY="your-api-key"
+            export MOS_TEXT_MEM_TYPE="general_text"
+
+            # Then use
+            memory = MOS.simple()
+            memory.add_memory("Hello world!")
+            response = memory.chat("What did I just say?")
+            ```
+        """
+        return cls()
+
+    def chat(self, query: str, user_id: str | None = None, base_prompt: str | None = None) -> str:
+        """
+        Enhanced chat method with optional CoT (Chain of Thought) enhancement.
+
+        Args:
+            query (str): The user's query.
+            user_id (str, optional): User ID for context.
+            base_prompt (str, optional): A custom base prompt to use for the chat.
+                It can be a template string with a `{memories}` placeholder.
+                If not provided, a default prompt is used.
+
+        Returns:
+            str: The response from the MOS.
+        """
+        # Check if CoT enhancement is enabled (either explicitly or via PRO mode)
+
+        if not self.enable_cot:
+            # Use the original chat method from core
+            return super().chat(query, user_id, base_prompt=base_prompt)
+
+        # Enhanced chat with CoT decomposition
+        return self._chat_with_cot_enhancement(query, user_id, base_prompt=base_prompt)
+
+    def _chat_with_cot_enhancement(
+        self, query: str, user_id: str | None = None, base_prompt: str | None = None
+    ) -> str:
+        """
+        Chat with CoT enhancement for complex query decomposition.
+        This method includes all the same validation and processing logic as the core chat method.
+
+        Args:
+            query (str): The user's query.
+            user_id (str, optional): User ID for context.
+
+        Returns:
+            str: The enhanced response.
+        """
+        # Step 1: Perform all the same validation and setup as core chat method
+        target_user_id = user_id if user_id is not None else self.user_id
+        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
+        user_cube_ids = [cube.cube_id for cube in accessible_cubes]
+
+        # Register chat history if needed
+        if target_user_id not in self.chat_history_manager:
+            self._register_chat_history(target_user_id)
+
+        chat_history = self.chat_history_manager[target_user_id]
+
+        try:
+            # Step 2: Decompose the query using CoT
+            logger.info(f"🔍 [CoT] Decomposing query: {query}")
+            decomposition_result = self.cot_decompose(
+                query, self.config.chat_model, target_user_id, self.chat_llm
+            )
+
+            # Check if the query is complex and needs decomposition
+            if not decomposition_result.get("is_complex", False):
+                logger.info("🔍 [CoT] Query is not complex, using standard chat")
+                return super().chat(query, user_id, base_prompt=base_prompt)
+
+            sub_questions = decomposition_result.get("sub_questions", [])
+            logger.info(f"🔍 [CoT] Decomposed into {len(sub_questions)} sub-questions")
+
+            # Step 3: Get search engine for sub-questions (with proper validation)
+            search_engine = self._get_search_engine_for_cot_with_validation(user_cube_ids)
+            if not search_engine:
+                logger.warning("🔍 [CoT] No search engine available, using standard chat")
+                return super().chat(query, user_id, base_prompt=base_prompt)
+
+            # Step 4: Get answers for sub-questions
+            logger.info("🔍 [CoT] Getting answers for sub-questions...")
+            sub_questions, sub_answers = self.get_sub_answers(
+                sub_questions=sub_questions,
+                search_engine=search_engine,
+                llm_config=self.config.chat_model,
+                user_id=target_user_id,
+                top_k=getattr(self.config, "cot_top_k", 3),
+                llm=self.chat_llm,
+            )
+
+            # Step 5: Generate enhanced response using sub-answers
+            logger.info("🔍 [CoT] Generating enhanced response...")
+            enhanced_response = self._generate_enhanced_response_with_context(
+                original_query=query,
+                sub_questions=sub_questions,
+                sub_answers=sub_answers,
+                chat_history=chat_history,
+                user_id=target_user_id,
+                search_engine=search_engine,
+                base_prompt=base_prompt,
+            )
+
+            # Step 6: Update chat history (same as core method)
+            chat_history.chat_history.append({"role": "user", "content": query})
+            chat_history.chat_history.append({"role": "assistant", "content": enhanced_response})
+            self.chat_history_manager[target_user_id] = chat_history
+
+            # Step 7: Submit message to scheduler (same as core method)
+            if len(accessible_cubes) == 1:
+                mem_cube_id = accessible_cubes[0].cube_id
+                if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                    from datetime import datetime
+
+                    from memos.mem_scheduler.schemas import (
+                        ANSWER_LABEL,
+                        ScheduleMessageItem,
+                    )
+
+                    message_item = ScheduleMessageItem(
+                        user_id=target_user_id,
+                        mem_cube_id=mem_cube_id,
+                        label=ANSWER_LABEL,
+                        content=enhanced_response,
+                        timestamp=datetime.now().isoformat(),
+                    )
+                    self.mem_scheduler.submit_messages(messages=[message_item])
+
+            return enhanced_response
+
+        except Exception as e:
+            logger.error(f"🔍 [CoT] Error in CoT enhancement: {e}")
+            logger.info("🔍 [CoT] Falling back to standard chat")
+            return super().chat(query, user_id, base_prompt=base_prompt)
+
+    def _get_search_engine_for_cot_with_validation(
+        self, user_cube_ids: list[str]
+    ) -> BaseTextMemory | None:
+        """
+        Get the best available search engine for CoT operations with proper validation.
+
+        Args:
+            user_cube_ids (list[str]): List of cube IDs the user has access to.
+
+        Returns:
+            BaseTextMemory or None: The search engine to use for CoT.
+        """
+        if not self.mem_cubes:
+            return None
+
+        # Get the first available text memory from user's accessible cubes
+        for mem_cube_id, mem_cube in self.mem_cubes.items():
+            if mem_cube_id not in user_cube_ids:
+                continue
+            if mem_cube.text_mem:
+                return mem_cube.text_mem
+
+        return None
+
+    def _generate_enhanced_response_with_context(
+        self,
+        original_query: str,
+        sub_questions: list[str],
+        sub_answers: list[str],
+        chat_history: Any,
+        user_id: str | None = None,
+        search_engine: BaseTextMemory | None = None,
+        base_prompt: str | None = None,
+    ) -> str:
+        """
+        Generate an enhanced response using sub-questions and their answers, with chat context.
+
+        Args:
+            original_query (str): The original user query.
+            sub_questions (list[str]): List of sub-questions.
+            sub_answers (list[str]): List of answers to sub-questions.
+            chat_history: The user's chat history.
+            user_id (str, optional): User ID for context.
+            search_engine (BaseTextMemory, optional): Search engine for context retrieval.
+            base_prompt (str, optional): A custom base prompt for the chat.
+
+        Returns:
+            str: The enhanced response.
+        """
+        # Build the synthesis prompt
+        qa_text = ""
+        for i, (question, answer) in enumerate(zip(sub_questions, sub_answers, strict=False), 1):
+            qa_text += f"Q{i}: {question}\nA{i}: {answer}\n\n"
+
+        # Build messages with chat history context (similar to core method)
+        if (search_engine is not None) and self.config.enable_textual_memory:
+            if self.enable_cot:
+                search_memories = search_engine.search(
+                    original_query, top_k=self.config.top_k, mode="fine"
+                )
+            else:
+                search_memories = search_engine.search(
+                    original_query, top_k=self.config.top_k, mode="fast"
+                )
+            system_prompt = self._build_system_prompt(
+                search_memories, base_prompt=base_prompt
+            )  # Use the same system prompt builder
+        else:
+            system_prompt = self._build_system_prompt(base_prompt=base_prompt)
+        current_messages = [
+            {"role": "system", "content": system_prompt + SYNTHESIS_PROMPT.format(qa_text=qa_text)},
+            *chat_history.chat_history,
+            {
+                "role": "user",
+                "content": original_query,
+            },
+        ]
+
+        # Handle activation memory if enabled (same as core method)
+        past_key_values = None
+        if self.config.enable_activation_memory:
+            if self.config.chat_model.backend not in ["huggingface", "huggingface_singleton"]:
+                logger.error(
+                    "Activation memory only used for huggingface backend. Skipping activation memory."
+                )
+            else:
+                # Get accessible cubes for the user
+                target_user_id = user_id if user_id is not None else self.user_id
+                accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
+                user_cube_ids = [cube.cube_id for cube in accessible_cubes]
+
+                for mem_cube_id, mem_cube in self.mem_cubes.items():
+                    if mem_cube_id not in user_cube_ids:
+                        continue
+                    if mem_cube.act_mem:
+                        kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
+                        past_key_values = (
+                            kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
+                        )
+                        break
+
+        try:
+            # Generate the enhanced response using the chat LLM with same parameters as core
+            if past_key_values is not None:
+                enhanced_response = self.chat_llm.generate(
+                    current_messages, past_key_values=past_key_values
+                )
+            else:
+                enhanced_response = self.chat_llm.generate(current_messages)
+
+            logger.info("🔍 [CoT] Generated enhanced response")
+            return enhanced_response
+        except Exception as e:
+            logger.error(f"🔍 [CoT] Error generating enhanced response: {e}")
+            # Fallback to standard chat
+            return super().chat(original_query, user_id, base_prompt=base_prompt)
+
+    @classmethod
+    def cot_decompose(
+        cls, query: str, llm_config: Any, user_id: str | None = None, llm: LLMFactory | None = None
+    ) -> list[str] | dict[str, Any]:
+        """
+        Decompose a complex query into sub-questions using Chain of Thought reasoning.
+
+        Args:
+            query (str): The complex query to decompose
+            llm_config: LLM configuration for decomposition
+            user_id (str, optional): User ID for context
+
+        Returns:
+            Union[List[str], Dict[str, Any]]: List of decomposed sub-questions or dict with complexity analysis
+        """
+        # Create a temporary LLM instance for decomposition
+        if llm is None:
+            llm = LLMFactory.from_config(llm_config)
+
+        # System prompt for CoT decomposition with complexity analysis
+        system_prompt = COT_DECOMPOSE_PROMPT.format(query=query)
+
+        messages = [{"role": "system", "content": system_prompt}]
+
+        try:
+            response = llm.generate(messages)
+            # Try to parse JSON response
+            result = json.loads(response)
+            return result
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse JSON response from LLM: {e}")
+            logger.warning(f"Raw response: {response}")
+
+            # Try to extract JSON-like content from the response
+            try:
+                # Look for JSON-like content between curly braces
+                import re
+
+                json_match = re.search(r"\{.*\}", response, re.DOTALL)
+                if json_match:
+                    json_str = json_match.group(0)
+                    result = json.loads(json_str)
+                    return result
+            except Exception:
+                pass
+
+            # If all parsing attempts fail, return default
+            return {"is_complex": False, "sub_questions": []}
+        except Exception as e:
+            logger.error(f"Unexpected error in cot_decompose: {e}")
+            return {"is_complex": False, "sub_questions": []}
+
+    @classmethod
+    def get_sub_answers(
+        cls,
+        sub_questions: list[str] | dict[str, Any],
+        search_results: dict[str, Any] | None = None,
+        search_engine: BaseTextMemory | None = None,
+        llm_config: LLMFactory | None = None,
+        user_id: str | None = None,
+        top_k: int = 5,
+        llm: LLMFactory | None = None,
+    ) -> tuple[list[str], list[str]]:
+        """
+        Get answers for sub-questions using either search results or a search engine.
+
+        Args:
+            sub_questions (Union[List[str], Dict[str, Any]]): List of sub-questions from cot_decompose or dict with analysis
+            search_results (Dict[str, Any], optional): Search results containing relevant information
+            search_engine (BaseTextMemory, optional): Text memory engine for searching
+            llm_config (Any, optional): LLM configuration for processing (required if search_engine is provided)
+            user_id (str, optional): User ID for context
+            top_k (int): Number of top results to retrieve from search engine
+
+        Returns:
+            Tuple[List[str], List[str]]: (sub_questions, sub_answers)
+        """
+        # Extract sub-questions from decomposition result if needed
+        if isinstance(sub_questions, dict):
+            if not sub_questions.get("is_complex", False):
+                return [], []
+            sub_questions = sub_questions.get("sub_questions", [])
+
+        if not sub_questions:
+            return [], []
+
+        # Validate inputs
+        if search_results is None and search_engine is None:
+            raise ValueError("Either search_results or search_engine must be provided")
+        if llm is None:
+            llm = LLMFactory.from_config(llm_config)
+
+        # Step 1: Get search results if search_engine is provided
+        if search_engine is not None:
+            search_results = cls._search_with_engine(sub_questions, search_engine, top_k)
+
+        # Step 2: Generate answers for each sub-question using LLM in parallel
+        def generate_answer_for_question(question_index: int, sub_question: str) -> tuple[int, str]:
+            """Generate answer for a single sub-question."""
+            # Extract relevant information from search results
+            relevant_info = []
+            if search_results and search_results.get("text_mem"):
+                for cube_result in search_results["text_mem"]:
+                    for memory in cube_result.get("memories", []):
+                        relevant_info.append(memory.memory)
+
+            # Build system prompt with memories (similar to MOSCore._build_system_prompt)
+            base_prompt = (
+                "You are a knowledgeable and helpful AI assistant. "
+                "You have access to relevant information that helps you provide accurate answers. "
+                "Use the provided information to answer the question comprehensively. "
+                "If the information is not sufficient, acknowledge the limitations."
+            )
+
+            # Add memory context if available
+            if relevant_info:
+                memory_context = "\n\n## Relevant Information:\n"
+                for j, info in enumerate(relevant_info[:top_k], 1):  # Take top 3 most relevant
+                    memory_context += f"{j}. {info}\n"
+                system_prompt = base_prompt + memory_context
+            else:
+                system_prompt = (
+                    base_prompt
+                    + "\n\n## Relevant Information:\nNo specific information found in memory."
+                )
+
+            # Create messages for LLM
+            messages = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": sub_question},
+            ]
+
+            try:
+                # Generate answer using LLM
+                response = llm.generate(messages)
+                return question_index, response
+            except Exception as e:
+                logger.error(f"Failed to generate answer for sub-question '{sub_question}': {e}")
+                return question_index, f"Unable to generate answer for: {sub_question}"
+
+        # Generate answers in parallel while maintaining order
+        sub_answers = [None] * len(sub_questions)
+        with ContextThreadPoolExecutor(max_workers=min(len(sub_questions), 10)) as executor:
+            # Submit all answer generation tasks
+            future_to_index = {
+                executor.submit(generate_answer_for_question, i, question): i
+                for i, question in enumerate(sub_questions)
+            }
+
+            # Collect results as they complete, but store them in the correct position
+            for future in concurrent.futures.as_completed(future_to_index):
+                try:
+                    question_index, answer = future.result()
+                    sub_answers[question_index] = answer
+                except Exception as e:
+                    question_index = future_to_index[future]
+                    logger.error(
+                        f"Exception occurred while generating answer for question at index {question_index}: {e}"
+                    )
+                    sub_answers[question_index] = (
+                        f"Error generating answer for question {question_index + 1}"
+                    )
+
+        return sub_questions, sub_answers
+
+    @classmethod
+    def _search_with_engine(
+        cls, sub_questions: list[str], search_engine: BaseTextMemory, top_k: int
+    ) -> dict[str, Any]:
+        """
+        Search for sub-questions using the provided search engine in parallel.
+
+        Args:
+            sub_questions (List[str]): List of sub-questions to search for
+            search_engine (BaseTextMemory): Text memory engine for searching
+            top_k (int): Number of top results to retrieve
+
+        Returns:
+            Dict[str, Any]: Search results in the expected format
+        """
+
+        def search_single_question(question: str) -> list[Any]:
+            """Search for a single question using the search engine."""
+            try:
+                # Handle different search method signatures
+                if hasattr(search_engine, "search"):
+                    # Try different parameter combinations based on the engine type
+                    try:
+                        # For tree_text memory
+                        return search_engine.search(question, top_k, mode="fast")
+                    except TypeError:
+                        try:
+                            # For general_text memory
+                            return search_engine.search(question, top_k)
+                        except TypeError:
+                            # For naive_text memory
+                            return search_engine.search(question, top_k)
+                else:
+                    return []
+            except Exception as e:
+                logger.error(f"Search failed for question '{question}': {e}")
+                return []
+
+        # Search in parallel while maintaining order
+        all_memories = []
+        with ContextThreadPoolExecutor(max_workers=min(len(sub_questions), 10)) as executor:
+            # Submit all search tasks and keep track of their order
+            future_to_index = {
+                executor.submit(search_single_question, question): i
+                for i, question in enumerate(sub_questions)
+            }
+
+            # Initialize results list with None values to maintain order
+            results = [None] * len(sub_questions)
+
+            # Collect results as they complete, but store them in the correct position
+            for future in concurrent.futures.as_completed(future_to_index):
+                index = future_to_index[future]
+                try:
+                    memories = future.result()
+                    results[index] = memories
+                except Exception as e:
+                    logger.error(
+                        f"Exception occurred while searching for question at index {index}: {e}"
+                    )
+                    results[index] = []
+
+            # Combine all results in the correct order
+            for result in results:
+                if result is not None:
+                    all_memories.extend(result)
+
+        # Format results in the expected structure
+        return {"text_mem": [{"cube_id": "search_engine", "memories": all_memories}]}