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

memos/mem_os/product.py

@@ -2,7 +2,6 @@ import asyncio
 import json
 import os
 import random
-import threading
 import time
 
 from collections.abc import Generator
@@ -14,6 +13,7 @@ from transformers import AutoTokenizer
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
+from memos.context.context import ContextThread
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_os.core import MOSCore
@@ -46,6 +46,7 @@ from memos.templates.mos_prompts import (
     get_memos_prompt,
 )
 from memos.types import MessageList
+from memos.utils import timed
 
 
 logger = get_logger(__name__)
@@ -257,6 +258,7 @@ class MOSProduct(MOSCore):
         except Exception as e:
             logger.error(f"Error pre-loading cubes for user {user_id}: {e}", exc_info=True)
 
+    @timed
     def _load_user_cubes(
         self, user_id: str, default_cube_config: GeneralMemCubeConfig | None = None
     ) -> None:
@@ -288,6 +290,7 @@ class MOSProduct(MOSCore):
                         )
                 except Exception as e:
                     logger.error(f"Failed to load cube {cube.cube_id} for user {user_id}: {e}")
+        logger.info(f"load user {user_id} cubes successfully")
 
     def _ensure_user_instance(self, user_id: str, max_instances: int | None = None) -> None:
         """
@@ -694,8 +697,8 @@ class MOSProduct(MOSCore):
                 else None
             )
         except RuntimeError:
-            # No event loop, run in a new thread
-            thread = threading.Thread(
+            # No event loop, run in a new thread with context propagation
+            thread = ContextThread(
                 target=run_async_in_thread,
                 name=f"PostChatProcessing-{user_id}",
                 # Set as a daemon thread to avoid blocking program exit
@@ -775,10 +778,14 @@ class MOSProduct(MOSCore):
                 return
 
             # Create MemCube from path
+            time_start = time.time()
             if os.path.exists(mem_cube_name_or_path):
                 mem_cube = GeneralMemCube.init_from_dir(
                     mem_cube_name_or_path, memory_types, default_config
                 )
+                logger.info(
+                    f"time register_mem_cube: init_from_dir time is: {time.time() - time_start}"
+                )
             else:
                 logger.warning(
                     f"MemCube {mem_cube_name_or_path} does not exist, try to init from remote repo."
@@ -791,7 +798,10 @@ class MOSProduct(MOSCore):
         logger.info(
             f"Registering MemCube {mem_cube_id} with cube config {mem_cube.config.model_dump(mode='json')}"
         )
+        time_start = time.time()
         self.mem_cubes[mem_cube_id] = mem_cube
+        time_end = time.time()
+        logger.info(f"time register_mem_cube: add mem_cube time is: {time_end - time_start}")
 
     def user_register(
         self,
@@ -801,6 +811,7 @@ class MOSProduct(MOSCore):
         interests: str | None = None,
         default_mem_cube: GeneralMemCube | None = None,
         default_cube_config: GeneralMemCubeConfig | None = None,
+        mem_cube_id: str | None = None,
     ) -> dict[str, str]:
         """Register a new user with configuration and default cube.
 
@@ -836,15 +847,19 @@ class MOSProduct(MOSCore):
             default_cube_name = f"{user_name}_{user_id}_default_cube"
             mem_cube_name_or_path = os.path.join(CUBE_PATH, default_cube_name)
             default_cube_id = self.create_cube_for_user(
-                cube_name=default_cube_name, owner_id=user_id, cube_path=mem_cube_name_or_path
+                cube_name=default_cube_name,
+                owner_id=user_id,
+                cube_path=mem_cube_name_or_path,
+                cube_id=mem_cube_id,
             )
-
+            time_start = time.time()
             if default_mem_cube:
                 try:
-                    default_mem_cube.dump(mem_cube_name_or_path)
+                    default_mem_cube.dump(mem_cube_name_or_path, memory_types=[])
                 except Exception as e:
                     logger.error(f"Failed to dump default cube: {e}")
-
+            time_end = time.time()
+            logger.info(f"time user_register: dump default cube time is: {time_end - time_start}")
             # Register the default cube with MOS
             self.register_mem_cube(
                 mem_cube_name_or_path_or_object=default_mem_cube,
@@ -924,6 +939,7 @@ class MOSProduct(MOSCore):
         moscube: bool = False,
         top_k: int = 10,
         threshold: float = 0.5,
+        session_id: str | None = None,
     ) -> str:
         """
         Chat with LLM with memory references and complete response.
@@ -938,6 +954,7 @@ class MOSProduct(MOSCore):
             mode="fine",
             internet_search=internet_search,
             moscube=moscube,
+            session_id=session_id,
         )["text_mem"]
 
         memories_list = []
@@ -982,6 +999,7 @@ class MOSProduct(MOSCore):
         top_k: int = 20,
         internet_search: bool = False,
         moscube: bool = False,
+        session_id: str | None = None,
     ) -> Generator[str, None, None]:
         """
         Chat with LLM with memory references and streaming output.
