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

memos/memories/textual/tree_text_memory/organize/reorganizer.py

@@ -4,19 +4,20 @@ import time
 import traceback
 
 from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed
 from queue import PriorityQueue
 from typing import Literal
 
 import numpy as np
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.dependency import require_python_package
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.item import GraphDBEdge, GraphDBNode
 from memos.graph_dbs.neo4j import Neo4jGraphDB
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
-from memos.memories.textual.item import TreeNodeTextualMemoryMetadata
+from memos.memories.textual.item import SourceMessage, TreeNodeTextualMemoryMetadata
 from memos.memories.textual.tree_text_memory.organize.handler import NodeHandler
 from memos.memories.textual.tree_text_memory.organize.relation_reason_detector import (
     RelationAndReasoningDetector,
@@ -27,6 +28,22 @@ from memos.templates.tree_reorganize_prompts import LOCAL_SUBCLUSTER_PROMPT, REO
 logger = get_logger(__name__)
 
 
+def build_summary_parent_node(cluster_nodes):
+    normalized_sources = []
+    for n in cluster_nodes:
+        sm = SourceMessage(
+            type="chat",
+            role=None,
+            chat_time=None,
+            message_id=None,
+            content=n.memory,
+            # extra
+            node_id=n.id,
+        )
+        normalized_sources.append(sm)
+    return normalized_sources
+
+
 class QueueMessage:
     def __init__(
         self,
@@ -51,6 +68,15 @@ class QueueMessage:
         return op_priority[self.op] < op_priority[other.op]
 
 
+def extract_first_to_last_brace(text: str):
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        return "", None
+    json_str = text[start : end + 1]
+    return json_str, json.loads(json_str)
+
+
 class GraphStructureReorganizer:
     def __init__(
         self, graph_store: Neo4jGraphDB, llm: BaseLLM, embedder: OllamaEmbedder, is_reorganize: bool
@@ -87,6 +113,7 @@ class GraphStructureReorganizer:
         1) queue is empty
         2) any running structure optimization is done
         """
+        deadline = time.time() + 600
         if not self.is_reorganize:
             return
 
@@ -96,6 +123,9 @@ class GraphStructureReorganizer:
 
         while any(self._is_optimizing.values()):
             logger.debug(f"Waiting for structure optimizer to finish... {self._is_optimizing}")
+            if time.time() > deadline:
+                logger.error(f"Wait timed out; flags={self._is_optimizing}")
+                break
             time.sleep(1)
         logger.debug("Structure optimizer is now idle.")
 
@@ -129,6 +159,9 @@ class GraphStructureReorganizer:
 
         logger.info("Structure optimizer schedule started.")
         while not getattr(self, "_stop_scheduler", False):
+            if any(self._is_optimizing.values()):
+                time.sleep(1)
+                continue
             if self._reorganize_needed:
                 logger.info("[Reorganizer] Triggering optimize_structure due to new nodes.")
                 self.optimize_structure(scope="LongTermMemory")
@@ -176,6 +209,7 @@ class GraphStructureReorganizer:
         local_tree_threshold: int = 10,
         min_cluster_size: int = 4,
         min_group_size: int = 20,
+        max_duration_sec: int = 600,
     ):
         """
         Periodically reorganize the graph:
@@ -183,8 +217,20 @@ class GraphStructureReorganizer:
         2. Summarize each cluster.
         3. Create parent nodes and build local PARENT trees.
         """
+        # --- Total time watch dog: check functions ---
+        start_ts = time.time()
+
+        def _check_deadline(where: str):
+            if time.time() - start_ts > max_duration_sec:
+                logger.error(
+                    f"[GraphStructureReorganize] {scope} surpass {max_duration_sec}s，time "
+                    f"over at {where}"
+                )
+                return True
+            return False
+
         if self._is_optimizing[scope]:
-            logger.info(f"Already optimizing for {scope}. Skipping.")
+            logger.info(f"[GraphStructureReorganize] Already optimizing for {scope}. Skipping.")
             return
 
         if self.graph_store.node_not_exist(scope):
@@ -198,32 +244,35 @@ class GraphStructureReorganizer:
             )
 
             logger.debug(
-                f"Num of scope in self.graph_store is {self.graph_store.get_memory_count(scope)}"
+                f"[GraphStructureReorganize] Num of scope in self.graph_store is"
+                f" {self.graph_store.get_memory_count(scope)}"
             )
             # Load candidate nodes
+            if _check_deadline("[GraphStructureReorganize] Before loading candidates"):
+                return
             raw_nodes = self.graph_store.get_structure_optimization_candidates(scope)
             nodes = [GraphDBNode(**n) for n in raw_nodes]
 
             if not nodes:
                 logger.info("[GraphStructureReorganize] No nodes to optimize. Skipping.")
                 return
-
             if len(nodes) < min_group_size:
                 logger.info(
                     f"[GraphStructureReorganize] Only {len(nodes)} candidate nodes found. Not enough to reorganize. Skipping."
                 )
                 return
 
-            logger.info(f"[GraphStructureReorganize] Loaded {len(nodes)} nodes.")
-
             # Step 2: Partition nodes
+            if _check_deadline("[GraphStructureReorganize] Before partition"):
+                return
             partitioned_groups = self._partition(nodes)
-
             logger.info(
                 f"[GraphStructureReorganize] Partitioned into {len(partitioned_groups)} clusters."
             )
 
-            with ThreadPoolExecutor(max_workers=4) as executor:
+            if _check_deadline("[GraphStructureReorganize] Before submit partition task"):
+                return
+            with ContextThreadPoolExecutor(max_workers=4) as executor:
                 futures = []
                 for cluster_nodes in partitioned_groups:
                     futures.append(
@@ -237,14 +286,17 @@ class GraphStructureReorganizer:
                     )
 
                 for f in as_completed(futures):
+                    if _check_deadline("[GraphStructureReorganize] Waiting clusters..."):
+                        for x in futures:
+                            x.cancel()
+                        return
                     try:
                         f.result()
                     except Exception as e:
                         logger.warning(
-                            f"[Reorganize] Cluster processing "
-                            f"failed: {e}, cluster_nodes: {cluster_nodes}, trace: {traceback.format_exc()}"
+                            f"[GraphStructureReorganize] Cluster processing failed: {e}, trace: {traceback.format_exc()}"
                         )
-                logger.info("[GraphStructure Reorganize] Structure optimization finished.")
+            logger.info("[GraphStructure Reorganize] Structure optimization finished.")
 
         finally:
             self._is_optimizing[scope] = False
@@ -282,7 +334,7 @@ class GraphStructureReorganizer:
         nodes_to_check = cluster_nodes
         exclude_ids = [n.id for n in nodes_to_check]
 
-        with ThreadPoolExecutor(max_workers=4) as executor:
+        with ContextThreadPoolExecutor(max_workers=4) as executor:
             futures = []
             for node in nodes_to_check:
                 futures.append(
@@ -294,7 +346,7 @@ class GraphStructureReorganizer:
                     )
                 )
 
-            for f in as_completed(futures):
+            for f in as_completed(futures, timeout=300):
                 results = f.result()
 
                 # 1) Add pairwise relations
@@ -331,11 +383,11 @@ class GraphStructureReorganizer:
                     for child_id in agg_node.metadata.sources:
                         self.graph_store.add_edge(agg_node.id, child_id, "AGGREGATE_TO")
 
-            logger.info("[Reorganizer] Cluster relation/reasoning done.")
+        logger.info("[Reorganizer] Cluster relation/reasoning done.")
 
     def _local_subcluster(
-        self, cluster_nodes: list[GraphDBNode], max_length: int = 8000
-    ) -> (list)[list[GraphDBNode]]:
+        self, cluster_nodes: list[GraphDBNode], max_length: int = 15000
+    ) -> list[list[GraphDBNode]]:
         """
         Use LLM to split a large cluster into semantically coherent sub-clusters.
         """
@@ -350,7 +402,7 @@ class GraphStructureReorganizer:
 
         joined_scene = "\n".join(scene_lines)
         if len(joined_scene) > max_length:
-            logger.warning(f"Sub-cluster too long: {joined_scene}")
+            logger.warning("Sub-cluster too long")
         prompt = LOCAL_SUBCLUSTER_PROMPT.replace("{joined_scene}", joined_scene[:max_length])
 
         messages = [{"role": "user", "content": prompt}]
@@ -499,17 +551,17 @@ class GraphStructureReorganizer:
         parent_node = GraphDBNode(
             memory=parent_value,
             metadata=TreeNodeTextualMemoryMetadata(
-                user_id="",  # TODO: summarized node: no user_id
-                session_id="",  # TODO: summarized node: no session_id
+                user_id=None,
+                session_id=None,
                 memory_type=scope,
                 status="activated",
                 key=parent_key,
                 tags=parent_tags,
                 embedding=embedding,
                 usage=[],
-                sources=[n.id for n in cluster_nodes],
+                sources=build_summary_parent_node(cluster_nodes),
                 background=parent_background,
-                confidence=0.99,
+                confidence=0.66,
                 type="topic",
             ),
         )
@@ -518,7 +570,7 @@ class GraphStructureReorganizer:
     def _parse_json_result(self, response_text):
         try:
             response_text = response_text.replace("```", "").replace("json", "")
-            response_json = json.loads(response_text)
+            response_json = extract_first_to_last_brace(response_text)[1]
             return response_json
         except json.JSONDecodeError as e:
             logger.warning(

memos/memories/textual/tree_text_memory/retrieve/bochasearch.py

@@ -2,15 +2,17 @@
 
 import json
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed
 from datetime import datetime
+from typing import Any
 
 import requests
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import OllamaEmbedder
 from memos.log import get_logger
 from memos.mem_reader.base import BaseMemReader
-from memos.memories.textual.item import TextualMemoryItem
+from memos.memories.textual.item import SourceMessage, TextualMemoryItem
 
 
 logger = get_logger(__name__)
@@ -177,7 +179,7 @@ class BochaAISearchRetriever:
         if not info:
             info = {"user_id": "", "session_id": ""}
 
-        with ThreadPoolExecutor(max_workers=8) as executor:
+        with ContextThreadPoolExecutor(max_workers=8) as executor:
             futures = [
                 executor.submit(self._process_result, r, query, parsed_goal, info)
                 for r in search_results
@@ -193,7 +195,7 @@ class BochaAISearchRetriever:
         return list(unique_memory_items.values())
 
     def _process_result(
-        self, result: dict, query: str, parsed_goal: str, info: None
+        self, result: dict, query: str, parsed_goal: str, info: dict[str, Any]
     ) -> list[TextualMemoryItem]:
         """Process one Bocha search result into TextualMemoryItem."""
         title = result.get("name", "")
@@ -218,12 +220,14 @@ class BochaAISearchRetriever:
         memory_items = []
         for read_item_i in read_items[0]:
             read_item_i.memory = (
-                f"Title: {title}\nNewsTime: {publish_time}\nSummary: {summary}\n"
+                f"[Outer internet view] Title: {title}\nNewsTime:"
+                f" {publish_time}\nSummary:"
+                f" {summary}\n"
                 f"Content: {read_item_i.memory}"
             )
             read_item_i.metadata.source = "web"
             read_item_i.metadata.memory_type = "OuterMemory"
-            read_item_i.metadata.sources = [url] if url else []
+            read_item_i.metadata.sources = [SourceMessage(type="web", url=url)] if url else []
             read_item_i.metadata.visibility = "public"
             memory_items.append(read_item_i)
         return memory_items

memos/memories/textual/tree_text_memory/retrieve/internet_retriever.py

@@ -7,7 +7,11 @@ from datetime import datetime
 import requests
 
 from memos.embedders.factory import OllamaEmbedder
-from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
+from memos.memories.textual.item import (
+    SourceMessage,
+    TextualMemoryItem,
+    TreeNodeTextualMemoryMetadata,
+)
 
 
 class GoogleCustomSearchAPI:
@@ -172,7 +176,7 @@ class InternetGoogleRetriever:
                 visibility="public",
                 memory_type="LongTermMemory",  # Internet search results as working memory
                 key=title,
-                sources=[link] if link else [],
+                sources=[SourceMessage(type="web", url=link)] if link else [],
                 embedding=self.embedder.embed([memory_content])[0],  # Can add embedding later
                 created_at=datetime.now().isoformat(),
                 usage=[],

memos/memories/textual/tree_text_memory/retrieve/internet_retriever_factory.py

@@ -10,6 +10,7 @@ from memos.memories.textual.tree_text_memory.retrieve.internet_retriever import
     InternetGoogleRetriever,
 )
 from memos.memories.textual.tree_text_memory.retrieve.xinyusearch import XinyuSearchRetriever
+from memos.memos_tools.singleton import singleton_factory
 
 
 class InternetRetrieverFactory:
@@ -23,6 +24,7 @@ class InternetRetrieverFactory:
     }
 
     @classmethod
+    @singleton_factory()
     def from_config(
         cls, config_factory: InternetRetrieverConfigFactory, embedder: BaseEmbedder
     ) -> InternetGoogleRetriever | None:

memos/memories/textual/tree_text_memory/retrieve/recall.py

@@ -1,11 +1,16 @@
 import concurrent.futures
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.neo4j import Neo4jGraphDB
+from memos.log import get_logger
 from memos.memories.textual.item import TextualMemoryItem
 from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
 
 
+logger = get_logger(__name__)
+
+
 class GraphMemoryRetriever:
     """
     Unified memory retriever that combines both graph-based and vector-based retrieval logic.
@@ -14,6 +19,8 @@ class GraphMemoryRetriever:
     def __init__(self, graph_store: Neo4jGraphDB, embedder: OllamaEmbedder):
         self.graph_store = graph_store
         self.embedder = embedder
+        self.max_workers = 10
+        self.filter_weight = 0.6
 
     def retrieve(
         self,
@@ -22,6 +29,7 @@ class GraphMemoryRetriever:
         top_k: int,
         memory_scope: str,
         query_embedding: list[list[float]] | None = None,
+        search_filter: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """
         Perform hybrid memory retrieval:
@@ -35,7 +43,7 @@ class GraphMemoryRetriever:
             top_k (int): Number of candidates to return.
             memory_scope (str): One of ['working', 'long_term', 'user'].
             query_embedding(list of embedding): list of embedding of query
-
+            search_filter (dict, optional): Optional metadata filters for search results.
         Returns:
             list: Combined memory items.
         """
@@ -45,16 +53,20 @@ class GraphMemoryRetriever:
         if memory_scope == "WorkingMemory":
             # For working memory, retrieve all entries (no filtering)
             working_memories = self.graph_store.get_all_memory_items(
-                scope="WorkingMemory", include_embedding=True
+                scope="WorkingMemory", include_embedding=False
             )
             return [TextualMemoryItem.from_dict(record) for record in working_memories]
 
-        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        with ContextThreadPoolExecutor(max_workers=2) as executor:
             # Structured graph-based retrieval
             future_graph = executor.submit(self._graph_recall, parsed_goal, memory_scope)
             # Vector similarity search
             future_vector = executor.submit(
-                self._vector_recall, query_embedding, memory_scope, top_k
+                self._vector_recall,
+                query_embedding or [],
+                memory_scope,
+                top_k,
+                search_filter=search_filter,
             )
 
             graph_results = future_graph.result()
@@ -74,6 +86,51 @@ class GraphMemoryRetriever:
 
         return list(combined.values())
 
+    def retrieve_from_cube(
+        self,
+        top_k: int,
+        memory_scope: str,
+        query_embedding: list[list[float]] | None = None,
+        cube_name: str = "memos_cube01",
+    ) -> list[TextualMemoryItem]:
+        """
+        Perform hybrid memory retrieval:
+        - Run graph-based lookup from dispatch plan.
+        - Run vector similarity search from embedded query.
+        - Merge and return combined result set.
+
+        Args:
+            top_k (int): Number of candidates to return.
+            memory_scope (str): One of ['working', 'long_term', 'user'].
+            query_embedding(list of embedding): list of embedding of query
+            cube_name: specify cube_name
+
+        Returns:
+            list: Combined memory items.
+        """
+        if memory_scope not in ["WorkingMemory", "LongTermMemory", "UserMemory"]:
+            raise ValueError(f"Unsupported memory scope: {memory_scope}")
+
+        graph_results = self._vector_recall(
+            query_embedding, memory_scope, top_k, cube_name=cube_name
+        )
+
+        for result_i in graph_results:
+            result_i.metadata.memory_type = "OuterMemory"
+        # Merge and deduplicate by ID
+        combined = {item.id: item for item in graph_results}
+
+        graph_ids = {item.id for item in graph_results}
+        combined_ids = set(combined.keys())
+        lost_ids = graph_ids - combined_ids
+
+        if lost_ids:
+            print(
+                f"[DEBUG] The following nodes were in graph_results but missing in combined: {lost_ids}"
+            )
+
+        return list(combined.values())
+
     def _graph_recall(
         self, parsed_goal: ParsedTaskGoal, memory_scope: str
     ) -> list[TextualMemoryItem]:
@@ -108,7 +165,7 @@ class GraphMemoryRetriever:
             return []
 
         # Load nodes and post-filter
-        node_dicts = self.graph_store.get_nodes(list(candidate_ids), include_embedding=True)
+        node_dicts = self.graph_store.get_nodes(list(candidate_ids), include_embedding=False)
 
         final_nodes = []
         for node in node_dicts:
@@ -134,31 +191,72 @@ class GraphMemoryRetriever:
         query_embedding: list[list[float]],
         memory_scope: str,
         top_k: int = 20,
-        max_num: int = 5,
+        max_num: int = 3,
+        cube_name: str | None = None,
+        search_filter: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """
-        # TODO: tackle with post-filter and pre-filter(5.18+) better.
         Perform vector-based similarity retrieval using query embedding.
+        # TODO: tackle with post-filter and pre-filter(5.18+) better.
         """
-        all_matches = []
+        if not query_embedding:
+            return []
 
-        def search_single(vec):
+        def search_single(vec, filt=None):
             return (
-                self.graph_store.search_by_embedding(vector=vec, top_k=top_k, scope=memory_scope)
+                self.graph_store.search_by_embedding(
+                    vector=vec,
+                    top_k=top_k,
+                    scope=memory_scope,
+                    cube_name=cube_name,
+                    search_filter=filt,
+                )
                 or []
             )
 
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            futures = [executor.submit(search_single, vec) for vec in query_embedding[:max_num]]
-            for future in concurrent.futures.as_completed(futures):
-                result = future.result()
-                all_matches.extend(result)
+        def search_path_a():
+            """Path A: search without filter"""
+            path_a_hits = []
+            with ContextThreadPoolExecutor() as executor:
+                futures = [
+                    executor.submit(search_single, vec, None) for vec in query_embedding[:max_num]
+                ]
+                for f in concurrent.futures.as_completed(futures):
+                    path_a_hits.extend(f.result() or [])
+            return path_a_hits
 
-        if not all_matches:
-            return []
+        def search_path_b():
+            """Path B: search with filter"""
+            if not search_filter:
+                return []
+            path_b_hits = []
+            with ContextThreadPoolExecutor() as executor:
+                futures = [
+                    executor.submit(search_single, vec, search_filter)
+                    for vec in query_embedding[:max_num]
+                ]
+                for f in concurrent.futures.as_completed(futures):
+                    path_b_hits.extend(f.result() or [])
+            return path_b_hits
+
+        # Execute both paths concurrently
+        all_hits = []
+        with ContextThreadPoolExecutor(max_workers=2) as executor:
+            path_a_future = executor.submit(search_path_a)
+            path_b_future = executor.submit(search_path_b)
 
-        # Step 3: Extract matched IDs and retrieve full nodes
-        unique_ids = set({r["id"] for r in all_matches})
-        node_dicts = self.graph_store.get_nodes(list(unique_ids), include_embedding=True)
+            all_hits.extend(path_a_future.result())
+            all_hits.extend(path_b_future.result())
 
-        return [TextualMemoryItem.from_dict(record) for record in node_dicts]
+        if not all_hits:
+            return []
+
+        # merge and deduplicate
+        unique_ids = {r["id"] for r in all_hits if r.get("id")}
+        node_dicts = (
+            self.graph_store.get_nodes(
+                list(unique_ids), include_embedding=False, cube_name=cube_name
+            )
+            or []
+        )
+        return [TextualMemoryItem.from_dict(n) for n in node_dicts]