MemoryOS 0.2.0__py3-none-any.whl → 0.2.2__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,105 @@ 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 = {
+                "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")
+            if "neo4j" not in default_text_config["graph_db"]["backend"]:
+                if "db_name" in merged_graph_config:
+                    merged_graph_config.pop("db_name")
+                    logger.info("neo4j is not supported, remove db_name")
+                else:
+                    logger.info("db_name is not in merged_graph_config, no need to remove")
+            else:
+                if "space" in merged_graph_config:
+                    merged_graph_config.pop("space")
+                    logger.info("neo4j is not supported, remove db_name")
+                else:
+                    logger.info("space is not in merged_graph_config, no need to remove")
+            preserved_graph_db = {
+                "backend": default_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 time
 
 from datetime import datetime
 from pathlib import Path
@@ -11,12 +13,18 @@ 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.scheduler_factory import SchedulerFactory
+from memos.mem_scheduler.schemas.general_schemas import (
+    ADD_LABEL,
+    ANSWER_LABEL,
+    QUERY_LABEL,
+)
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 from memos.mem_user.user_manager import UserManager, UserRole
 from memos.memories.activation.item import ActivationMemoryItem
 from memos.memories.parametric.item import ParametricMemoryItem
 from memos.memories.textual.item import TextualMemoryItem, TextualMemoryMetadata
+from memos.templates.mos_prompts import QUERY_REWRITING_PROMPT
 from memos.types import ChatHistory, MessageList, MOSSearchResult
 
 
@@ -30,7 +38,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 +47,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):
@@ -47,10 +60,15 @@ class MOSCore:
                 f"User '{self.user_id}' does not exist or is inactive. Please create user first."
             )
 
-        # Lazy initialization marker
+        # Initialize mem_scheduler
         self._mem_scheduler_lock = Lock()
         self.enable_mem_scheduler = self.config.get("enable_mem_scheduler", False)
-        self._mem_scheduler = None
+        if self.enable_mem_scheduler:
+            self._mem_scheduler = self._initialize_mem_scheduler()
+            self._mem_scheduler.mem_cubes = self.mem_cubes
+        else:
+            self._mem_scheduler: GeneralScheduler = None
+
         logger.info(f"MOS initialized for user: {self.user_id}")
 
     @property
@@ -58,6 +76,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,26 +93,43 @@ 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")
             else:
                 logger.debug("Memory scheduler cleared")
 
-    def _initialize_mem_scheduler(self):
+    def _initialize_mem_scheduler(self) -> GeneralScheduler:
         """Initialize the memory scheduler on first access."""
         if not self.config.enable_mem_scheduler:
             logger.debug("Memory scheduler is disabled in config")
             self._mem_scheduler = None
+            return self._mem_scheduler
         elif not hasattr(self.config, "mem_scheduler"):
             logger.error("Config of Memory scheduler is not available")
             self._mem_scheduler = None
+            return self._mem_scheduler
         else:
             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, process_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()
+            return self._mem_scheduler
 
     def mem_scheduler_on(self) -> bool:
         if not self.config.enable_mem_scheduler or self._mem_scheduler is None:
@@ -123,6 +159,25 @@ 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()
+                mem_cube.text_mem.memory_manager.wait_reorganizer()
+
+    def mem_reorganizer_wait(self) -> bool:
+        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.wait_reorganizer()
+
     def _register_chat_history(self, user_id: str | None = None) -> None:
         """Initialize chat history with user ID."""
         if user_id is None:
@@ -186,12 +241,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.
@@ -224,12 +283,20 @@ class MOSCore:
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 
-                memories = mem_cube.text_mem.search(query, top_k=self.config.top_k)
+                memories = mem_cube.text_mem.search(
+                    query,
+                    top_k=self.config.top_k,
+                    info={
+                        "user_id": target_user_id,
+                        "session_id": self.session_id,
+                        "chat_history": chat_history.chat_history,
+                    },
+                )
                 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 +328,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 +344,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 +450,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 +466,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 +489,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 +510,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 +534,13 @@ 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,
+        mode: Literal["fast", "fine"] = "fast",
+        internet_search: bool = False,
     ) -> MOSSearchResult:
         """
         Search for textual memories across all registered MemCubes.
@@ -451,6 +564,10 @@ class MOSCore:
         logger.info(
             f"User {target_user_id} has access to {len(user_cube_ids)} cubes: {user_cube_ids}"
         )
+        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]
+
         result: MOSSearchResult = {
             "text_mem": [],
             "act_mem": [],
@@ -464,20 +581,25 @@ 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)
+                time_start = time.time()
+                memories = mem_cube.text_mem.search(
+                    query,
+                    top_k=top_k if top_k else self.config.top_k,
+                    mode=mode,
+                    manual_close_internet=not internet_search,
+                    info={
+                        "user_id": target_user_id,
+                        "session_id": self.session_id,
+                        "chat_history": chat_history.chat_history,
+                    },
+                )
                 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"
                 )
-            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]})
+                search_time_end = time.time()
                 logger.info(
-                    f"🧠 [Memory] Searched memories from {mem_cube_id}:\n{self._str_memories(memories)}\n"
+                    f"time search graph: search graph time user_id: {target_user_id} time is: {search_time_end - time_start}"
                 )
         return result
 
@@ -501,6 +623,7 @@ class MOSCore:
             user_id (str, optional): The identifier of the user to add the memories to.
                 If None, the default user is used.
         """
