MemoryOS 1.0.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

memos/graph_dbs/neo4j.py

@@ -1,3 +1,4 @@
+import json
 import time
 
 from datetime import datetime
@@ -174,6 +175,12 @@ class Neo4jGraphDB(BaseGraphDB):
                 n.updated_at = datetime($updated_at),
                 n += $metadata
         """
+
+        # serialization
+        if metadata["sources"]:
+            for idx in range(len(metadata["sources"])):
+                metadata["sources"][idx] = json.dumps(metadata["sources"][idx])
+
         with self.driver.session(database=self.db_name) as session:
             session.run(
                 query,
@@ -606,6 +613,7 @@ class Neo4jGraphDB(BaseGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        search_filter: dict | None = None,
         **kwargs,
     ) -> list[dict]:
         """
@@ -618,6 +626,8 @@ class Neo4jGraphDB(BaseGraphDB):
             status (str, optional): Node status filter (e.g., 'active', 'archived').
                             If provided, restricts results to nodes with matching status.
             threshold (float, optional): Minimum similarity score threshold (0 ~ 1).
+            search_filter (dict, optional): Additional metadata filters for search results.
+                            Keys should match node properties, values are the expected values.
 
         Returns:
             list[dict]: A list of dicts with 'id' and 'score', ordered by similarity.
@@ -627,6 +637,7 @@ class Neo4jGraphDB(BaseGraphDB):
             - If scope is provided, it restricts results to nodes with matching memory_type.
             - If 'status' is provided, only nodes with the matching status will be returned.
             - If threshold is provided, only results with score >= threshold will be returned.
+            - If search_filter is provided, additional WHERE clauses will be added for metadata filtering.
             - Typical use case: restrict to 'status = activated' to avoid
             matching archived or merged nodes.
         """
@@ -639,6 +650,12 @@ class Neo4jGraphDB(BaseGraphDB):
         if not self.config.use_multi_db and self.config.user_name:
             where_clauses.append("node.user_name = $user_name")
 
+        # Add search_filter conditions
+        if search_filter:
+            for key, _ in search_filter.items():
+                param_name = f"filter_{key}"
+                where_clauses.append(f"node.{key} = ${param_name}")
+
         where_clause = ""
         if where_clauses:
             where_clause = "WHERE " + " AND ".join(where_clauses)
@@ -650,7 +667,8 @@ class Neo4jGraphDB(BaseGraphDB):
             RETURN node.id AS id, score
         """
 
-        parameters = {"embedding": vector, "k": top_k, "scope": scope}
+        parameters = {"embedding": vector, "k": top_k}
+
         if scope:
             parameters["scope"] = scope
         if status:
@@ -661,6 +679,12 @@ class Neo4jGraphDB(BaseGraphDB):
             else:
                 parameters["user_name"] = self.config.user_name
 
+        # Add search_filter parameters
+        if search_filter:
+            for key, value in search_filter.items():
+                param_name = f"filter_{key}"
+                parameters[param_name] = value
+
         with self.driver.session(database=self.db_name) as session:
             result = session.run(query, parameters)
             records = [{"id": record["id"], "score": record["score"]} for record in result]
@@ -1111,4 +1135,14 @@ class Neo4jGraphDB(BaseGraphDB):
                 node[time_field] = node[time_field].isoformat()
         node.pop("user_name", None)
 
+        # serialization
+        if node["sources"]:
+            for idx in range(len(node["sources"])):
+                if not (
+                    isinstance(node["sources"][idx], str)
+                    and node["sources"][idx][0] == "{"
+                    and node["sources"][idx][0] == "}"
+                ):
+                    break
+                node["sources"][idx] = json.loads(node["sources"][idx])
         return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}

memos/graph_dbs/neo4j_community.py

