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

memos/graph_dbs/neo4j.py

@@ -114,14 +114,14 @@ class Neo4jGraphDB(BaseGraphDB):
             )
             return result.single()["count"]
 
-    def count_nodes(self, scope: str) -> int:
+    def node_not_exist(self, scope: str) -> int:
         query = """
         MATCH (n:Memory)
         WHERE n.memory_type = $scope
         """
         if not self.config.use_multi_db and self.config.user_name:
             query += "\nAND n.user_name = $user_name"
-        query += "\nRETURN count(n) AS count"
+        query += "\nRETURN n LIMIT 1"
 
         with self.driver.session(database=self.db_name) as session:
             result = session.run(
@@ -131,7 +131,7 @@ class Neo4jGraphDB(BaseGraphDB):
                     "user_name": self.config.user_name if self.config.user_name else None,
                 },
             )
-            return result.single()["count"]
+            return result.single() is None
 
     def remove_oldest_memory(self, memory_type: str, keep_latest: int) -> None:
         """
@@ -323,14 +323,16 @@ class Neo4jGraphDB(BaseGraphDB):
             return result.single() is not None
 
     # Graph Query & Reasoning
-    def get_node(self, id: str) -> dict[str, Any] | None:
+    def get_node(self, id: str, include_embedding: bool = True) -> dict[str, Any] | None:
         """
         Retrieve the metadata and memory of a node.
         Args:
             id: Node identifier.
+            include_embedding (bool): Whether to include the large embedding field.
         Returns:
             Dictionary of node fields, or None if not found.
         """
+
         where_user = ""
         params = {"id": id}
         if not self.config.use_multi_db and self.config.user_name:
@@ -343,11 +345,12 @@ class Neo4jGraphDB(BaseGraphDB):
             record = session.run(query, params).single()
             return self._parse_node(dict(record["n"])) if record else None
 
-    def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]:
+    def get_nodes(self, ids: list[str], include_embedding: bool = True) -> list[dict[str, Any]]:
         """
         Retrieve the metadata and memory of a list of nodes.
         Args:
             ids: List of Node identifier.
+            include_embedding (bool): Whether to include the large embedding field.
         Returns:
         list[dict]: Parsed node records containing 'id', 'memory', and 'metadata'.
 
@@ -355,6 +358,7 @@ class Neo4jGraphDB(BaseGraphDB):
             - Assumes all provided IDs are valid and exist.
             - Returns empty list if input is empty.
         """
+
         if not ids:
             return []
 
@@ -829,7 +833,7 @@ class Neo4jGraphDB(BaseGraphDB):
             logger.error(f"[ERROR] Failed to clear database '{self.db_name}': {e}")
             raise
 
-    def export_graph(self) -> dict[str, Any]:
+    def export_graph(self, include_embedding: bool = True) -> dict[str, Any]:
         """
         Export all graph nodes and edges in a structured form.
 
@@ -910,17 +914,19 @@ class Neo4jGraphDB(BaseGraphDB):
                     target_id=edge["target"],
                 )
 
-    def get_all_memory_items(self, scope: str) -> list[dict]:
+    def get_all_memory_items(self, scope: str, include_embedding: bool = True) -> list[dict]:
         """
         Retrieve all memory items of a specific memory_type.
 
         Args:
             scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'.
+            include_embedding (bool): Whether to include the large embedding field.
+        Returns:
 
         Returns:
             list[dict]: Full list of memory items under this scope.
         """
-        if scope not in {"WorkingMemory", "LongTermMemory", "UserMemory"}:
+        if scope not in {"WorkingMemory", "LongTermMemory", "UserMemory", "OuterMemory"}:
             raise ValueError(f"Unsupported memory type scope: {scope}")
 
         where_clause = "WHERE n.memory_type = $scope"
@@ -940,12 +946,15 @@ class Neo4jGraphDB(BaseGraphDB):
             results = session.run(query, params)
             return [self._parse_node(dict(record["n"])) for record in results]
 
-    def get_structure_optimization_candidates(self, scope: str) -> list[dict]:
+    def get_structure_optimization_candidates(
+        self, scope: str, include_embedding: bool = True
+    ) -> list[dict]:
         """
         Find nodes that are likely candidates for structure optimization:
         - Isolated nodes, nodes with empty background, or nodes with exactly one child.
         - Plus: the child of any parent node that has exactly one child.
         """
