turn off graph recall

Author: 黑布林

src/memos/api/config.py

@@ -866,6 +866,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
                                 "UserMemory": os.getenv("NEBULAR_USER_MEMORY", 1e6),
                             },
                             "search_strategy": {
+                                "fast_graph": bool(os.getenv("FAST_GRAPH", "false") == "true"),
                                 "bm25": bool(os.getenv("BM25_CALL", "false") == "true"),
                                 "cot": bool(os.getenv("VEC_COT_CALL", "false") == "true"),
                             },
@@ -937,6 +938,7 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None:
                                 "UserMemory": os.getenv("NEBULAR_USER_MEMORY", 1e6),
                             },
                             "search_strategy": {
+                                "fast_graph": bool(os.getenv("FAST_GRAPH", "false") == "true"),
                                 "bm25": bool(os.getenv("BM25_CALL", "false") == "true"),
                                 "cot": bool(os.getenv("VEC_COT_CALL", "false") == "true"),
                             },

src/memos/memories/textual/simple_tree.py

@@ -70,12 +70,6 @@ def __init__(
         )
         logger.info(f"time init: bm25_retriever time is: {time.time() - time_start_bm}")
 
-        self.vec_cot = (
-            self.search_strategy["cot"]
-            if self.search_strategy and "cot" in self.search_strategy
-            else False
-        )
-
         time_start_rr = time.time()
         self.reranker = reranker
         logger.info(f"time init: reranker time is: {time.time() - time_start_rr}")
@@ -189,7 +183,7 @@ def search(
                 bm25_retriever=self.bm25_retriever,
                 internet_retriever=None,
                 moscube=moscube,
-                vec_cot=self.vec_cot,
+                search_strategy=self.search_strategy,
             )
         else:
             searcher = Searcher(
@@ -200,7 +194,7 @@ def search(
                 bm25_retriever=self.bm25_retriever,
                 internet_retriever=self.internet_retriever,
                 moscube=moscube,
-                vec_cot=self.vec_cot,
+                search_strategy=self.search_strategy,
             )
         return searcher.search(
             query, top_k, info, mode, memory_type, search_filter, user_name=user_name

src/memos/memories/textual/tree.py

@@ -51,11 +51,6 @@ def __init__(self, config: TreeTextMemoryConfig):
         self.bm25_retriever = (
             EnhancedBM25() if self.search_strategy and self.search_strategy["bm25"] else None
         )
-        self.vec_cot = (
-            self.search_strategy["cot"]
-            if self.search_strategy and "cot" in self.search_strategy
-            else False
-        )
 
         if config.reranker is None:
             default_cfg = RerankerConfigFactory.model_validate(
@@ -143,6 +138,7 @@ def get_searcher(
                 self.reranker,
                 internet_retriever=None,
                 moscube=moscube,
+                search_strategy=self.search_strategy,
             )
         else:
             searcher = Searcher(
@@ -152,6 +148,7 @@ def get_searcher(
                 self.reranker,
                 internet_retriever=self.internet_retriever,
                 moscube=moscube,
+                search_strategy=self.search_strategy,
             )
         return searcher
 

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

@@ -1,5 +1,4 @@
 import concurrent.futures
-import os
 
 from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import OllamaEmbedder
@@ -41,6 +40,7 @@ def retrieve(
         search_filter: dict | None = None,
         user_name: str | None = None,
         id_filter: dict | None = None,
+        use_fast_graph: bool = False,
     ) -> list[TextualMemoryItem]:
         """
         Perform hybrid memory retrieval:
@@ -70,7 +70,13 @@ def retrieve(
 
         with ContextThreadPoolExecutor(max_workers=3) as executor:
             # Structured graph-based retrieval
-            future_graph = executor.submit(self._graph_recall, parsed_goal, memory_scope, user_name)
+            future_graph = executor.submit(
+                self._graph_recall,
+                parsed_goal,
+                memory_scope,
+                user_name,
+                use_fast_graph=use_fast_graph,
+            )
             # Vector similarity search
             future_vector = executor.submit(
                 self._vector_recall,
@@ -156,14 +162,15 @@ def retrieve_from_cube(
         return list(combined.values())
 
     def _graph_recall(
-        self, parsed_goal: ParsedTaskGoal, memory_scope: str, user_name: str | None = None
+        self, parsed_goal: ParsedTaskGoal, memory_scope: str, user_name: str | None = None, **kwargs
     ) -> list[TextualMemoryItem]:
         """
         Perform structured node-based retrieval from Neo4j.
         - keys must match exactly (n.key IN keys)
         - tags must overlap with at least 2 input tags
         - scope filters by memory_type if provided
         """
+        use_fast_graph = kwargs.get("use_fast_graph", False)
 
         def process_node(node):
             meta = node.get("metadata", {})
@@ -185,7 +192,7 @@ def process_node(node):
                 return TextualMemoryItem.from_dict(node)
             return None
 
-        if os.getenv("FAST_GRAPH", "false") == "true":
+        if not use_fast_graph:
             candidate_ids = set()
 
             # 1) key-based OR branch

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

@@ -45,7 +45,7 @@ def __init__(
         bm25_retriever: EnhancedBM25 | None = None,
         internet_retriever: None = None,
         moscube: bool = False,
-        vec_cot: bool = False,
+        search_strategy: dict | None = None,
     ):
         self.graph_store = graph_store
         self.embedder = embedder
@@ -59,7 +59,9 @@ def __init__(
         # Create internet retriever from config if provided
         self.internet_retriever = internet_retriever
         self.moscube = moscube
-        self.vec_cot = vec_cot
+        self.use_fast_graph = (
+            search_strategy.get("fast_graph", "false") == "true" if search_strategy else False
+        )
 
         self._usage_executor = ContextThreadPoolExecutor(max_workers=4, thread_name_prefix="usage")
 
@@ -226,6 +228,7 @@ def _parse_task(
             context="\n".join(context),
             conversation=info.get("chat_history", []),
             mode=mode,
+            use_fast_graph=self.use_fast_graph,
         )
 
         query = parsed_goal.rephrased_query or query
@@ -340,6 +343,7 @@ def _retrieve_from_working_memory(
             search_filter=search_filter,
             user_name=user_name,
             id_filter=id_filter,
+            use_fast_graph=self.use_fast_graph,
         )
         return self.reranker.rerank(
             query=query,
@@ -369,7 +373,7 @@ def _retrieve_from_long_term_and_user(
 
         # chain of thinking
         cot_embeddings = []
-        if self.vec_cot:
+        if self.search_strategy["vec_cot"]:
             queries = self._cot_query(query)
             if len(queries) > 1:
                 cot_embeddings = self.embedder.embed(queries)
@@ -390,6 +394,7 @@ def _retrieve_from_long_term_and_user(
                         search_filter=search_filter,
                         user_name=user_name,
                         id_filter=id_filter,
+                        use_fast_graph=self.use_fast_graph,
                     )
                 )
             if memory_type in ["All", "UserMemory"]:
@@ -404,6 +409,7 @@ def _retrieve_from_long_term_and_user(
                         search_filter=search_filter,
                         user_name=user_name,
                         id_filter=id_filter,
+                        use_fast_graph=self.use_fast_graph,
                     )
                 )
 

src/memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py

@@ -1,4 +1,3 @@
-import os
 import traceback
 
 from string import Template
@@ -30,6 +29,7 @@ def parse(
         context: str = "",
         conversation: list[dict] | None = None,
         mode: str = "fast",
+        **kwargs,
     ) -> ParsedTaskGoal:
         """
         Parse user input into structured semantic layers.
@@ -39,19 +39,20 @@ def parse(
         - mode == 'fine': use LLM to parse structured topic/keys/tags
         """
         if mode == "fast":
-            return self._parse_fast(task_description)
+            return self._parse_fast(task_description, **kwargs)
         elif mode == "fine":
             if not self.llm:
                 raise ValueError("LLM not provided for slow mode.")
             return self._parse_fine(task_description, context, conversation)
         else:
             raise ValueError(f"Unknown mode: {mode}")
 
-    def _parse_fast(self, task_description: str, limit_num: int = 5) -> ParsedTaskGoal:
+    def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
         """
         Fast mode: simple jieba word split.
         """
-        if os.getenv("FAST_GRAPH", "false") == "true":
+        use_fast_graph = kwargs.get("use_fast_graph", False)
+        if use_fast_graph:
             desc_tokenized = self.tokenizer.tokenize_mixed(task_description)
             return ParsedTaskGoal(
                 memories=[task_description],