@@ -1008,6 +1026,7 @@ class MOSProduct(MOSCore):
             mode="fine",
             internet_search=internet_search,
             moscube=moscube,
+            session_id=session_id,
         )["text_mem"]
 
         yield f"data: {json.dumps({'type': 'status', 'data': '1'})}\n\n"
@@ -1028,7 +1047,7 @@ class MOSProduct(MOSCore):
         system_prompt = self._build_enhance_system_prompt(user_id, memories_list)
         # Get chat history
         if user_id not in self.chat_history_manager:
-            self._register_chat_history(user_id)
+            self._register_chat_history(user_id, session_id)
 
         chat_history = self.chat_history_manager[user_id]
         if history:
@@ -1296,6 +1315,7 @@ class MOSProduct(MOSCore):
         install_cube_ids: list[str] | None = None,
         top_k: int = 10,
         mode: Literal["fast", "fine"] = "fast",
+        session_id: str | None = None,
     ):
         """Search memories for a specific user."""
 
@@ -1306,7 +1326,9 @@ class MOSProduct(MOSCore):
         logger.info(
             f"time search: load_user_cubes time user_id: {user_id} time is: {load_user_cubes_time_end - time_start}"
         )
-        search_result = super().search(query, user_id, install_cube_ids, top_k, mode=mode)
+        search_result = super().search(
+            query, user_id, install_cube_ids, top_k, mode=mode, session_id=session_id
+        )
         search_time_end = time.time()
         logger.info(
             f"time search: search text_mem time user_id: {user_id} time is: {search_time_end - load_user_cubes_time_end}"
@@ -1342,13 +1364,15 @@ class MOSProduct(MOSCore):
         mem_cube_id: str | None = None,
         source: str | None = None,
         user_profile: bool = False,
+        session_id: str | None = None,
     ):
         """Add memory for a specific user."""
 
         # Load user cubes if not already loaded
         self._load_user_cubes(user_id, self.default_cube_config)
-
-        result = super().add(messages, memory_content, doc_path, mem_cube_id, user_id)
+        result = super().add(
+            messages, memory_content, doc_path, mem_cube_id, user_id, session_id=session_id
+        )
         if user_profile:
             try:
                 user_interests = memory_content.split("'userInterests': '")[1].split("', '")[0]

memos/mem_reader/factory.py

@@ -3,6 +3,7 @@ from typing import Any, ClassVar
 from memos.configs.mem_reader import MemReaderConfigFactory
 from memos.mem_reader.base import BaseMemReader
 from memos.mem_reader.simple_struct import SimpleStructMemReader
+from memos.memos_tools.singleton import singleton_factory
 
 
 class MemReaderFactory(BaseMemReader):
@@ -13,6 +14,7 @@ class MemReaderFactory(BaseMemReader):
     }
 
     @classmethod
+    @singleton_factory()
     def from_config(cls, config_factory: MemReaderConfigFactory) -> BaseMemReader:
         backend = config_factory.backend
         if backend not in cls.backend_to_class:

memos/mem_reader/simple_struct.py

@@ -13,6 +13,7 @@ from memos import log
 from memos.chunkers import ChunkerFactory
 from memos.configs.mem_reader import SimpleStructMemReaderConfig
 from memos.configs.parser import ParserConfigFactory
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import EmbedderFactory
 from memos.llms.factory import LLMFactory
 from memos.mem_reader.base import BaseMemReader
@@ -26,6 +27,7 @@ from memos.templates.mem_reader_prompts import (
     SIMPLE_STRUCT_MEM_READER_PROMPT,
     SIMPLE_STRUCT_MEM_READER_PROMPT_ZH,
 )
+from memos.utils import timed
 
 
 logger = log.get_logger(__name__)
@@ -55,45 +57,60 @@ def detect_lang(text):
 
 def _build_node(idx, message, info, scene_file, llm, parse_json_result, embedder):
     # generate
-    raw = llm.generate(message)
-    if not raw:
+    try:
+        raw = llm.generate(message)
+        if not raw:
+            logger.warning(f"[LLM] Empty generation for input: {message}")
+            return None
+    except Exception as e:
+        logger.error(f"[LLM] Exception during generation: {e}")
         return None
 
     # parse_json_result
-    chunk_res = parse_json_result(raw)
-    if not chunk_res:
+    try:
+        chunk_res = parse_json_result(raw)
+        if not chunk_res:
+            logger.warning(f"[Parse] Failed to parse result: {raw}")
+            return None
+    except Exception as e:
+        logger.error(f"[Parse] Exception during JSON parsing: {e}")
         return None
 
