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

memos/graph_dbs/neo4j.py

@@ -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,12 +914,14 @@ 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.
@@ -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
 
 

memos/mem_os/core.py

@@ -24,6 +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.templates.mos_prompts import QUERY_REWRITING_PROMPT
 from memos.types import ChatHistory, MessageList, MOSSearchResult
 
@@ -42,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
@@ -124,7 +128,7 @@ class MOSCore:
                     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
                 )
@@ -185,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=[],
         )
@@ -279,7 +283,7 @@ class MOSCore:
                         mem_cube=mem_cube,
                         label=QUERY_LABEL,
                         content=query,
-                        timestamp=datetime.now(),
+                        timestamp=datetime.utcnow(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 
@@ -338,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])
 
@@ -575,7 +579,13 @@ 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)
@@ -681,7 +691,7 @@ class MOSCore:
                         mem_cube=mem_cube,
                         label=ADD_LABEL,
                         content=json.dumps(mem_ids),
-                        timestamp=datetime.now(),
+                        timestamp=datetime.utcnow(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 
@@ -725,7 +735,7 @@ class MOSCore:
                         mem_cube=mem_cube,
                         label=ADD_LABEL,
                         content=json.dumps(mem_ids),
-                        timestamp=datetime.now(),
+                        timestamp=datetime.utcnow(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 
@@ -756,7 +766,7 @@ class MOSCore:
                     mem_cube=mem_cube,
                     label=ADD_LABEL,
                     content=json.dumps(mem_ids),
-                    timestamp=datetime.now(),
+                    timestamp=datetime.utcnow(),
                 )
                 self.mem_scheduler.submit_messages(messages=[message_item])
 

memos/mem_os/main.py

@@ -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])
 

memos/mem_os/product.py

