refactor: add searcher to handler_init; remove info log from task_queue

Author: tangg555

examples/mem_scheduler/api_w_scheduler.py

@@ -25,6 +25,10 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
     print(f"My test handler received {len(messages)} messages:")
     for msg in messages:
         print(f" my_test_handler - {msg.item_id}: {msg.content}")
+        user_status_running = handle_scheduler_status(
+            user_name=USER_MEM_CUBE, mem_scheduler=mem_scheduler, instance_id="api_w_scheduler"
+        )
+        print(f"[Monitor] Status for {USER_MEM_CUBE} after submit:", user_status_running)
 
 
 # 2. Register the handler
@@ -56,10 +60,6 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
 
 # 5.1 Monitor status for specific mem_cube while running
 USER_MEM_CUBE = "test_mem_cube"
-user_status_running = handle_scheduler_status(
-    user_name=USER_MEM_CUBE, mem_scheduler=mem_scheduler, instance_id="api_w_scheduler"
-)
-print(f"[Monitor] Status for {USER_MEM_CUBE} after submit:", user_status_running)
 
 # 6. Wait for messages to be processed (limited to 100 checks)
 print("Waiting for messages to be consumed (max 100 checks)...")

src/memos/api/handlers/base_handler.py

@@ -9,6 +9,7 @@
 
 from memos.log import get_logger
 from memos.mem_scheduler.base_scheduler import BaseScheduler
+from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
 
 
 logger = get_logger(__name__)
