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

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", "")
@@ -225,7 +227,7 @@ class BochaAISearchRetriever:
             )
             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()
@@ -153,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:
@@ -181,34 +193,70 @@ class GraphMemoryRetriever:
         top_k: int = 20,
         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, cube_name=cube_name
+                    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)
-
-        if not all_matches:
+        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
+
+        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)
+
+            all_hits.extend(path_a_future.result())
+            all_hits.extend(path_b_future.result())
+
+        if not all_hits:
             return []
 
-        # 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, cube_name=cube_name
+        # 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(record) for record in node_dicts]
+        return [TextualMemoryItem.from_dict(n) for n in node_dicts]

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

@@ -1,8 +1,9 @@
-import concurrent.futures
 import json
+import traceback
 
 from datetime import datetime
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.factory import Neo4jGraphDB
 from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
@@ -42,13 +43,17 @@ class Searcher:
         self.internet_retriever = internet_retriever
         self.moscube = moscube
 
-        self._usage_executor = concurrent.futures.ThreadPoolExecutor(
-            max_workers=4, thread_name_prefix="usage"
-        )
+        self._usage_executor = ContextThreadPoolExecutor(max_workers=4, thread_name_prefix="usage")
 
     @timed
     def search(
-        self, query: str, top_k: int, info=None, mode="fast", memory_type="All"
+        self,
+        query: str,
+        top_k: int,
+        info=None,
+        mode="fast",
+        memory_type="All",
+        search_filter: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """
         Search for memories based on a query.
@@ -63,6 +68,7 @@ class Searcher:
             - 'fine': Uses a more detailed search process, invoking large models for higher precision, but slower performance.
             memory_type (str): Type restriction for search.
             ['All', 'WorkingMemory', 'LongTermMemory', 'UserMemory']
+            search_filter (dict, optional): Optional metadata filters for search results.
         Returns:
             list[TextualMemoryItem]: List of matching memories.
         """
@@ -78,9 +84,11 @@ class Searcher:
         else:
             logger.debug(f"[SEARCH] Received info dict: {info}")
 
-        parsed_goal, query_embedding, context, query = self._parse_task(query, info, mode)
+        parsed_goal, query_embedding, context, query = self._parse_task(
+            query, info, mode, search_filter=search_filter
+        )
         results = self._retrieve_paths(
-            query, parsed_goal, query_embedding, info, top_k, mode, memory_type
+            query, parsed_goal, query_embedding, info, top_k, mode, memory_type, search_filter
         )
         deduped = self._deduplicate_results(results)
         final_results = self._sort_and_trim(deduped, top_k)
@@ -96,7 +104,7 @@ class Searcher:
         return final_results
 
     @timed
-    def _parse_task(self, query, info, mode, top_k=5):
+    def _parse_task(self, query, info, mode, top_k=5, search_filter: dict | None = None):
         """Parse user query, do embedding search and create context"""
         context = []
         query_embedding = None
@@ -109,9 +117,24 @@ class Searcher:
             # retrieve related nodes by embedding
             related_nodes = [
                 self.graph_store.get_node(n["id"])
-                for n in self.graph_store.search_by_embedding(query_embedding, top_k=top_k)
+                for n in self.graph_store.search_by_embedding(
+                    query_embedding, top_k=top_k, search_filter=search_filter
+                )
             ]
-            context = list({node["memory"] for node in related_nodes})
+            memories = []
+            for node in related_nodes:
+                try:
+                    m = (
+                        node.get("memory")
+                        if isinstance(node, dict)
+                        else (getattr(node, "memory", None))
+                    )
+                    if isinstance(m, str) and m:
+                        memories.append(m)
+                except Exception:
+                    logger.error(f"[SEARCH] Error during search: {traceback.format_exc()}")
+                    continue
+            context = list(dict.fromkeys(memories))
 
             # optional: supplement context with internet knowledge
             """if self.internet_retriever:
@@ -135,10 +158,20 @@ class Searcher:
         return parsed_goal, query_embedding, context, query
 
     @timed
-    def _retrieve_paths(self, query, parsed_goal, query_embedding, info, top_k, mode, memory_type):
+    def _retrieve_paths(
+        self,
+        query,
+        parsed_goal,
+        query_embedding,
+        info,
+        top_k,
+        mode,
+        memory_type,
+        search_filter: dict | None = None,
+    ):
         """Run A/B/C retrieval paths in parallel"""
         tasks = []
-        with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        with ContextThreadPoolExecutor(max_workers=3) as executor:
             tasks.append(
                 executor.submit(
                     self._retrieve_from_working_memory,
@@ -147,6 +180,7 @@ class Searcher:
                     query_embedding,
                     top_k,
                     memory_type,
+                    search_filter,
                 )
             )
             tasks.append(
@@ -157,6 +191,7 @@ class Searcher:
                     query_embedding,
                     top_k,
                     memory_type,
+                    search_filter,
                 )
             )
             tasks.append(
@@ -193,14 +228,24 @@ class Searcher:
     # --- Path A
     @timed
     def _retrieve_from_working_memory(
-        self, query, parsed_goal, query_embedding, top_k, memory_type
+        self,
+        query,
+        parsed_goal,
+        query_embedding,
+        top_k,
+        memory_type,
+        search_filter: dict | None = None,
     ):
         """Retrieve and rerank from WorkingMemory"""
         if memory_type not in ["All", "WorkingMemory"]:
             logger.info(f"[PATH-A] '{query}'Skipped (memory_type does not match)")
             return []
         items = self.graph_retriever.retrieve(
-            query=query, parsed_goal=parsed_goal, top_k=top_k, memory_scope="WorkingMemory"
+            query=query,
+            parsed_goal=parsed_goal,
+            top_k=top_k,
+            memory_scope="WorkingMemory",
+            search_filter=search_filter,
         )
         return self.reranker.rerank(
             query=query,
@@ -208,37 +253,61 @@ class Searcher:
             graph_results=items,
             top_k=top_k,
             parsed_goal=parsed_goal,
+            search_filter=search_filter,
         )
 
     # --- Path B
     @timed
     def _retrieve_from_long_term_and_user(
-        self, query, parsed_goal, query_embedding, top_k, memory_type
+        self,
+        query,
+        parsed_goal,
+        query_embedding,
+        top_k,
+        memory_type,
+        search_filter: dict | None = None,
     ):
         """Retrieve and rerank from LongTermMemory and UserMemory"""
         results = []
-        if memory_type in ["All", "LongTermMemory"]:
-            results += self.graph_retriever.retrieve(
-                query=query,
-                parsed_goal=parsed_goal,
-                query_embedding=query_embedding,
-                top_k=top_k * 2,
-                memory_scope="LongTermMemory",
-            )
-        if memory_type in ["All", "UserMemory"]:
-            results += self.graph_retriever.retrieve(
-                query=query,
-                parsed_goal=parsed_goal,
-                query_embedding=query_embedding,
-                top_k=top_k * 2,
-                memory_scope="UserMemory",
-            )
+        tasks = []
+
+        with ContextThreadPoolExecutor(max_workers=2) as executor:
+            if memory_type in ["All", "LongTermMemory"]:
+                tasks.append(
+                    executor.submit(
+                        self.graph_retriever.retrieve,
+                        query=query,
+                        parsed_goal=parsed_goal,
+                        query_embedding=query_embedding,
+                        top_k=top_k * 2,
+                        memory_scope="LongTermMemory",
+                        search_filter=search_filter,
+                    )
+                )
+            if memory_type in ["All", "UserMemory"]:
+                tasks.append(
+                    executor.submit(
+                        self.graph_retriever.retrieve,
+                        query=query,
+                        parsed_goal=parsed_goal,
+                        query_embedding=query_embedding,
+                        top_k=top_k * 2,
+                        memory_scope="UserMemory",
+                        search_filter=search_filter,
+                    )
+                )
+
+            # Collect results from all tasks
+            for task in tasks:
+                results.extend(task.result())
+
         return self.reranker.rerank(
             query=query,
             query_embedding=query_embedding[0],
             graph_results=results,
             top_k=top_k,
             parsed_goal=parsed_goal,
+            search_filter=search_filter,
         )
 
     @timed
@@ -300,8 +369,7 @@ class Searcher:
         final_items = []
         for item, score in sorted_results:
             meta_data = item.metadata.model_dump()
-            if "relativity" not in meta_data:
-                meta_data["relativity"] = score
+            meta_data["relativity"] = score
             final_items.append(
                 TextualMemoryItem(
                     id=item.id,

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

@@ -3,15 +3,16 @@
 import json
 import uuid
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed
 from datetime import datetime
 
 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__)
@@ -150,7 +151,7 @@ class XinyuSearchRetriever:
         # Convert to TextualMemoryItem format
         memory_items: list[TextualMemoryItem] = []
 
-        with ThreadPoolExecutor(max_workers=8) as executor:
+        with ContextThreadPoolExecutor(max_workers=8) as executor:
             futures = [
                 executor.submit(self._process_result, result, query, parsed_goal, info)
                 for result in search_results
@@ -332,7 +333,7 @@ class XinyuSearchRetriever:
             )
             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)