refactor: refactor deep search feature, now only allowing one-round deep search

Author: tangg555

src/memos/api/handlers/component_init.py

@@ -237,9 +237,6 @@ def init_server() -> dict[str, Any]:
     # Initialize SchedulerAPIModule
     api_module = mem_scheduler.api_module
 
-    # TODO: must remove!
-    mem_scheduler.memos_message_queue.debug_mode_on()
-
     # Start scheduler if enabled
     if os.getenv("API_SCHEDULER_ON", "true").lower() == "true":
         mem_scheduler.start()

src/memos/api/handlers/search_handler.py

@@ -224,7 +224,7 @@ def _deep_search(
             "chat_history": search_req.chat_history,
         }
 
-        return self.searcher.deep_search(
+        enhanced_memories = self.searcher.deep_search(
             query=search_req.query,
             user_name=user_context.mem_cube_id,
             top_k=search_req.top_k,
@@ -234,12 +234,14 @@ def _deep_search(
             search_filter=search_filter,
             info=info,
         )
+        formatted_memories = [format_memory_item(data) for data in enhanced_memories]
+        return formatted_memories
 
     def _fine_search(
         self,
         search_req: APISearchRequest,
         user_context: UserContext,
-    ) -> list[str]:
+    ) -> list:
         """
         Fine-grained search with query enhancement.
 

src/memos/mem_scheduler/base_scheduler.py

@@ -41,6 +41,7 @@
 )
 from memos.mem_scheduler.schemas.monitor_schemas import MemoryMonitorItem
 from memos.mem_scheduler.task_schedule_modules.dispatcher import SchedulerDispatcher
+from memos.mem_scheduler.task_schedule_modules.local_queue import SchedulerLocalQueue
 from memos.mem_scheduler.task_schedule_modules.redis_queue import SchedulerRedisQueue
 from memos.mem_scheduler.task_schedule_modules.task_queue import ScheduleTaskQueue
 from memos.mem_scheduler.utils.db_utils import get_utc_now
@@ -824,25 +825,65 @@ def get_running_tasks(self, filter_func: Callable | None = None) -> dict[str, di
 
         return result
 
+    @staticmethod
+    def init_task_status():
+        return {
+            "running": 0,
+            "remaining": 0,
+            "completed": 0,
+        }
+
+    def get_tasks_status(self):
+        task_status = self.init_task_status()
+        memos_message_queue = self.memos_message_queue.memos_message_queue
+        if isinstance(memos_message_queue, SchedulerRedisQueue):
+            stream_keys = memos_message_queue.get_stream_keys(
+                stream_key_prefix=memos_message_queue.stream_key_prefix
+            )
+            for stream_key in stream_keys:
+                if stream_key not in task_status:
+                    task_status[stream_key] = self.init_task_status()
+                # For Redis queue, prefer XINFO GROUPS to compute pending
+                groups_info = memos_message_queue.redis.xinfo_groups(stream_key)
+                if groups_info:
+                    for group in groups_info:
+                        if group.get("name") == memos_message_queue.consumer_group:
+                            task_status[stream_key]["running"] += int(group.get("pending", 0))
+                            task_status[stream_key]["remaining"] += int(group.get("remaining", 0))
+                            task_status["running"] += int(group.get("pending", 0))
+                            task_status["remaining"] += int(group.get("remaining", 0))
+                            break
+
+        elif isinstance(memos_message_queue, SchedulerLocalQueue):
+            running_task_count = self.dispatcher.get_running_task_count()
+            task_status["running"] = running_task_count
+            task_status["remaining"] = sum(memos_message_queue.qsize().values())
+        else:
+            logger.error(
+                f"type of self.memos_message_queue is {memos_message_queue}, which is not supported"
+            )
+            raise NotImplementedError()
+
     def mem_scheduler_wait(
         self, timeout: float = 180.0, poll: float = 0.1, log_every: float = 0.01
     ) -> bool:
         """
         Uses EWMA throughput, detects leaked `unfinished_tasks`, and waits for dispatcher.
         """
         deadline = time.monotonic() + timeout
+        memos_message_queue = self.memos_message_queue.memos_message_queue
 
         # --- helpers (local, no external deps) ---
         def _unfinished() -> int:
             """Prefer `unfinished_tasks`; fallback to `qsize()`."""
             try:
-                u = getattr(self.memos_message_queue, "unfinished_tasks", None)
+                u = getattr(memos_message_queue, "unfinished_tasks", None)
                 if u is not None:
                     return int(u)
             except Exception:
                 pass
             try:
-                return int(self.memos_message_queue.qsize())
+                return int(memos_message_queue.qsize())
             except Exception:
                 return 0
 
@@ -876,7 +917,7 @@ def _fmt_eta(seconds: float | None) -> str:
             # 1) read counters
             curr_unfinished = _unfinished()
             try:
-                qsz = int(self.memos_message_queue.qsize())
+                qsz = int(memos_message_queue.qsize())
             except Exception:
                 qsz = -1
 
@@ -892,14 +933,14 @@ def _fmt_eta(seconds: float | None) -> str:
                 except Exception:
                     pass
 
-            if isinstance(self.memos_message_queue, SchedulerRedisQueue):
+            if isinstance(memos_message_queue, SchedulerRedisQueue):
                 # For Redis queue, prefer XINFO GROUPS to compute pending
-                groups_info = self.memos_message_queue.redis.xinfo_groups(
-                    self.memos_message_queue.stream_key_prefix
+                groups_info = memos_message_queue.redis.xinfo_groups(
+                    memos_message_queue.stream_key_prefix
                 )
                 if groups_info:
                     for group in groups_info:
-                        if group.get("name") == self.memos_message_queue.consumer_group:
+                        if group.get("name") == memos_message_queue.consumer_group:
                             pend = int(group.get("pending", pend))
                             break
             else:
@@ -975,18 +1016,19 @@ def _fmt_eta(seconds: float | None) -> str:
 
     def _gather_queue_stats(self) -> dict:
         """Collect queue/dispatcher stats for reporting."""
+        memos_message_queue = self.memos_message_queue.memos_message_queue
         stats: dict[str, int | float | str] = {}
         stats["use_redis_queue"] = bool(self.use_redis_queue)
         # local queue metrics
         if not self.use_redis_queue:
             try:
-                stats["qsize"] = int(self.memos_message_queue.qsize())
+                stats["qsize"] = int(memos_message_queue.qsize())
             except Exception:
                 stats["qsize"] = -1
             # unfinished_tasks if available
             try:
                 stats["unfinished_tasks"] = int(
-                    getattr(self.memos_message_queue, "unfinished_tasks", 0) or 0
+                    getattr(memos_message_queue, "unfinished_tasks", 0) or 0
                 )
             except Exception:
                 stats["unfinished_tasks"] = -1

src/memos/mem_scheduler/task_schedule_modules/redis_queue.py

@@ -5,6 +5,7 @@
 the local memos_message_queue functionality in BaseScheduler.
 """
 
+import os
 import re
 import time
 
@@ -33,7 +34,9 @@ class SchedulerRedisQueue(RedisSchedulerModule):
 
     def __init__(
         self,
-        stream_key_prefix: str = "scheduler:messages:stream",
+        stream_key_prefix: str = os.getenv(
+            "MEMSCHEDULER_REDIS_STREAM_KEY_PREFIX", "scheduler:messages:stream"
+        ),
         consumer_group: str = "scheduler_group",
         consumer_name: str | None = "scheduler_consumer",
         max_len: int = 10000,
@@ -283,7 +286,7 @@ def qsize(self) -> int:
             logger.error(f"Failed to get Redis queue size: {e}")
             return 0
 
-    def get_stream_keys(self) -> list[str]:
+    def get_stream_keys(self, stream_key_prefix: str | None = None) -> list[str]:
         """
         List all Redis stream keys that match this queue's prefix.
 
@@ -293,16 +296,18 @@ def get_stream_keys(self) -> list[str]:
         if not self._redis_conn:
             return []
 
+        if stream_key_prefix is None:
+            stream_key_prefix = self.stream_key_prefix
         # First, get all keys that might match (using Redis pattern matching)
-        redis_pattern = f"{self.stream_key_prefix}:*"
+        redis_pattern = f"{stream_key_prefix}:*"
         raw_keys = [
             key.decode("utf-8") if isinstance(key, bytes) else key
             for key in self._redis_conn.scan_iter(match=redis_pattern)
         ]
 
         # Second, filter using Python regex to ensure exact prefix match
         # Escape special regex characters in the prefix, then add :.*
-        escaped_prefix = re.escape(self.stream_key_prefix)
+        escaped_prefix = re.escape(stream_key_prefix)
         regex_pattern = f"^{escaped_prefix}:"
         stream_keys = [key for key in raw_keys if re.match(regex_pattern, key)]