+        # user input messages
         assert (messages is not None) or (memory_content is not None) or (doc_path is not None), (
             "messages_or_doc_path or memory_content or doc_path must be provided."
         )
@@ -540,8 +663,29 @@ class MOSCore:
                     type="chat",
                     info={"user_id": target_user_id, "session_id": self.session_id},
                 )
+
+                mem_ids = []
                 for mem in memories:
-                    self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_ids.extend(mem_id_list)
+                    logger.info(
+                        f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_id_list}"
+                    )
+
+                # submit messages for scheduler
+                if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                    mem_cube = self.mem_cubes[mem_cube_id]
+                    message_item = ScheduleMessageItem(
+                        user_id=target_user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=mem_cube,
+                        label=ADD_LABEL,
+                        content=json.dumps(mem_ids),
+                        timestamp=datetime.now(),
+                    )
+                    self.mem_scheduler.submit_messages(messages=[message_item])
+
+        # user profile
         if (
             (memory_content is not None)
             and self.config.enable_textual_memory
@@ -556,34 +700,66 @@ class MOSCore:
                 )
             else:
                 messages_list = [
-                    [
-                        {"role": "user", "content": memory_content},
-                        {
-                            "role": "assistant",
-                            "content": "",
-                        },  # add by str to keep the format,assistant role is empty
-                    ]
-                ]
+                    [{"role": "user", "content": memory_content}]
+                ]  # for only user-str input and convert message
                 memories = self.mem_reader.get_memory(
                     messages_list,
                     type="chat",
                     info={"user_id": target_user_id, "session_id": self.session_id},
                 )
+
+                mem_ids = []
                 for mem in memories:
-                    self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    logger.info(
+                        f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_id_list}"
+                    )
+                    mem_ids.extend(mem_id_list)
+
+                # submit messages for scheduler
+                if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                    mem_cube = self.mem_cubes[mem_cube_id]
+                    message_item = ScheduleMessageItem(
+                        user_id=target_user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=mem_cube,
+                        label=ADD_LABEL,
+                        content=json.dumps(mem_ids),
+                        timestamp=datetime.now(),
+                    )
+                    self.mem_scheduler.submit_messages(messages=[message_item])
+
+        # user doc input
         if (
             (doc_path is not None)
             and self.config.enable_textual_memory
             and self.mem_cubes[mem_cube_id].text_mem
         ):
             documents = self._get_all_documents(doc_path)
-            doc_memory = self.mem_reader.get_memory(
+            doc_memories = self.mem_reader.get_memory(
                 documents,
                 type="doc",
                 info={"user_id": target_user_id, "session_id": self.session_id},
             )
-            for mem in doc_memory:
-                self.mem_cubes[mem_cube_id].text_mem.add(mem)
+
+            mem_ids = []
+            for mem in doc_memories:
+                mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                mem_ids.extend(mem_id_list)
+
+            # submit messages for scheduler
+            if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                mem_cube = self.mem_cubes[mem_cube_id]
+                message_item = ScheduleMessageItem(
+                    user_id=target_user_id,
+                    mem_cube_id=mem_cube_id,
+                    mem_cube=mem_cube,
+                    label=ADD_LABEL,
+                    content=json.dumps(mem_ids),
+                    timestamp=datetime.now(),
+                )
+                self.mem_scheduler.submit_messages(messages=[message_item])
+
         logger.info(f"Add memory to {mem_cube_id} successfully")
 
     def get(
@@ -817,3 +993,27 @@ class MOSCore:
             raise ValueError(f"Target user '{target_user_id}' does not exist or is inactive.")
 
         return self.user_manager.add_user_to_cube(target_user_id, cube_id)
+
+    def get_query_rewrite(self, query: str, user_id: str | None = None):
+        """
+        Rewrite user's query according the context.
+        Args:
+            query (str): The search query that needs rewriting.
+            user_id(str, optional): The identifier of the user that the query belongs to.
+                If None, the default user is used.
+
+        Returns:
+            str: query after rewriting process.
+        """
+        target_user_id = user_id if user_id is not None else self.user_id
+        chat_history = self.chat_history_manager[target_user_id]
+
+        dialogue = "————{}".format("\n————".join(chat_history.chat_history))
+        user_prompt = QUERY_REWRITING_PROMPT.format(dialogue=dialogue, query=query)
+        messages = {"role": "user", "content": user_prompt}
+        rewritten_result = self.chat_llm.generate(messages=messages)
+        rewritten_result = json.loads(rewritten_result)
+        if rewritten_result.get("former_dialogue_related", False):
+            rewritten_query = rewritten_result.get("rewritten_question")
+            return rewritten_query if len(rewritten_query) > 0 else query
+        return query