MemoryOS 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

memos/mem_cube/utils.py

@@ -1,6 +1,13 @@
+import copy
+import logging
 import subprocess
 import tempfile
 
+from memos.configs.mem_cube import GeneralMemCubeConfig
+
+
+logger = logging.getLogger(__name__)
+
 
 def download_repo(repo: str, base_url: str, dir: str | None = None) -> str:
     """Download a repository from a remote source.
@@ -22,3 +29,98 @@ def download_repo(repo: str, base_url: str, dir: str | None = None) -> str:
     subprocess.run(["git", "clone", repo_url, dir], check=True)
 
     return dir
+
+
+def merge_config_with_default(
+    existing_config: GeneralMemCubeConfig, default_config: GeneralMemCubeConfig
+) -> GeneralMemCubeConfig:
+    """
+    Merge existing cube config with default config, preserving critical fields.
+
+    This method updates general configuration fields (like API keys, model parameters)
+    while preserving critical user-specific fields (like user_id, cube_id, graph_db settings).
+
+    Args:
+        existing_config (GeneralMemCubeConfig): The existing cube configuration loaded from file
+        default_config (GeneralMemCubeConfig): The default configuration to merge from
+
+    Returns:
+        GeneralMemCubeConfig: Merged configuration
+    """
+
+    # Convert configs to dictionaries
+    existing_dict = existing_config.model_dump(mode="json")
+    default_dict = default_config.model_dump(mode="json")
+
+    logger.info(
+        f"Starting config merge for user {existing_config.user_id}, cube {existing_config.cube_id}"
+    )
+
+    # Define fields that should be preserved from existing config
+    preserve_fields = {"user_id", "cube_id", "config_filename", "model_schema"}
+
+    # Preserve graph_db from existing config if it exists, but merge some fields
+    preserved_graph_db = None
+    if "text_mem" in existing_dict and "text_mem" in default_dict:
+        existing_text_config = existing_dict["text_mem"].get("config", {})
+        default_text_config = default_dict["text_mem"].get("config", {})
+
+        if "graph_db" in existing_text_config and "graph_db" in default_text_config:
+            existing_graph_config = existing_text_config["graph_db"]["config"]
+            default_graph_config = default_text_config["graph_db"]["config"]
+
+            # Define graph_db fields to preserve (user-specific)
+            preserve_graph_fields = {
+                "uri",
+                "user",
+                "password",
+                "db_name",
+                "auto_create",
+                "user_name",
+                "use_multi_db",
+            }
+
+            # Create merged graph_db config
+            merged_graph_config = copy.deepcopy(existing_graph_config)
+            for key, value in default_graph_config.items():
+                if key not in preserve_graph_fields:
+                    merged_graph_config[key] = value
+                    logger.debug(
+                        f"Updated graph_db field '{key}': {existing_graph_config.get(key)} -> {value}"
+                    )
+            if not default_graph_config.get("use_multi_db", True):
+                # set original use_multi_db to False if default_graph_config.use_multi_db is False
+                if merged_graph_config.get("use_multi_db", True):
+                    merged_graph_config["use_multi_db"] = False
+                    merged_graph_config["user_name"] = merged_graph_config.get("db_name")
+                    merged_graph_config["db_name"] = default_graph_config.get("db_name")
+                else:
+                    logger.info("use_multi_db is already False, no need to change")
+
+            preserved_graph_db = {
+                "backend": existing_text_config["graph_db"]["backend"],
+                "config": merged_graph_config,
+            }
+
+    # Use default config as base
+    merged_dict = copy.deepcopy(default_dict)
+
+    # Restore preserved fields from existing config
+    for field in preserve_fields:
+        if field in existing_dict:
+            merged_dict[field] = existing_dict[field]
+            logger.debug(f"Preserved field '{field}': {existing_dict[field]}")
+
+    # Restore graph_db if it was preserved
+    if preserved_graph_db and "text_mem" in merged_dict:
+        merged_dict["text_mem"]["config"]["graph_db"] = preserved_graph_db
+        logger.debug(f"Preserved graph_db with merged config: {preserved_graph_db}")
+
+    # Create new config from merged dictionary
+    merged_config = GeneralMemCubeConfig.model_validate(merged_dict)
+
+    logger.info(
+        f"Successfully merged cube config for user {merged_config.user_id}, cube {merged_config.cube_id}"
+    )
+
+    return merged_config

memos/mem_os/core.py

@@ -1,4 +1,6 @@
+import json
 import os
+import uuid
 
 from datetime import datetime
 from pathlib import Path
@@ -11,7 +13,11 @@ from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_reader.factory import MemReaderFactory
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
-from memos.mem_scheduler.modules.schemas import ANSWER_LABEL, QUERY_LABEL, ScheduleMessageItem
+from memos.mem_scheduler.modules.schemas import (
+    ADD_LABEL,
+    ANSWER_LABEL,
+    ScheduleMessageItem,
+)
 from memos.mem_scheduler.scheduler_factory import SchedulerFactory
 from memos.mem_user.user_manager import UserManager, UserRole
 from memos.memories.activation.item import ActivationMemoryItem
@@ -30,7 +36,7 @@ class MOSCore:
     MOSCore acts as an operating system layer for handling and orchestrating MemCube instances.
     """
 