@@ -22,7 +22,9 @@ from memos.mem_os.utils.format_utils import (
     filter_nodes_by_tree_ids,
     remove_embedding_recursive,
     sort_children_by_memory_type,
-    split_continuous_references,
+)
+from memos.mem_os.utils.reference_utils import (
+    process_streaming_references_complete,
 )
 from memos.mem_scheduler.schemas.general_schemas import (
     ANSWER_LABEL,
@@ -406,71 +408,6 @@ class MOSProduct(MOSCore):
             return MEMOS_PRODUCT_ENHANCE_PROMPT + personal_memory_context + outer_memory_context
         return MEMOS_PRODUCT_ENHANCE_PROMPT
 
-    def _process_streaming_references_complete(self, text_buffer: str) -> tuple[str, str]:
-        """
-        Complete streaming reference processing to ensure reference tags are never split.
-
-        Args:
-            text_buffer (str): The accumulated text buffer.
-
-        Returns:
-            tuple[str, str]: (processed_text, remaining_buffer)
-        """
-        import re
-
-        # Pattern to match complete reference tags: [refid:memoriesID]
-        complete_pattern = r"\[\d+:[^\]]+\]"
-
-        # Find all complete reference tags
-        complete_matches = list(re.finditer(complete_pattern, text_buffer))
-
-        if complete_matches:
-            # Find the last complete tag
-            last_match = complete_matches[-1]
-            end_pos = last_match.end()
-
-            # Get text up to the end of the last complete tag
-            processed_text = text_buffer[:end_pos]
-            remaining_buffer = text_buffer[end_pos:]
-
-            # Apply reference splitting to the processed text
-            processed_text = split_continuous_references(processed_text)
-
-            return processed_text, remaining_buffer
-
-        # Check for incomplete reference tags
-        # Look for opening bracket with number and colon
-        opening_pattern = r"\[\d+:"
-        opening_matches = list(re.finditer(opening_pattern, text_buffer))
-
-        if opening_matches:
-            # Find the last opening tag
-            last_opening = opening_matches[-1]
-            opening_start = last_opening.start()
-
-            # Check if we have a complete opening pattern
-            if last_opening.end() <= len(text_buffer):
-                # We have a complete opening pattern, keep everything in buffer
-                return "", text_buffer
-            else:
-                # Incomplete opening pattern, return text before it
-                processed_text = text_buffer[:opening_start]
-                # Apply reference splitting to the processed text
-                processed_text = split_continuous_references(processed_text)
-                return processed_text, text_buffer[opening_start:]
-
-        # Check for partial opening pattern (starts with [ but not complete)
-        if "[" in text_buffer:
-            ref_start = text_buffer.find("[")
-            processed_text = text_buffer[:ref_start]
-            # Apply reference splitting to the processed text
-            processed_text = split_continuous_references(processed_text)
-            return processed_text, text_buffer[ref_start:]
-
-        # No reference tags found, apply reference splitting and return all text
-        processed_text = split_continuous_references(text_buffer)
-        return processed_text, ""
-
     def _extract_references_from_response(self, response: str) -> tuple[str, list[dict]]:
         """
         Extract reference information from the response and return clean text.
@@ -554,7 +491,7 @@ class MOSProduct(MOSCore):
                 mem_cube=self.mem_cubes[mem_cube_id],
                 label=label,
                 content=query,
-                timestamp=datetime.now(),
+                timestamp=datetime.utcnow(),
             )
             self.mem_scheduler.submit_messages(messages=[message_item])
 
@@ -868,7 +805,7 @@ class MOSProduct(MOSCore):
             full_response += chunk
 
             # Process buffer to ensure complete reference tags
-            processed_chunk, remaining_buffer = self._process_streaming_references_complete(buffer)
+            processed_chunk, remaining_buffer = process_streaming_references_complete(buffer)
 
             if processed_chunk:
                 chunk_data = f"data: {json.dumps({'type': 'text', 'data': processed_chunk}, ensure_ascii=False)}\n\n"
@@ -877,7 +814,7 @@ class MOSProduct(MOSCore):
 
         # Process any remaining buffer
         if buffer:
-            processed_chunk, remaining_buffer = self._process_streaming_references_complete(buffer)
+            processed_chunk, remaining_buffer = process_streaming_references_complete(buffer)
             if processed_chunk:
                 chunk_data = f"data: {json.dumps({'type': 'text', 'data': processed_chunk}, ensure_ascii=False)}\n\n"
                 yield chunk_data

memos/mem_os/utils/default_config.py

@@ -112,7 +112,7 @@ def get_default_config(
                 "thread_pool_max_workers": kwargs.get("scheduler_thread_pool_max_workers", 10),
                 "consume_interval_seconds": kwargs.get("scheduler_consume_interval_seconds", 3),
                 "enable_parallel_dispatch": kwargs.get("scheduler_enable_parallel_dispatch", True),
-                "enable_act_memory_update": True,
+                "enable_activation_memory": True,
             },
         }
 

memos/mem_os/utils/format_utils.py

@@ -570,15 +570,23 @@ def convert_graph_to_tree_forworkmem(
         else:
             other_roots.append(root_id)
 
-    def build_tree(node_id: str) -> dict[str, Any]:
-        """Recursively build tree structure"""
+    def build_tree(node_id: str, visited=None) -> dict[str, Any] | None:
+        """Recursively build tree structure with cycle detection"""
+        if visited is None:
+            visited = set()
+
+        if node_id in visited:
+            logger.warning(f"[build_tree] Detected cycle at node {node_id}, skipping.")
+            return None
+        visited.add(node_id)
+
         if node_id not in node_map:
             return None
 
         children_ids = children_map.get(node_id, [])
         children = []
         for child_id in children_ids:
-            child_tree = build_tree(child_id)
+            child_tree = build_tree(child_id, visited)
             if child_tree:
                 children.append(child_tree)
 
@@ -1355,47 +1363,3 @@ def clean_json_response(response: str) -> str:
         str: Clean JSON string without markdown formatting
     """
     return response.replace("```json", "").replace("```", "").strip()
-
-
-def split_continuous_references(text: str) -> str:
-    """
-    Split continuous reference tags into individual reference tags.
-
-    Converts patterns like [1:92ff35fb, 4:bfe6f044] to [1:92ff35fb] [4:bfe6f044]
-
-    Only processes text if:
-    1. '[' appears exactly once
-    2. ']' appears exactly once
-    3. Contains commas between '[' and ']'
-
-    Args:
-        text (str): Text containing reference tags
-
-    Returns:
-        str: Text with split reference tags, or original text if conditions not met
-    """
-    # Early return if text is empty
-    if not text:
-        return text
-    # Check if '[' appears exactly once
-    if text.count("[") != 1:
-        return text
-    # Check if ']' appears exactly once
-    if text.count("]") != 1:
-        return text
-    # Find positions of brackets
-    open_bracket_pos = text.find("[")
-    close_bracket_pos = text.find("]")
-
-    # Check if brackets are in correct order
-    if open_bracket_pos >= close_bracket_pos:
-        return text
-    # Extract content between brackets
-    content_between_brackets = text[open_bracket_pos + 1 : close_bracket_pos]
-    # Check if there's a comma between brackets
-    if "," not in content_between_brackets:
-        return text
-    text = text.replace(content_between_brackets, content_between_brackets.replace(", ", "]["))
-    text = text.replace(content_between_brackets, content_between_brackets.replace(",", "]["))
-
-    return text

memos/mem_os/utils/reference_utils.py

@@ -0,0 +1,133 @@
+def split_continuous_references(text: str) -> str:
+    """
+    Split continuous reference tags into individual reference tags.
+
+    Converts patterns like [1:92ff35fb, 4:bfe6f044] to [1:92ff35fb] [4:bfe6f044]
+
+    Only processes text if:
+    1. '[' appears exactly once
+    2. ']' appears exactly once
+    3. Contains commas between '[' and ']'
+
+    Args:
+        text (str): Text containing reference tags
+
+    Returns:
+        str: Text with split reference tags, or original text if conditions not met
+    """
+    # Early return if text is empty
+    if not text:
+        return text
+    # Check if '[' appears exactly once
+    if text.count("[") != 1:
+        return text
+    # Check if ']' appears exactly once
+    if text.count("]") != 1:
+        return text
+    # Find positions of brackets
+    open_bracket_pos = text.find("[")
+    close_bracket_pos = text.find("]")
+
+    # Check if brackets are in correct order
+    if open_bracket_pos >= close_bracket_pos:
+        return text
+    # Extract content between brackets
+    content_between_brackets = text[open_bracket_pos + 1 : close_bracket_pos]
+    # Check if there's a comma between brackets
+    if "," not in content_between_brackets:
+        return text
+    text = text.replace(content_between_brackets, content_between_brackets.replace(", ", "]["))
+    text = text.replace(content_between_brackets, content_between_brackets.replace(",", "]["))
+
+    return text
+
+
+def process_streaming_references_complete(text_buffer: str) -> tuple[str, str]:
+    """
+    Complete streaming reference processing to ensure reference tags are never split.
+
+    Args:
+        text_buffer (str): The accumulated text buffer.
+
+    Returns:
+        tuple[str, str]: (processed_text, remaining_buffer)
+    """
+    import re
+
+    # Pattern to match complete reference tags: [refid:memoriesID]
+    complete_pattern = r"\[\d+:[^\]]+\]"
+
+    # Find all complete reference tags
+    complete_matches = list(re.finditer(complete_pattern, text_buffer))
+
+    if complete_matches:
+        # Find the last complete tag
+        last_match = complete_matches[-1]
+        end_pos = last_match.end()
+
+        # Check if there's any incomplete reference after the last complete one
+        remaining_text = text_buffer[end_pos:]
+
+        # Look for potential incomplete reference patterns after the last complete tag
+        incomplete_pattern = r"\[\d*:?[^\]]*$"
+        if re.search(incomplete_pattern, remaining_text):
+            # There's a potential incomplete reference, find where it starts
+            incomplete_match = re.search(incomplete_pattern, remaining_text)
+            if incomplete_match:
+                incomplete_start = end_pos + incomplete_match.start()
+                processed_text = text_buffer[:incomplete_start]
+                remaining_buffer = text_buffer[incomplete_start:]
+
+                # Apply reference splitting to the processed text
+                processed_text = split_continuous_references(processed_text)
+                return processed_text, remaining_buffer
+
+        # No incomplete reference after the last complete tag, process all
+        processed_text = split_continuous_references(text_buffer)
+        return processed_text, ""
+
+    # Check for incomplete reference tags - be more specific about what constitutes a potential reference
+    # Look for opening bracket with number and colon that could be a reference tag
+    opening_pattern = r"\[\d+:"
+    opening_matches = list(re.finditer(opening_pattern, text_buffer))
+
+    if opening_matches:
+        # Find the last opening tag
+        last_opening = opening_matches[-1]
+        opening_start = last_opening.start()
+
+        # Check if this might be a complete reference tag (has closing bracket after the pattern)
+        remaining_text = text_buffer[last_opening.end() :]
+        if "]" in remaining_text:
+            # This looks like a complete reference tag, process it
+            processed_text = split_continuous_references(text_buffer)
+            return processed_text, ""
+        else:
+            # Incomplete reference tag, keep it in buffer
+            processed_text = text_buffer[:opening_start]
+            processed_text = split_continuous_references(processed_text)
+            return processed_text, text_buffer[opening_start:]
+
+    # More sophisticated check for potential reference patterns
+    # Only hold back text if we see a pattern that could be the start of a reference tag
+    potential_ref_pattern = r"\[\d*:?$"  # Matches [, [1, [12:, etc. at end of buffer
+    if re.search(potential_ref_pattern, text_buffer):
+        # Find the position of the potential reference start
+        match = re.search(potential_ref_pattern, text_buffer)
+        if match:
+            ref_start = match.start()
+            processed_text = text_buffer[:ref_start]
+            processed_text = split_continuous_references(processed_text)
+            return processed_text, text_buffer[ref_start:]
+
+    # Check for standalone [ only at the very end of the buffer
+    # This prevents cutting off mathematical expressions like [ \Delta U = Q - W ]
+    if text_buffer.endswith("["):
+        # Only hold back the single [ character
+        processed_text = text_buffer[:-1]
+        processed_text = split_continuous_references(processed_text)
+        return processed_text, "["
+
+    # No reference-like patterns found, process all text
+    processed_text = split_continuous_references(text_buffer)
+    return processed_text, ""