@@ -28,6 +29,7 @@ def __init__(
         naive_mem_cube: Any | None = None,
         mem_reader: Any | None = None,
         mem_scheduler: Any | None = None,
+        searcher: Any | None = None,
         embedder: Any | None = None,
         reranker: Any | None = None,
         graph_db: Any | None = None,
@@ -58,6 +60,7 @@ def __init__(
         self.naive_mem_cube = naive_mem_cube
         self.mem_reader = mem_reader
         self.mem_scheduler = mem_scheduler
+        self.searcher = searcher
         self.embedder = embedder
         self.reranker = reranker
         self.graph_db = graph_db
@@ -128,6 +131,11 @@ def mem_scheduler(self) -> BaseScheduler:
         """Get scheduler instance."""
         return self.deps.mem_scheduler
 
+    @property
+    def searcher(self) -> Searcher:
+        """Get scheduler instance."""
+        return self.deps.searcher
+
     @property
     def embedder(self):
         """Get embedder instance."""

src/memos/api/handlers/component_init.py

@@ -5,6 +5,8 @@
 including databases, LLMs, memory systems, and schedulers.
 """
 
+import os
+
 from typing import TYPE_CHECKING, Any
 
 from memos.api.config import APIConfig
@@ -38,6 +40,10 @@
 from memos.memories.textual.simple_preference import SimplePreferenceTextMemory
 from memos.memories.textual.simple_tree import SimpleTreeTextMemory
 from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.tree import TreeTextMemory
 from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
     InternetRetrieverFactory,
 )
@@ -47,7 +53,7 @@
 
 if TYPE_CHECKING:
     from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
-
+    from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
 logger = get_logger(__name__)
 
 
@@ -205,6 +211,13 @@ def init_server() -> dict[str, Any]:
 
     logger.debug("MemCube created")
 
+    tree_mem: TreeTextMemory = naive_mem_cube.text_mem
+    searcher: Searcher = tree_mem.get_searcher(
+        manual_close_internet=os.getenv("ENABLE_INTERNET", "true").lower() == "false",
+        moscube=False,
+    )
+    logger.debug("Searcher created")
+
     # Initialize Scheduler
     scheduler_config_dict = APIConfig.get_scheduler_config()
     scheduler_config = SchedulerConfigFactory(
@@ -217,16 +230,14 @@ def init_server() -> dict[str, Any]:
         db_engine=BaseDBManager.create_default_sqlite_engine(),
         mem_reader=mem_reader,
     )
-    mem_scheduler.init_mem_cube(mem_cube=naive_mem_cube)
+    mem_scheduler.init_mem_cube(mem_cube=naive_mem_cube, searcher=searcher)
     logger.debug("Scheduler initialized")
 
     # Initialize SchedulerAPIModule
     api_module = mem_scheduler.api_module
 
     # Start scheduler if enabled
-    import os
-
-    if os.getenv("API_SCHEDULER_ON", True):
+    if os.getenv("API_SCHEDULER_ON", "true").lower() == "true":
         mem_scheduler.start()
         logger.info("Scheduler started")
 
@@ -253,6 +264,7 @@ def init_server() -> dict[str, Any]:
         "mos_server": mos_server,
         "mem_scheduler": mem_scheduler,
         "naive_mem_cube": naive_mem_cube,
+        "searcher": searcher,
         "api_module": api_module,
         "vector_db": vector_db,
         "pref_extractor": pref_extractor,

src/memos/api/handlers/search_handler.py

@@ -18,7 +18,7 @@
 from memos.api.product_models import APISearchRequest, SearchResponse
 from memos.context.context import ContextThreadPoolExecutor
 from memos.log import get_logger
-from memos.mem_scheduler.schemas.general_schemas import SearchMode
+from memos.mem_scheduler.schemas.general_schemas import FINE_STRATEGY, FineStrategy, SearchMode
 from memos.types import MOSSearchResult, UserContext
 
 
@@ -40,7 +40,7 @@ def __init__(self, dependencies: HandlerDependencies):
             dependencies: HandlerDependencies instance
         """
         super().__init__(dependencies)
-        self._validate_dependencies("naive_mem_cube", "mem_scheduler")
+        self._validate_dependencies("naive_mem_cube", "mem_scheduler", "searcher")
 
     def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse:
         """
@@ -211,11 +211,17 @@ def _fast_search(
 
         return formatted_memories
 
+    def _deep_search(
+        self, search_req: APISearchRequest, user_context: UserContext, max_thinking_depth: int
+    ) -> list:
+        logger.error("waiting to be implemented")
+        return []
+
     def _fine_search(
         self,
         search_req: APISearchRequest,
         user_context: UserContext,
-    ) -> list:
+    ) -> list[str]:
         """
         Fine-grained search with query enhancement.
 
@@ -226,19 +232,22 @@ def _fine_search(
         Returns:
             List of enhanced search results
         """
+        if FINE_STRATEGY == FineStrategy.DEEP_SEARCH:
+            return self._deep_search(
+                search_req=search_req, user_context=user_context, max_thinking_depth=3
+            )
+
         target_session_id = search_req.session_id or "default_session"
         search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
 
-        searcher = self.mem_scheduler.searcher
-
         info = {
             "user_id": search_req.user_id,
             "session_id": target_session_id,
             "chat_history": search_req.chat_history,
         }
 
         # Fine retrieve
-        fast_retrieved_memories = searcher.retrieve(
+        raw_retrieved_memories = self.searcher.retrieve(
             query=search_req.query,
             user_name=user_context.mem_cube_id,
             top_k=search_req.top_k,
@@ -250,8 +259,8 @@ def _fine_search(
         )
 
         # Post retrieve
-        fast_memories = searcher.post_retrieve(
-            retrieved_results=fast_retrieved_memories,
+        raw_memories = self.searcher.post_retrieve(
+            retrieved_results=raw_retrieved_memories,
             top_k=search_req.top_k,
             user_name=user_context.mem_cube_id,
             info=info,
@@ -260,22 +269,22 @@ def _fine_search(
         # Enhance with query
         enhanced_memories, _ = self.mem_scheduler.retriever.enhance_memories_with_query(
             query_history=[search_req.query],
-            memories=fast_memories,
+            memories=raw_memories,
         )
 
-        if len(enhanced_memories) < len(fast_memories):
+        if len(enhanced_memories) < len(raw_memories):
             logger.info(
-                f"Enhanced memories ({len(enhanced_memories)}) are less than fast memories ({len(fast_memories)}). Recalling for more."
+                f"Enhanced memories ({len(enhanced_memories)}) are less than raw memories ({len(raw_memories)}). Recalling for more."
             )
             missing_info_hint, trigger = self.mem_scheduler.retriever.recall_for_missing_memories(
                 query=search_req.query,
-                memories=fast_memories,
+                memories=raw_memories,
             )
-            retrieval_size = len(fast_memories) - len(enhanced_memories)
+            retrieval_size = len(raw_memories) - len(enhanced_memories)
             logger.info(f"Retrieval size: {retrieval_size}")
             if trigger:
                 logger.info(f"Triggering additional search with hint: {missing_info_hint}")
-                additional_memories = searcher.search(
+                additional_memories = self.searcher.search(
                     query=missing_info_hint,
                     user_name=user_context.mem_cube_id,
                     top_k=retrieval_size,
@@ -286,7 +295,7 @@ def _fine_search(
                 )
             else:
                 logger.info("Not triggering additional search, using fast memories.")
-                additional_memories = fast_memories[:retrieval_size]
+                additional_memories = raw_memories[:retrieval_size]
 
             enhanced_memories += additional_memories
             logger.info(

src/memos/mem_scheduler/base_scheduler.py

@@ -54,11 +54,11 @@
 from memos.memories.activation.kv import KVCacheMemory
 from memos.memories.activation.vllmkv import VLLMKVCacheItem, VLLMKVCacheMemory
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
+from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
 from memos.templates.mem_scheduler_prompts import MEMORY_ASSEMBLY_TEMPLATE
 
 
 if TYPE_CHECKING:
-    from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
     from memos.reranker.http_bge import HTTPBGEReranker
 
 
@@ -141,14 +141,21 @@ def __init__(self, config: BaseSchedulerConfig):
         self.auth_config = None
         self.rabbitmq_config = None
 
-    def init_mem_cube(self, mem_cube):
+    def init_mem_cube(
+        self,
+        mem_cube: BaseMemCube,
+        searcher: Searcher | None = None,
+    ):
         self.mem_cube = mem_cube
         self.text_mem: TreeTextMemory = self.mem_cube.text_mem
-        self.searcher: Searcher = self.text_mem.get_searcher(
-            manual_close_internet=os.getenv("ENABLE_INTERNET", "true").lower() == "false",
-            moscube=False,
-        )
         self.reranker: HTTPBGEReranker = self.text_mem.reranker
+        if searcher is None:
+            self.searcher: Searcher = self.text_mem.get_searcher(
+                manual_close_internet=os.getenv("ENABLE_INTERNET", "true").lower() == "false",
+                moscube=False,
+            )
+        else:
+            self.searcher = searcher
 
     def initialize_modules(
         self,

src/memos/mem_scheduler/schemas/general_schemas.py

@@ -18,6 +18,7 @@ class FineStrategy(str, Enum):
 
     REWRITE = "rewrite"
     RECREATE = "recreate"
+    DEEP_SEARCH = "deep_search"
 
 
 FILE_PATH = Path(__file__).absolute()

src/memos/mem_scheduler/task_schedule_modules/task_queue.py

@@ -112,10 +112,10 @@ def get_messages(self, batch_size: int) -> list[ScheduleMessageItem]:
             )
 
             messages.extend(fetched)
-
-        logger.info(
-            f"Fetched {len(messages)} messages across users with per-user batch_size={batch_size}"
-        )
+        if len(messages) > 0:
+            logger.debug(
+                f"Fetched {len(messages)} messages across users with per-user batch_size={batch_size}"
+            )
         return messages
 
     def clear(self):