-    value = chunk_res.get("value")
-    if not value:
+    try:
+        value = chunk_res.get("value", "").strip()
+        if not value:
+            logger.warning("[BuildNode] value is empty")
+            return None
+
+        tags = chunk_res.get("tags", [])
+        if not isinstance(tags, list):
+            tags = []
+
+        key = chunk_res.get("key", None)
+
+        embedding = embedder.embed([value])[0]
+
+        return TextualMemoryItem(
+            memory=value,
+            metadata=TreeNodeTextualMemoryMetadata(
+                user_id=info.get("user_id", ""),
+                session_id=info.get("session_id", ""),
+                memory_type="LongTermMemory",
+                status="activated",
+                tags=tags,
+                key=key,
+                embedding=embedding,
+                usage=[],
+                sources=[{"type": "doc", "doc_path": f"{scene_file}_{idx}"}],
+                background="",
+                confidence=0.99,
+                type="fact",
+            ),
+        )
+    except Exception as e:
+        logger.error(f"[BuildNode] Error building node: {e}")
         return None
 
-    # embed
-    embedding = embedder.embed([value])[0]
-
-    # TextualMemoryItem
-    tags = chunk_res["tags"] if isinstance(chunk_res.get("tags"), list) else []
-    key = chunk_res.get("key", None)
-
-    node_i = TextualMemoryItem(
-        memory=value,
-        metadata=TreeNodeTextualMemoryMetadata(
-            user_id=info.get("user_id"),
-            session_id=info.get("session_id"),
-            memory_type="LongTermMemory",
-            status="activated",
-            tags=tags,
-            key=key,
-            embedding=embedding,
-            usage=[],
-            sources=[f"{scene_file}_{idx}"],
-            background="",
-            confidence=0.99,
-            type="fact",
-        ),
-    )
-    return node_i
-
 
 class SimpleStructMemReader(BaseMemReader, ABC):
     """Naive implementation of MemReader."""
@@ -110,44 +127,77 @@ class SimpleStructMemReader(BaseMemReader, ABC):
         self.embedder = EmbedderFactory.from_config(config.embedder)
         self.chunker = ChunkerFactory.from_config(config.chunker)
 
+    @timed
     def _process_chat_data(self, scene_data_info, info):
-        lang = detect_lang("\n".join(scene_data_info))
+        mem_list = []
+        for item in scene_data_info:
+            if "chat_time" in item:
+                mem = item["role"] + ": " + f"[{item['chat_time']}]: " + item["content"]
+                mem_list.append(mem)
+            else:
+                mem = item["role"] + ":" + item["content"]
+                mem_list.append(mem)
+        lang = detect_lang("\n".join(mem_list))
         template = PROMPT_DICT["chat"][lang]
         examples = PROMPT_DICT["chat"][f"{lang}_example"]
 
-        prompt = template.replace("${conversation}", "\n".join(scene_data_info))
+        prompt = template.replace("${conversation}", "\n".join(mem_list))
         if self.config.remove_prompt_example:
             prompt = prompt.replace(examples, "")
 
         messages = [{"role": "user", "content": prompt}]
 
-        response_text = self.llm.generate(messages)
-        response_json = self.parse_json_result(response_text)
+        try:
+            response_text = self.llm.generate(messages)
+            response_json = self.parse_json_result(response_text)
+        except Exception as e:
+            logger.error(f"[LLM] Exception during chat generation: {e}")
+            response_json = {
+                "memory list": [
+                    {
+                        "key": "\n".join(mem_list)[:10],
+                        "memory_type": "UserMemory",
+                        "value": "\n".join(mem_list),
+                        "tags": [],
+                    }
+                ],
+                "summary": "\n".join(mem_list),
+            }
 
         chat_read_nodes = []
         for memory_i_raw in response_json.get("memory list", []):
-            node_i = TextualMemoryItem(
-                memory=memory_i_raw.get("value", ""),
-                metadata=TreeNodeTextualMemoryMetadata(
-                    user_id=info.get("user_id"),
-                    session_id=info.get("session_id"),
-                    memory_type=memory_i_raw.get("memory_type", "")
+            try:
+                memory_type = (
+                    memory_i_raw.get("memory_type", "LongTermMemory")
                     .replace("长期记忆", "LongTermMemory")
-                    .replace("用户记忆", "UserMemory"),
-                    status="activated",
-                    tags=memory_i_raw.get("tags", [])
-                    if type(memory_i_raw.get("tags", [])) is list
-                    else [],
-                    key=memory_i_raw.get("key", ""),
-                    embedding=self.embedder.embed([memory_i_raw.get("value", "")])[0],
-                    usage=[],
-                    sources=scene_data_info,
-                    background=response_json.get("summary", ""),
-                    confidence=0.99,
-                    type="fact",
-                ),
-            )
-            chat_read_nodes.append(node_i)
+                    .replace("用户记忆", "UserMemory")
+                )
+
+                if memory_type not in ["LongTermMemory", "UserMemory"]:
+                    memory_type = "LongTermMemory"
+
+                node_i = TextualMemoryItem(
+                    memory=memory_i_raw.get("value", ""),
+                    metadata=TreeNodeTextualMemoryMetadata(
+                        user_id=info.get("user_id"),
+                        session_id=info.get("session_id"),
+                        memory_type=memory_type,
+                        status="activated",
+                        tags=memory_i_raw.get("tags", [])
+                        if type(memory_i_raw.get("tags", [])) is list
+                        else [],
+                        key=memory_i_raw.get("key", ""),
+                        embedding=self.embedder.embed([memory_i_raw.get("value", "")])[0],
+                        usage=[],
+                        sources=scene_data_info,
+                        background=response_json.get("summary", ""),
+                        confidence=0.99,
+                        type="fact",
+                    ),
+                )
+                chat_read_nodes.append(node_i)
+            except Exception as e:
+                logger.error(f"[ChatReader] Error parsing memory item: {e}")
 
         return chat_read_nodes
 