-    def __init__(self, config: MOSConfig):
+    def __init__(self, config: MOSConfig, user_manager: UserManager | None = None):
         self.config = config
         self.user_id = config.user_id
         self.session_id = config.session_id
@@ -39,7 +45,12 @@ class MOSCore:
         self.mem_reader = MemReaderFactory.from_config(config.mem_reader)
         self.chat_history_manager: dict[str, ChatHistory] = {}
         self._register_chat_history()
-        self.user_manager = UserManager(user_id=self.user_id if self.user_id else "root")
+
+        # Use provided user_manager or create a new one
+        if user_manager is not None:
+            self.user_manager = user_manager
+        else:
+            self.user_manager = UserManager(user_id=self.user_id if self.user_id else "root")
 
         # Validate user exists
         if not self.user_manager.validate_user(self.user_id):
@@ -50,7 +61,7 @@ class MOSCore:
         # Lazy initialization marker
         self._mem_scheduler_lock = Lock()
         self.enable_mem_scheduler = self.config.get("enable_mem_scheduler", False)
-        self._mem_scheduler = None
+        self._mem_scheduler: GeneralScheduler = None
         logger.info(f"MOS initialized for user: {self.user_id}")
 
     @property
@@ -58,6 +69,7 @@ class MOSCore:
         """Lazy-loaded property for memory scheduler."""
         if self.enable_mem_scheduler and self._mem_scheduler is None:
             self._initialize_mem_scheduler()
+        self._mem_scheduler.mem_cubes = self.mem_cubes
         return self._mem_scheduler
 
     @mem_scheduler.setter
@@ -74,6 +86,7 @@ class MOSCore:
                 raise TypeError(f"Expected GeneralScheduler or None, got {type(value)}")
 
             self._mem_scheduler = value
+            self._mem_scheduler.mem_cubes = self.mem_cubes
 
             if value:
                 logger.info("Memory scheduler manually set")
@@ -92,7 +105,18 @@ class MOSCore:
             logger.info("Initializing memory scheduler...")
             scheduler_config = self.config.mem_scheduler
             self._mem_scheduler = SchedulerFactory.from_config(scheduler_config)
-            self._mem_scheduler.initialize_modules(chat_llm=self.chat_llm)
+            # Validate required components
+            if not hasattr(self.mem_reader, "llm"):
+                raise AttributeError(
+                    f"Memory reader of type {type(self.mem_reader).__name__} "
+                    "missing required 'llm' attribute"
+                )
+                self._mem_scheduler.initialize_modules(chat_llm=self.chat_llm)
+            else:
+                # Configure scheduler modules
+                self._mem_scheduler.initialize_modules(
+                    chat_llm=self.chat_llm, process_llm=self.mem_reader.llm
+                )
             self._mem_scheduler.start()
 
     def mem_scheduler_on(self) -> bool:
@@ -123,6 +147,17 @@ class MOSCore:
             logger.error(f"Failed to stop scheduler: {e!s}")
             return False
 