+
         where_clause = """
                 WHERE n.memory_type = $scope
                   AND n.status = 'activated'

memos/log.py

@@ -4,9 +4,14 @@ from logging.config import dictConfig
 from pathlib import Path
 from sys import stdout
 
+from dotenv import load_dotenv
+
 from memos import settings
 
 
+# Load environment variables
+load_dotenv()
+
 selected_log_level = logging.DEBUG if settings.DEBUG else logging.WARNING
 
 
@@ -48,7 +53,7 @@ LOGGING_CONFIG = {
             "class": "logging.handlers.RotatingFileHandler",
             "filename": _setup_logfile(),
             "maxBytes": 1024**2 * 10,
-            "backupCount": 3,
+            "backupCount": 10,
             "formatter": "standard",
         },
     },

memos/mem_cube/utils.py

@@ -71,10 +71,6 @@ def merge_config_with_default(
 
             # Define graph_db fields to preserve (user-specific)
             preserve_graph_fields = {
-                "uri",
-                "user",
-                "password",
-                "db_name",
                 "auto_create",
                 "user_name",
                 "use_multi_db",
@@ -96,9 +92,20 @@ def merge_config_with_default(
                     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": existing_text_config["graph_db"]["backend"],
+                "backend": default_text_config["graph_db"]["backend"],
                 "config": merged_graph_config,
             }
 

memos/mem_os/core.py

@@ -1,6 +1,6 @@
 import json
 import os
-import uuid
+import time
 
 from datetime import datetime
 from pathlib import Path
@@ -13,16 +13,19 @@ 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 (
+from memos.mem_scheduler.scheduler_factory import SchedulerFactory
+from memos.mem_scheduler.schemas.general_schemas import (
     ADD_LABEL,
     ANSWER_LABEL,
-    ScheduleMessageItem,
+    QUERY_LABEL,
 )
-from memos.mem_scheduler.scheduler_factory import SchedulerFactory
+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.memos_tools.thread_safe_dict import ThreadSafeDict
+from memos.templates.mos_prompts import QUERY_REWRITING_PROMPT
 from memos.types import ChatHistory, MessageList, MOSSearchResult
 
 
@@ -40,10 +43,13 @@ class MOSCore:
         self.config = config
         self.user_id = config.user_id
         self.session_id = config.session_id
-        self.mem_cubes: dict[str, GeneralMemCube] = {}
         self.chat_llm = LLMFactory.from_config(config.chat_model)
         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._register_chat_history()
 
         # Use provided user_manager or create a new one
@@ -58,10 +64,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: GeneralScheduler = 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
@@ -93,14 +104,16 @@ class MOSCore:
             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
@@ -111,13 +124,16 @@ class MOSCore:
                     f"Memory reader of type {type(self.mem_reader).__name__} "
                     "missing required 'llm' attribute"
                 )
-                self._mem_scheduler.initialize_modules(chat_llm=self.chat_llm)
+                self._mem_scheduler.initialize_modules(
+                    chat_llm=self.chat_llm, process_llm=self.chat_llm
+                )
             else:
-                # Configure scheduler modules
+                # Configure scheduler general_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:
@@ -157,6 +173,14 @@ class MOSCore:
             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."""
@@ -165,7 +189,7 @@ class MOSCore:
         self.chat_history_manager[user_id] = ChatHistory(
             user_id=user_id,
             session_id=self.session_id,
-            created_at=datetime.now(),
+            created_at=datetime.utcnow(),
             total_messages=0,
             chat_history=[],
         )
@@ -257,13 +281,21 @@ class MOSCore:
                         user_id=target_user_id,
                         mem_cube_id=mem_cube_id,
                         mem_cube=mem_cube,
-                        label=ADD_LABEL,
+                        label=QUERY_LABEL,
                         content=query,
-                        timestamp=datetime.now(),
+                        timestamp=datetime.utcnow(),
                     )
                     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, base_prompt=base_prompt)
@@ -310,7 +342,7 @@ class MOSCore:
                     mem_cube=mem_cube,
                     label=ANSWER_LABEL,
                     content=response,
-                    timestamp=datetime.now(),
+                    timestamp=datetime.utcnow(),
                 )
                 self.mem_scheduler.submit_messages(messages=[message_item])
 
@@ -511,6 +543,8 @@ class MOSCore:
         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.
@@ -534,6 +568,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": [],
@@ -541,19 +579,38 @@ class MOSCore:
         }
         if install_cube_ids is None:
             install_cube_ids = user_cube_ids
-        for mem_cube_id, mem_cube in self.mem_cubes.items():
+        # create exist dict in mem_cubes and avoid  one search slow
+        tmp_mem_cubes = {}
+        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)
+
+        for mem_cube_id, mem_cube in tmp_mem_cubes.items():
             if (
                 (mem_cube_id in install_cube_ids)
                 and (mem_cube.text_mem is not None)
                 and self.config.enable_textual_memory
             ):
+                time_start = time.time()
                 memories = mem_cube.text_mem.search(
-                    query, top_k=top_k if top_k else self.config.top_k
+                    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"
                 )
+                search_time_end = time.time()
+                logger.info(
+                    f"time search graph: search graph time user_id: {target_user_id} time is: {search_time_end - time_start}"
+                )
         return result
 
     def add(
@@ -576,6 +633,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."
         )
@@ -613,25 +671,31 @@ class MOSCore:
                 memories = self.mem_reader.get_memory(
                     messages_list,
                     type="chat",
-                    info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
+                    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
-                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]
+                    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(text_messages),
-                        timestamp=datetime.now(),
+                        content=json.dumps(mem_ids),
+                        timestamp=datetime.utcnow(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 
+        # user profile
         if (
             (memory_content is not None)
             and self.config.enable_textual_memory
@@ -646,34 +710,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": str(uuid.uuid4())},
+                    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.utcnow(),
+                    )
+                    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": str(uuid.uuid4())},
+                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.utcnow(),
+                )
+                self.mem_scheduler.submit_messages(messages=[message_item])
+
         logger.info(f"Add memory to {mem_cube_id} successfully")
 
     def get(
@@ -907,3 +1003,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

memos/mem_os/main.py

@@ -208,7 +208,7 @@ class MOS(MOSCore):
                 if self.enable_mem_scheduler and self.mem_scheduler is not None:
                     from datetime import datetime
 
-                    from memos.mem_scheduler.modules.schemas import (
+                    from memos.mem_scheduler.schemas import (
                         ANSWER_LABEL,
                         ScheduleMessageItem,
                     )
@@ -219,7 +219,7 @@ class MOS(MOSCore):
                         mem_cube=mem_cube,
                         label=ANSWER_LABEL,
                         content=enhanced_response,
-                        timestamp=datetime.now(),
+                        timestamp=datetime.now().isoformat(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])