@@ -129,6 +129,7 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        search_filter: dict | None = None,
         **kwargs,
     ) -> list[dict]:
         """
@@ -140,6 +141,7 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
             scope (str, optional): Memory type filter (e.g., 'WorkingMemory', 'LongTermMemory').
             status (str, optional): Node status filter (e.g., 'activated', 'archived').
             threshold (float, optional): Minimum similarity score threshold (0 ~ 1).
+            search_filter (dict, optional): Additional metadata filters to apply.
 
         Returns:
             list[dict]: A list of dicts with 'id' and 'score', ordered by similarity.
@@ -149,6 +151,7 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
             - If 'scope' is provided, it restricts results to nodes with matching memory_type.
             - If 'status' is provided, it further filters nodes by status.
             - If 'threshold' is provided, only results with score >= threshold will be returned.
+            - If 'search_filter' is provided, it applies additional metadata-based filtering.
             - The returned IDs can be used to fetch full node data from Neo4j if needed.
         """
         # Build VecDB filter
@@ -163,6 +166,10 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
         else:
             vec_filter["user_name"] = self.config.user_name
 
+        # Add search_filter conditions
+        if search_filter:
+            vec_filter.update(search_filter)
+
         # Perform vector search
         results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter)
 

memos/llms/factory.py

@@ -9,6 +9,7 @@ from memos.llms.ollama import OllamaLLM
 from memos.llms.openai import AzureLLM, OpenAILLM
 from memos.llms.qwen import QwenLLM
 from memos.llms.vllm import VLLMLLM
+from memos.memos_tools.singleton import singleton_factory
 
 
 class LLMFactory(BaseLLM):
@@ -26,6 +27,7 @@ class LLMFactory(BaseLLM):
     }
 
     @classmethod
+    @singleton_factory()
     def from_config(cls, config_factory: LLMConfigFactory) -> BaseLLM:
         backend = config_factory.backend
         if backend not in cls.backend_to_class:

memos/llms/openai.py

@@ -1,4 +1,8 @@
+import hashlib
+import json
+
 from collections.abc import Generator
+from typing import ClassVar
 
 import openai
 
@@ -13,11 +17,44 @@ logger = get_logger(__name__)
 
 
 class OpenAILLM(BaseLLM):
-    """OpenAI LLM class."""
+    """OpenAI LLM class with singleton pattern."""
+
+    _instances: ClassVar[dict] = {}  # Class variable to store instances
+
+    def __new__(cls, config: OpenAILLMConfig) -> "OpenAILLM":
+        config_hash = cls._get_config_hash(config)
+
+        if config_hash not in cls._instances:
+            logger.info(f"Creating new OpenAI LLM instance for config hash: {config_hash}")
+            instance = super().__new__(cls)
+            cls._instances[config_hash] = instance
+        else:
+            logger.info(f"Reusing existing OpenAI LLM instance for config hash: {config_hash}")
+
+        return cls._instances[config_hash]
 
     def __init__(self, config: OpenAILLMConfig):
+        # Avoid duplicate initialization
+        if hasattr(self, "_initialized"):
+            return
+
         self.config = config
         self.client = openai.Client(api_key=config.api_key, base_url=config.api_base)
+        self._initialized = True
+        logger.info("OpenAI LLM instance initialized")
+
+    @classmethod
+    def _get_config_hash(cls, config: OpenAILLMConfig) -> str:
+        """Generate hash value of configuration"""
+        config_dict = config.model_dump()
+        config_str = json.dumps(config_dict, sort_keys=True)
+        return hashlib.md5(config_str.encode()).hexdigest()
+
+    @classmethod
+    def clear_cache(cls):
+        """Clear all cached instances"""
+        cls._instances.clear()
+        logger.info("OpenAI LLM instance cache cleared")
 
     def generate(self, messages: MessageList) -> str:
         """Generate a response from OpenAI LLM."""
@@ -71,15 +108,50 @@ class OpenAILLM(BaseLLM):
 
 
 class AzureLLM(BaseLLM):
-    """Azure OpenAI LLM class."""
+    """Azure OpenAI LLM class with singleton pattern."""
+
+    _instances: ClassVar[dict] = {}  # Class variable to store instances
+
+    def __new__(cls, config: AzureLLMConfig):
+        # Generate hash value of config as cache key
+        config_hash = cls._get_config_hash(config)
+
+        if config_hash not in cls._instances:
+            logger.info(f"Creating new Azure LLM instance for config hash: {config_hash}")
+            instance = super().__new__(cls)
+            cls._instances[config_hash] = instance
+        else:
+            logger.info(f"Reusing existing Azure LLM instance for config hash: {config_hash}")
+
+        return cls._instances[config_hash]
 
     def __init__(self, config: AzureLLMConfig):
+        # Avoid duplicate initialization
+        if hasattr(self, "_initialized"):
+            return
+
         self.config = config
         self.client = openai.AzureOpenAI(
             azure_endpoint=config.base_url,
             api_version=config.api_version,
             api_key=config.api_key,
         )
+        self._initialized = True
+        logger.info("Azure LLM instance initialized")
+
+    @classmethod
+    def _get_config_hash(cls, config: AzureLLMConfig) -> str:
+        """Generate hash value of configuration"""
+        # Convert config to dict and sort to ensure consistency
+        config_dict = config.model_dump()
+        config_str = json.dumps(config_dict, sort_keys=True)
+        return hashlib.md5(config_str.encode()).hexdigest()
+
+    @classmethod
+    def clear_cache(cls):
+        """Clear all cached instances"""
+        cls._instances.clear()
+        logger.info("Azure LLM instance cache cleared")
 
     def generate(self, messages: MessageList) -> str:
         """Generate a response from Azure OpenAI LLM."""

memos/log.py

@@ -2,7 +2,9 @@ import atexit
 import logging
 import os
 import threading
+import time
 
+from concurrent.futures import ThreadPoolExecutor
 from logging.config import dictConfig
 from pathlib import Path
 from sys import stdout
@@ -12,8 +14,7 @@ import requests
 from dotenv import load_dotenv
 
 from memos import settings
-from memos.api.context.context import get_current_trace_id
-from memos.api.context.context_thread import ContextThreadPoolExecutor
+from memos.context.context import get_current_api_path, get_current_trace_id
 
 
 # Load environment variables
@@ -39,9 +40,9 @@ class TraceIDFilter(logging.Filter):
     def filter(self, record):
         try:
             trace_id = get_current_trace_id()
-            record.trace_id = trace_id if trace_id else "no-trace-id"
+            record.trace_id = trace_id if trace_id else "trace-id"
         except Exception:
-            record.trace_id = "no-trace-id"
+            record.trace_id = "trace-id"
         return True
 
 
@@ -65,7 +66,7 @@ class CustomLoggerRequestHandler(logging.Handler):
         if not self._initialized:
             super().__init__()
             workers = int(os.getenv("CUSTOM_LOGGER_WORKERS", "2"))
-            self._executor = ContextThreadPoolExecutor(
+            self._executor = ThreadPoolExecutor(
                 max_workers=workers, thread_name_prefix="log_sender"
             )
             self._is_shutting_down = threading.Event()
@@ -78,21 +79,32 @@ class CustomLoggerRequestHandler(logging.Handler):
         if os.getenv("CUSTOM_LOGGER_URL") is None or self._is_shutting_down.is_set():
             return
 
+        # Only process INFO and ERROR level logs
+        if record.levelno < logging.INFO:  # Skip DEBUG and lower
+            return
+
         try:
-            trace_id = get_current_trace_id() or "no-trace-id"
-            self._executor.submit(self._send_log_sync, record.getMessage(), trace_id)
+            trace_id = get_current_trace_id() or "trace-id"
+            api_path = get_current_api_path()
+            if api_path is not None:
+                self._executor.submit(self._send_log_sync, record.getMessage(), trace_id, api_path)
         except Exception as e:
             if not self._is_shutting_down.is_set():
                 print(f"Error sending log: {e}")
 
-    def _send_log_sync(self, message, trace_id):
+    def _send_log_sync(self, message, trace_id, api_path):
         """Send log message synchronously in a separate thread"""
         try:
             logger_url = os.getenv("CUSTOM_LOGGER_URL")
             token = os.getenv("CUSTOM_LOGGER_TOKEN")
 
             headers = {"Content-Type": "application/json"}
-            post_content = {"message": message, "trace_id": trace_id}
+            post_content = {
+                "message": message,
+                "trace_id": trace_id,
+                "action": api_path,
+                "current_time": round(time.time(), 3),
+            }
 
             # Add auth token if exists
             if token:
@@ -139,7 +151,7 @@ LOGGING_CONFIG = {
             "format": "[%(trace_id)s] - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
         },
         "simplified": {
-            "format": "%(asctime)s | %(trace_id)s | %(levelname)s | %(filename)s | %(message)s"
+            "format": "%(asctime)s | %(trace_id)s | %(levelname)s | %(filename)s:%(lineno)d: %(funcName)s | %(message)s"
         },
     },
     "filters": {
@@ -151,7 +163,7 @@ LOGGING_CONFIG = {
             "level": selected_log_level,
             "class": "logging.StreamHandler",
             "stream": stdout,
-            "formatter": "simplified",
+            "formatter": "no_datetime",
             "filters": ["package_tree_filter", "trace_id_filter"],
         },
         "file": {
@@ -160,18 +172,18 @@ LOGGING_CONFIG = {
             "filename": _setup_logfile(),
             "maxBytes": 1024**2 * 10,
             "backupCount": 10,
-            "formatter": "simplified",
+            "formatter": "standard",
             "filters": ["trace_id_filter"],
         },
         "custom_logger": {
-            "level": selected_log_level,
+            "level": "INFO",
             "class": "memos.log.CustomLoggerRequestHandler",
             "formatter": "simplified",
         },
     },
     "root": {  # Root logger handles all logs
-        "level": selected_log_level,
-        "handlers": ["console", "file", "custom_logger"],
+        "level": logging.DEBUG if settings.DEBUG else logging.INFO,
+        "handlers": ["console", "file"],
     },
     "loggers": {
         "memos": {

memos/mem_cube/general.py

@@ -1,4 +1,5 @@
 import os
+import time
 
 from typing import Literal
 
@@ -23,11 +24,13 @@ class GeneralMemCube(BaseMemCube):
     def __init__(self, config: GeneralMemCubeConfig):
         """Initialize the MemCube with a configuration."""
         self.config = config
+        time_start = time.time()
         self._text_mem: BaseTextMemory | None = (
             MemoryFactory.from_config(config.text_mem)
             if config.text_mem.backend != "uninitialized"
             else None
         )
+        logger.info(f"init_text_mem in {time.time() - time_start} seconds")
         self._act_mem: BaseActMemory | None = (
             MemoryFactory.from_config(config.act_mem)
             if config.act_mem.backend != "uninitialized"
@@ -137,7 +140,6 @@ class GeneralMemCube(BaseMemCube):
         if default_config is not None:
             config = merge_config_with_default(config, default_config)
             logger.info(f"Applied default config to cube {config.cube_id}")
-
         mem_cube = GeneralMemCube(config)
         mem_cube.load(dir, memory_types)
         return mem_cube

memos/mem_os/core.py

@@ -24,7 +24,7 @@ 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.memos_tools.thread_safe_dict import ThreadSafeDict
+from memos.memos_tools.thread_safe_dict_segment import OptimizedThreadSafeDict
 from memos.templates.mos_prompts import QUERY_REWRITING_PROMPT
 from memos.types import ChatHistory, MessageList, MOSSearchResult
 
@@ -47,8 +47,8 @@ class MOSCore:
         self.mem_reader = MemReaderFactory.from_config(config.mem_reader)
         self.chat_history_manager: dict[str, ChatHistory] = {}
         # use thread safe dict for multi-user product-server scenario
-        self.mem_cubes: ThreadSafeDict[str, GeneralMemCube] = (
-            ThreadSafeDict() if user_manager is not None else {}
+        self.mem_cubes: OptimizedThreadSafeDict[str, GeneralMemCube] = (
+            OptimizedThreadSafeDict() if user_manager is not None else {}
         )
         self._register_chat_history()
 
@@ -125,12 +125,16 @@ class MOSCore:
                     "missing required 'llm' attribute"
                 )
                 self._mem_scheduler.initialize_modules(
-                    chat_llm=self.chat_llm, process_llm=self.chat_llm
+                    chat_llm=self.chat_llm,
+                    process_llm=self.chat_llm,
+                    db_engine=self.user_manager.engine,
                 )
             else:
                 # Configure scheduler general_modules
                 self._mem_scheduler.initialize_modules(
-                    chat_llm=self.chat_llm, process_llm=self.mem_reader.llm
+                    chat_llm=self.chat_llm,
+                    process_llm=self.mem_reader.llm,
+                    db_engine=self.user_manager.engine,
                 )
             self._mem_scheduler.start()
             return self._mem_scheduler
@@ -182,13 +186,13 @@ class MOSCore:
                 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:
+    def _register_chat_history(
+        self, user_id: str | None = None, session_id: str | None = None
+    ) -> None:
         """Initialize chat history with user ID."""
-        if user_id is None:
-            user_id = self.user_id
         self.chat_history_manager[user_id] = ChatHistory(
-            user_id=user_id,
-            session_id=self.session_id,
+            user_id=user_id if user_id is not None else self.user_id,
+            session_id=session_id if session_id is not None else self.session_id,
             created_at=datetime.utcnow(),
             total_messages=0,
             chat_history=[],
@@ -483,14 +487,14 @@ class MOSCore:
                 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)
+                mem_cube_obj = GeneralMemCube.init_from_dir(mem_cube_name_or_path)
+                self.mem_cubes[mem_cube_id] = mem_cube_obj
             else:
                 logger.warning(
                     f"MemCube {mem_cube_name_or_path} does not exist, try to init from remote repo."
                 )
-                self.mem_cubes[mem_cube_id] = GeneralMemCube.init_from_remote_repo(
-                    mem_cube_name_or_path
-                )
+                mem_cube_obj = GeneralMemCube.init_from_remote_repo(mem_cube_name_or_path)
+                self.mem_cubes[mem_cube_id] = mem_cube_obj
         # Check if cube already exists in database
         existing_cube = self.user_manager.get_cube(mem_cube_id)
 
@@ -547,6 +551,7 @@ class MOSCore:
         mode: Literal["fast", "fine"] = "fast",
         internet_search: bool = False,
         moscube: bool = False,
+        session_id: str | None = None,
         **kwargs,
     ) -> MOSSearchResult:
         """
@@ -562,7 +567,9 @@ class MOSCore:
         Returns:
             MemoryResult: A dictionary containing the search results.
         """
+        target_session_id = session_id if session_id is not None else self.session_id
         target_user_id = user_id if user_id is not None else self.user_id
+
         self._validate_user_exists(target_user_id)
         # Get all cubes accessible by the target user
         accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
@@ -575,6 +582,11 @@ class MOSCore:
             self._register_chat_history(target_user_id)
         chat_history = self.chat_history_manager[target_user_id]
 
+        # Create search filter if session_id is provided
+        search_filter = None
+        if session_id is not None:
+            search_filter = {"session_id": session_id}
+
         result: MOSSearchResult = {
             "text_mem": [],
             "act_mem": [],
@@ -584,9 +596,13 @@ class MOSCore:
             install_cube_ids = user_cube_ids
         # create exist dict in mem_cubes and avoid  one search slow
         tmp_mem_cubes = {}
+        time_start_cube_get = time.time()
         for mem_cube_id in install_cube_ids:
             if mem_cube_id in self.mem_cubes:
                 tmp_mem_cubes[mem_cube_id] = self.mem_cubes.get(mem_cube_id)
+        logger.info(
+            f"time search: transform cube time user_id: {target_user_id} time is: {time.time() - time_start_cube_get}"
+        )
 
         for mem_cube_id, mem_cube in tmp_mem_cubes.items():
             if (
@@ -602,10 +618,11 @@ class MOSCore:
                     manual_close_internet=not internet_search,
                     info={
                         "user_id": target_user_id,
-                        "session_id": self.session_id,
+                        "session_id": target_session_id,
                         "chat_history": chat_history.chat_history,
                     },
                     moscube=moscube,
+                    search_filter=search_filter,
                 )
                 result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories})
                 logger.info(
@@ -624,6 +641,8 @@ class MOSCore:
         doc_path: str | None = None,
         mem_cube_id: str | None = None,
         user_id: str | None = None,
+        session_id: str | None = None,
+        **kwargs,
     ) -> None:
         """
         Add textual memories to a MemCube.
@@ -636,11 +655,16 @@ class MOSCore:
                 If None, the default MemCube for the user is used.
             user_id (str, optional): The identifier of the user to add the memories to.
                 If None, the default user is used.
+            session_id (str, optional): session_id
         """
         # 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."
         )
+        # TODO: asure that session_id is a valid string
+        time_start = time.time()
+
+        target_session_id = session_id if session_id else self.session_id
         target_user_id = user_id if user_id is not None else self.user_id
         if mem_cube_id is None:
             # Try to find a default cube for the user
@@ -652,18 +676,29 @@ class MOSCore:
             mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
         else:
             self._validate_cube_access(target_user_id, mem_cube_id)
+        logger.info(
+            f"time add: get mem_cube_id time user_id: {target_user_id} time is: {time.time() - time_start}"
+        )
 
+        time_start_0 = time.time()
         if mem_cube_id not in self.mem_cubes:
             raise ValueError(f"MemCube '{mem_cube_id}' is not loaded. Please register.")
+        logger.info(
+            f"time add: get mem_cube_id check in mem_cubes time user_id: {target_user_id} time is: {time.time() - time_start_0}"
+        )
+        time_start_1 = time.time()
         if (
             (messages is not None)
             and self.config.enable_textual_memory
             and self.mem_cubes[mem_cube_id].text_mem
         ):
+            logger.info(
+                f"time add: messages is not None and enable_textual_memory and text_mem is not None time user_id: {target_user_id} time is: {time.time() - time_start_1}"
+            )
             if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                 add_memory = []
                 metadata = TextualMemoryMetadata(
-                    user_id=target_user_id, session_id=self.session_id, source="conversation"
+                    user_id=target_user_id, session_id=target_session_id, source="conversation"
                 )
                 for message in messages:
                     add_memory.append(
@@ -672,12 +707,15 @@ class MOSCore:
                 self.mem_cubes[mem_cube_id].text_mem.add(add_memory)
             else:
                 messages_list = [messages]
+                time_start_2 = time.time()
                 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": target_session_id},
+                )
+                logger.info(
+                    f"time add: get mem_reader time user_id: {target_user_id} time is: {time.time() - time_start_2}"
                 )
-
                 mem_ids = []
                 for mem in memories:
                     mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
@@ -707,7 +745,7 @@ class MOSCore:
         ):
             if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                 metadata = TextualMemoryMetadata(
-                    user_id=self.user_id, session_id=self.session_id, source="conversation"
+                    user_id=target_user_id, session_id=target_session_id, source="conversation"
                 )
                 self.mem_cubes[mem_cube_id].text_mem.add(
                     [TextualMemoryItem(memory=memory_content, metadata=metadata)]
@@ -719,7 +757,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": target_session_id},
                 )
 
                 mem_ids = []
@@ -753,7 +791,7 @@ class MOSCore:
             doc_memories = 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": target_session_id},
             )
 
             mem_ids = []
@@ -986,7 +1024,7 @@ class MOSCore:
 
     def get_user_info(self) -> dict[str, Any]:
         """Get current user information including accessible cubes.
-
+        TODO: maybe input user_id
         Returns:
             dict: User information and accessible cubes.
         """

memos/mem_os/main.py

@@ -5,6 +5,7 @@ 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
@@ -487,9 +488,7 @@ class MOS(MOSCore):
 
         # Generate answers in parallel while maintaining order
         sub_answers = [None] * len(sub_questions)
-        with concurrent.futures.ThreadPoolExecutor(
-            max_workers=min(len(sub_questions), 10)
-        ) as executor:
+        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
@@ -552,9 +551,7 @@ class MOS(MOSCore):
 
         # Search in parallel while maintaining order
         all_memories = []
-        with concurrent.futures.ThreadPoolExecutor(
-            max_workers=min(len(sub_questions), 10)
-        ) as executor:
+        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