+    def mem_reorganizer_on(self) -> bool:
+        pass
+
+    def mem_reorganizer_off(self) -> bool:
+        """temporally implement"""
+        for mem_cube in self.mem_cubes.values():
+            logger.info(f"try to close reorganizer for {mem_cube.text_mem.config.cube_id}")
+            if mem_cube.text_mem and mem_cube.text_mem.is_reorganize:
+                logger.info(f"close reorganizer for {mem_cube.text_mem.config.cube_id}")
+                mem_cube.text_mem.memory_manager.close()
+
     def _register_chat_history(self, user_id: str | None = None) -> None:
         """Initialize chat history with user ID."""
         if user_id is None:
@@ -186,12 +221,16 @@ class MOSCore:
                 documents.append(str(file_path))
         return documents
 
-    def chat(self, query: str, user_id: str | None = None) -> str:
+    def chat(self, query: str, user_id: str | None = None, base_prompt: str | None = None) -> str:
         """
         Chat with the MOS.
 
         Args:
             query (str): The user's query.
+            user_id (str, optional): The user ID for the chat session. Defaults to the user ID from the config.
+            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.
@@ -218,7 +257,7 @@ class MOSCore:
                         user_id=target_user_id,
                         mem_cube_id=mem_cube_id,
                         mem_cube=mem_cube,
-                        label=QUERY_LABEL,
+                        label=ADD_LABEL,
                         content=query,
                         timestamp=datetime.now(),
                     )
@@ -227,9 +266,9 @@ class MOSCore:
                 memories = mem_cube.text_mem.search(query, top_k=self.config.top_k)
                 memories_all.extend(memories)
             logger.info(f"🧠 [Memory] Searched memories:\n{self._str_memories(memories_all)}\n")
-            system_prompt = self._build_system_prompt(memories_all)
+            system_prompt = self._build_system_prompt(memories_all, base_prompt=base_prompt)
         else:
-            system_prompt = self._build_system_prompt()
+            system_prompt = self._build_system_prompt(base_prompt=base_prompt)
         current_messages = [
             {"role": "system", "content": system_prompt},
             *chat_history.chat_history,
@@ -261,8 +300,8 @@ class MOSCore:
         self.chat_history_manager[user_id] = chat_history
 
         # submit message to scheduler
-        if len(accessible_cubes) == 1:
-            mem_cube_id = accessible_cubes[0].cube_id
+        for accessible_mem_cube in accessible_cubes:
+            mem_cube_id = accessible_mem_cube.cube_id
             mem_cube = self.mem_cubes[mem_cube_id]
             if self.enable_mem_scheduler and self.mem_scheduler is not None:
                 message_item = ScheduleMessageItem(
@@ -277,20 +316,39 @@ class MOSCore:
 
         return response
 
-    def _build_system_prompt(self, memories: list | None = None) -> str:
+    def _build_system_prompt(
+        self,
+        memories: list[TextualMemoryItem] | list[str] | None = None,
+        base_prompt: str | 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 base_prompt is None:
+            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."
+            )
 
+        memory_context = ""
         if memories:
-            memory_context = "\n\n## Memories:\n"
+            memory_list = []
             for i, memory in enumerate(memories, 1):
-                memory_context += f"{i}. {memory.memory}\n"
-            return base_prompt + memory_context
+                if isinstance(memory, TextualMemoryItem):
+                    text_memory = memory.memory
+                else:
+                    if not isinstance(memory, str):
+                        logger.error("Unexpected memory type.")
+                    text_memory = memory
+                memory_list.append(f"{i}. {text_memory}")
+            memory_context = "\n".join(memory_list)
+
+        if "{memories}" in base_prompt:
+            return base_prompt.format(memories=memory_context)
+        elif memories:
+            # For backward compatibility, append memories if no placeholder is found
+            memory_context_with_header = "\n\n## Memories:\n" + memory_context
+            return base_prompt + memory_context_with_header
         return base_prompt
 
     def _str_memories(
@@ -364,7 +422,10 @@ class MOSCore:
         return self.user_manager.create_cube(cube_name, owner_id, cube_path, cube_id)
 
     def register_mem_cube(
-        self, mem_cube_name_or_path: str, mem_cube_id: str | None = None, user_id: str | None = None
+        self,
+        mem_cube_name_or_path: str | GeneralMemCube,
+        mem_cube_id: str | None = None,
+        user_id: str | None = None,
     ) -> None:
         """
         Register a MemCube with the MOS.