@@ -200,8 +250,8 @@ class SimpleStructMemReader(BaseMemReader, ABC):
         else:
             processing_func = self._process_doc_data
 
-        # Process Q&A pairs concurrently
-        with concurrent.futures.ThreadPoolExecutor() as executor:
+        # Process Q&A pairs concurrently with context propagation
+        with ContextThreadPoolExecutor() as executor:
             futures = [
                 executor.submit(processing_func, scene_data_info, info)
                 for scene_data_info in list_scene_data_info
@@ -239,11 +289,9 @@ class SimpleStructMemReader(BaseMemReader, ABC):
                 for item in items:
                     # Convert dictionary to string
                     if "chat_time" in item:
-                        mem = item["role"] + ": " + f"[{item['chat_time']}]: " + item["content"]
-                        result.append(mem)
+                        result.append(item)
                     else:
-                        mem = item["role"] + ":" + item["content"]
-                        result.append(mem)
+                        result.append(item)
                     if len(result) >= 10:
                         results.append(result)
                         context = copy.deepcopy(result[-2:])
@@ -254,17 +302,21 @@ class SimpleStructMemReader(BaseMemReader, ABC):
             for item in scene_data:
                 try:
                     if os.path.exists(item):
-                        parsed_text = parser.parse(item)
-                        results.append({"file": "pure_text", "text": parsed_text})
+                        try:
+                            parsed_text = parser.parse(item)
+                            results.append({"file": item, "text": parsed_text})
+                        except Exception as e:
+                            logger.error(f"[SceneParser] Error parsing {item}: {e}")
+                            continue
                     else:
                         parsed_text = item
-                        results.append({"file": item, "text": parsed_text})
+                        results.append({"file": "pure_text", "text": parsed_text})
                 except Exception as e:
                     print(f"Error parsing file {item}: {e!s}")
 
         return results
 
-    def _process_doc_data(self, scene_data_info, info):
+    def _process_doc_data(self, scene_data_info, info, **kwargs):
         chunks = self.chunker.chunk(scene_data_info["text"])
         messages = []
         for chunk in chunks:
@@ -277,7 +329,7 @@ class SimpleStructMemReader(BaseMemReader, ABC):
         doc_nodes = []
         scene_file = scene_data_info["file"]
 
-        with concurrent.futures.ThreadPoolExecutor(max_workers=50) as executor:
+        with ContextThreadPoolExecutor(max_workers=50) as executor:
             futures = {
                 executor.submit(
                     _build_node,
@@ -302,6 +354,7 @@ class SimpleStructMemReader(BaseMemReader, ABC):
                         doc_nodes.append(node)
                 except Exception as e:
                     tqdm.write(f"[ERROR] {e}")
+                    logger.error(f"[DocReader] Future task failed: {e}")
         return doc_nodes
 
     def parse_json_result(self, response_text):
@@ -309,14 +362,14 @@ class SimpleStructMemReader(BaseMemReader, ABC):
             json_start = response_text.find("{")
             response_text = response_text[json_start:]
             response_text = response_text.replace("```", "").strip()
-            if response_text[-1] != "}":
+            if not response_text.endswith("}"):
                 response_text += "}"
-            response_json = json.loads(response_text)
-            return response_json
+            return json.loads(response_text)
         except json.JSONDecodeError as e:
-            logger.warning(
-                f"Failed to parse LLM response as JSON: {e}\nRaw response:\n{response_text}"
-            )
+            logger.error(f"[JSONParse] Failed to decode JSON: {e}\nRaw:\n{response_text}")
+            return {}
+        except Exception as e:
+            logger.error(f"[JSONParse] Unexpected error: {e}")
             return {}
 
     def transform_memreader(self, data: dict) -> list[TextualMemoryItem]: