refactor add_before_search from mem_reader to SingleCubeView

Author: tangg555

src/memos/mem_reader/simple_struct.py

@@ -561,6 +561,7 @@ def add_before_search(
         self,
         messages: list[dict],
         memory_list: list[TextualMemoryItem],
+        user_name: str,
         info: dict[str, Any],
     ) -> list[TextualMemoryItem]:
         # Build input objects with memory text and metadata (timestamps, sources, etc.)
@@ -580,7 +581,7 @@ def add_before_search(
         for idx, mem in enumerate(memory_list):
             try:
                 related_memories = self.searcher.search(
-                    query=mem.memory, top_k=3, mode="fast", info=info
+                    query=mem.memory, top_k=3, mode="fast", user_nam=user_name, info=info
                 )
                 related_text = "None"
                 if related_memories:

src/memos/mem_reader/utils.py

@@ -1,16 +1,7 @@
 import json
-import os
 import re
 
-from typing import Any
-
 from memos import log
-from memos.api.config import APIConfig
-from memos.configs.graph_db import GraphDBConfigFactory
-from memos.configs.reranker import RerankerConfigFactory
-from memos.graph_dbs.factory import GraphStoreFactory
-from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
-from memos.reranker.factory import RerankerFactory
 
 
 logger = log.get_logger(__name__)
@@ -164,47 +155,3 @@ def parse_keep_filter_response(text: str) -> tuple[bool, dict[int, dict]]:
                 "reason": reason,
             }
     return (len(result) > 0), result
-
-
-def build_graph_db_config(user_id: str = "default") -> dict[str, Any]:
-    graph_db_backend_map = {
-        "neo4j-community": APIConfig.get_neo4j_community_config(user_id=user_id),
-        "neo4j": APIConfig.get_neo4j_config(user_id=user_id),
-        "nebular": APIConfig.get_nebular_config(user_id=user_id),
-        "polardb": APIConfig.get_polardb_config(user_id=user_id),
-    }
-
-    graph_db_backend = os.getenv("NEO4J_BACKEND", "nebular").lower()
-    return GraphDBConfigFactory.model_validate(
-        {
-            "backend": graph_db_backend,
-            "config": graph_db_backend_map[graph_db_backend],
-        }
-    )
-
-
-def build_reranker_config() -> dict[str, Any]:
-    return RerankerConfigFactory.model_validate(APIConfig.get_reranker_config())
-
-
-def init_searcher(llm, embedder) -> Searcher:
-    """Initialize a Searcher instance for SimpleStructMemReader."""
-
-    # Build configs
-    graph_db_config = build_graph_db_config()
-    reranker_config = build_reranker_config()
-
-    # Create instances
-    graph_db = GraphStoreFactory.from_config(graph_db_config)
-    reranker = RerankerFactory.from_config(reranker_config)
-
-    # Create Searcher
-    searcher = Searcher(
-        dispatcher_llm=llm,
-        graph_store=graph_db,
-        embedder=embedder,
-        reranker=reranker,
-        manual_close_internet=os.getenv("ENABLE_INTERNET", "true").lower() == "false",
-    )
-
-    return searcher

src/memos/multi_mem_cube/single_cube.py

@@ -15,6 +15,7 @@
 )
 from memos.context.context import ContextThreadPoolExecutor
 from memos.log import get_logger
+from memos.mem_reader.utils import parse_keep_filter_response
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 from memos.mem_scheduler.schemas.task_schemas import (
     ADD_TASK_LABEL,
@@ -23,6 +24,7 @@
     PREF_ADD_TASK_LABEL,
 )
 from memos.multi_mem_cube.views import MemCubeView
+from memos.templates.mem_reader_prompts import PROMPT_MAPPING
 from memos.types.general_types import (
     FINE_STRATEGY,
     FineStrategy,
@@ -41,6 +43,7 @@
     from memos.mem_cube.navie import NaiveMemCube
     from memos.mem_reader.simple_struct import SimpleStructMemReader
     from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
+    from memos.memories.textual.item import TextualMemoryItem
 
 
 @dataclass
@@ -631,6 +634,104 @@ def _process_pref_mem(
                 for memory_id, memory in zip(pref_ids_local, pref_memories_local, strict=False)
             ]
 
+    def add_before_search(
+        self,
+        messages: list[dict],
+        memory_list: list[TextualMemoryItem],
+        user_name: str,
+        info: dict[str, Any],
+    ) -> list[TextualMemoryItem]:
+        # Build input objects with memory text and metadata (timestamps, sources, etc.)
+        template = PROMPT_MAPPING["add_before_search"]
+
+        if not self.searcher:
+            self.logger.warning("[add_before_search] Searcher is not initialized, skipping check.")
+            return memory_list
+
+        # 1. Gather candidates and search for related memories
+        candidates_data = []
+        for idx, mem in enumerate(memory_list):
+            try:
+                related_memories = self.searcher.search(
+                    query=mem.memory, top_k=3, mode="fast", user_name=user_name, info=info
+                )
+                related_text = "None"
+                if related_memories:
+                    related_text = "\n".join([f"- {r.memory}" for r in related_memories])
+
+                candidates_data.append(
+                    {"idx": idx, "new_memory": mem.memory, "related_memories": related_text}
+                )
+            except Exception as e:
+                self.logger.error(
+                    f"[add_before_search] Search error for memory '{mem.memory}': {e}"
+                )
+                # If search fails, we can either skip this check or treat related as empty
+                candidates_data.append(
+                    {
+                        "idx": idx,
+                        "new_memory": mem.memory,
+                        "related_memories": "None (Search Failed)",
+                    }
+                )
+
+        if not candidates_data:
+            return memory_list
+
+        # 2. Build Prompt
+        messages_inline = "\n".join(
+            [
+                f"- [{message.get('role', 'unknown')}]: {message.get('content', '')}"
+                for message in messages
+            ]
+        )
+
+        candidates_inline_dict = {
+            str(item["idx"]): {
+                "new_memory": item["new_memory"],
+                "related_memories": item["related_memories"],
+            }
+            for item in candidates_data
+        }
+
+        candidates_inline = json.dumps(candidates_inline_dict, ensure_ascii=False, indent=2)
+
+        prompt = template.format(
+            messages_inline=messages_inline, candidates_inline=candidates_inline
+        )
+
+        # 3. Call LLM
+        try:
+            raw = self.mem_reader.llm.generate([{"role": "user", "content": prompt}])
+            success, parsed_result = parse_keep_filter_response(raw)
+
+            if not success:
+                self.logger.warning(
+                    "[add_before_search] Failed to parse LLM response, keeping all."
+                )
+                return memory_list
+
+            # 4. Filter
+            filtered_list = []
+            for idx, mem in enumerate(memory_list):
+                res = parsed_result.get(idx)
+                if not res:
+                    filtered_list.append(mem)
+                    continue
+
+                if res.get("keep", True):
+                    filtered_list.append(mem)
+                else:
+                    self.logger.info(
+                        f"[add_before_search] Dropping memory: '{mem.memory}', reason: '{res.get('reason')}'"
+                    )
+
+            return filtered_list
+
+        except Exception as e:
+            self.logger.error(f"[add_before_search] LLM execution error: {e}")
+            return memory_list
+
     def _process_text_mem(
         self,
         add_req: APIADDRequest,