@@ -377,12 +438,18 @@ class MOSCore:
         self._validate_user_exists(target_user_id)
 
         if mem_cube_id is None:
-            mem_cube_id = mem_cube_name_or_path
+            if isinstance(mem_cube_name_or_path, GeneralMemCube):
+                mem_cube_id = f"cube_{target_user_id}"
+            else:
+                mem_cube_id = mem_cube_name_or_path
 
         if mem_cube_id in self.mem_cubes:
             logger.info(f"MemCube with ID {mem_cube_id} already in MOS, skip install.")
         else:
-            if os.path.exists(mem_cube_name_or_path):
+            if isinstance(mem_cube_name_or_path, GeneralMemCube):
+                self.mem_cubes[mem_cube_id] = mem_cube_name_or_path
+                logger.info(f"register new cube {mem_cube_id} for user {target_user_id}")
+            elif os.path.exists(mem_cube_name_or_path):
                 self.mem_cubes[mem_cube_id] = GeneralMemCube.init_from_dir(mem_cube_name_or_path)
             else:
                 logger.warning(
@@ -394,6 +461,14 @@ class MOSCore:
         # Check if cube already exists in database
         existing_cube = self.user_manager.get_cube(mem_cube_id)
 
+        # check the embedder is it consistent with MOSConfig
+        if self.config.mem_reader.config.embedder != (
+            cube_embedder := self.mem_cubes[mem_cube_id].text_mem.config.embedder
+        ):
+            logger.warning(
+                f"Cube Embedder is not consistent with MOSConfig for cube: {mem_cube_id}, will use Cube Embedder: {cube_embedder}"
+            )
+
         if existing_cube:
             # Cube exists, just add user to cube if not already associated
             if not self.user_manager.validate_user_cube_access(target_user_id, mem_cube_id):
@@ -407,10 +482,14 @@ class MOSCore:
         else:
             # Cube doesn't exist, create it
             self.create_cube_for_user(
-                cube_name=mem_cube_name_or_path,
+                cube_name=mem_cube_name_or_path
+                if not isinstance(mem_cube_name_or_path, GeneralMemCube)
+                else mem_cube_id,
                 owner_id=target_user_id,
                 cube_id=mem_cube_id,
-                cube_path=mem_cube_name_or_path,
+                cube_path=mem_cube_name_or_path
+                if not isinstance(mem_cube_name_or_path, GeneralMemCube)
+                else "init",
             )
             logger.info(f"register new cube {mem_cube_id} for user {target_user_id}")
 
@@ -427,7 +506,11 @@ class MOSCore:
             raise ValueError(f"MemCube with ID {mem_cube_id} does not exist.")
 
     def search(
-        self, query: str, user_id: str | None = None, install_cube_ids: list[str] | None = None
+        self,
+        query: str,
+        user_id: str | None = None,
+        install_cube_ids: list[str] | None = None,
+        top_k: int | None = None,
     ) -> MOSSearchResult:
         """
         Search for textual memories across all registered MemCubes.
@@ -464,18 +547,10 @@ class MOSCore:
                 and (mem_cube.text_mem is not None)
                 and self.config.enable_textual_memory
             ):
-                memories = mem_cube.text_mem.search(query, top_k=self.config.top_k)
-                result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories})
-                logger.info(
-                    f"🧠 [Memory] Searched memories from {mem_cube_id}:\n{self._str_memories(memories)}\n"
+                memories = mem_cube.text_mem.search(
+                    query, top_k=top_k if top_k else self.config.top_k
                 )
-            if (
-                (mem_cube_id in install_cube_ids)
-                and (mem_cube.act_mem is not None)
-                and self.config.enable_activation_memory
-            ):
-                memories = mem_cube.act_mem.extract(query)
-                result["act_mem"].append({"cube_id": mem_cube_id, "memories": [memories]})
+                result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories})
                 logger.info(
                     f"🧠 [Memory] Searched memories from {mem_cube_id}:\n{self._str_memories(memories)}\n"
                 )
@@ -538,10 +613,25 @@ class MOSCore:
                 memories = self.mem_reader.get_memory(
                     messages_list,
                     type="chat",
-                    info={"user_id": target_user_id, "session_id": self.session_id},
+                    info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
                 )
                 for mem in memories:
                     self.mem_cubes[mem_cube_id].text_mem.add(mem)
+
+                # submit messages for scheduler
+                mem_cube = self.mem_cubes[mem_cube_id]
+                if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                    text_messages = [message["content"] for message in messages]
+                    message_item = ScheduleMessageItem(
+                        user_id=target_user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=mem_cube,
+                        label=ADD_LABEL,
+                        content=json.dumps(text_messages),
+                        timestamp=datetime.now(),
+                    )
+                    self.mem_scheduler.submit_messages(messages=[message_item])
+
         if (
             (memory_content is not None)
             and self.config.enable_textual_memory
@@ -567,7 +657,7 @@ class MOSCore:
                 memories = self.mem_reader.get_memory(
                     messages_list,
                     type="chat",
-                    info={"user_id": target_user_id, "session_id": self.session_id},
+                    info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
                 )
                 for mem in memories:
                     self.mem_cubes[mem_cube_id].text_mem.add(mem)
@@ -580,7 +670,7 @@ class MOSCore:
             doc_memory = self.mem_reader.get_memory(
                 documents,
                 type="doc",
-                info={"user_id": target_user_id, "session_id": self.session_id},
+                info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
             )
             for mem in doc_memory:
                 self.mem_cubes[mem_cube_id].text_mem.add(mem)

memos/mem_os/main.py

@@ -1,5 +1,6 @@
 import concurrent.futures
 import json
+import os
 
 from typing import Any
 
@@ -7,6 +8,7 @@ from memos.configs.mem_os import MOSConfig
 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,
@@ -24,20 +26,94 @@ class MOS(MOSCore):
     This class maintains backward compatibility with the original MOS interface.
     """
 
-    def __init__(self, config: MOSConfig):
+    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)
 
-    def chat(self, query: str, user_id: str | None = None) -> str:
+        # 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.
@@ -46,12 +122,14 @@ class MOS(MOSCore):
 
         if not self.enable_cot:
             # Use the original chat method from core
-            return super().chat(query, user_id)
+            return super().chat(query, user_id, base_prompt=base_prompt)
 
         # Enhanced chat with CoT decomposition
-        return self._chat_with_cot_enhancement(query, user_id)
+        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) -> str:
+    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.
@@ -84,7 +162,7 @@ class MOS(MOSCore):
             # 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)
+                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")
@@ -93,7 +171,7 @@ class MOS(MOSCore):
             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)
+                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...")
@@ -115,6 +193,7 @@ class MOS(MOSCore):
                 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)
@@ -149,7 +228,7 @@ class MOS(MOSCore):
         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)
+            return super().chat(query, user_id, base_prompt=base_prompt)
 
     def _get_search_engine_for_cot_with_validation(
         self, user_cube_ids: list[str]
@@ -183,6 +262,7 @@ class MOS(MOSCore):
         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.
@@ -193,6 +273,8 @@ class MOS(MOSCore):
             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.
@@ -213,10 +295,10 @@ class MOS(MOSCore):
                     original_query, top_k=self.config.top_k, mode="fast"
                 )
             system_prompt = self._build_system_prompt(
-                search_memories
+                search_memories, base_prompt=base_prompt
             )  # Use the same system prompt builder
         else:
-            system_prompt = self._build_system_prompt()
+            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,
@@ -261,7 +343,7 @@ class MOS(MOSCore):
         except Exception as e:
             logger.error(f"🔍 [CoT] Error generating enhanced response: {e}")
             # Fallback to standard chat
-            return super().chat(original_query, user_id)
+            return super().chat(original_query, user_id, base_prompt=base_prompt)
 
     @classmethod
     def cot_decompose(