fix bugs: modify mos_for_test_scheduler.py and fix bugs of scheduler dispatch

Author: tangg555

examples/data/config/mem_scheduler/general_scheduler_config.yaml

@@ -5,5 +5,5 @@ config:
   act_mem_update_interval: 30
   context_window_size: 5
   thread_pool_max_workers: 5
-  consume_interval_seconds: 3
+  consume_interval_seconds: 1
   enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml

@@ -37,7 +37,7 @@ mem_scheduler:
     act_mem_update_interval: 30
     context_window_size: 5
     thread_pool_max_workers: 10
-    consume_interval_seconds: 3
+    consume_interval_seconds: 1
     enable_parallel_dispatch: true
 max_turns_window: 20
 top_k: 5

examples/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml

@@ -39,7 +39,7 @@ mem_scheduler:
     act_mem_update_interval: 30
     context_window_size: 5
     thread_pool_max_workers: 10
-    consume_interval_seconds: 3
+    consume_interval_seconds: 1
     enable_parallel_dispatch: true
 max_turns_window: 20
 top_k: 5

…amples/mem_scheduler/schedule_w_memos.py …mples/mem_scheduler/memos_w_scheduler.pyexamples/mem_scheduler/schedule_w_memos.py renamed to examples/mem_scheduler/memos_w_scheduler.py examples/mem_scheduler/schedule_w_memos.py renamed to examples/mem_scheduler/memos_w_scheduler.py

@@ -1,6 +1,7 @@
 import shutil
 import sys
 
+from datetime import datetime
 from pathlib import Path
 from queue import Queue
 from typing import TYPE_CHECKING
@@ -13,7 +14,11 @@
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
-from memos.mem_scheduler.modules.schemas import NOT_APPLICABLE_TYPE
+from memos.mem_scheduler.modules.schemas import (
+    NOT_APPLICABLE_TYPE,
+    QUERY_LABEL,
+    ScheduleMessageItem,
+)
 from memos.mem_scheduler.mos_for_test_scheduler import MOSForTestScheduler
 
 
@@ -184,6 +189,17 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
             query_history=None,
         )
 
+        # test query_consume
+        message_item = ScheduleMessageItem(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+            label=QUERY_LABEL,
+            content=query,
+            timestamp=datetime.now(),
+        )
+        mos.mem_scheduler._query_message_consumer(messages=[message_item])
+
         # test activation memory update
         mos.mem_scheduler.update_activation_memory_periodically(
             interval_seconds=0,

examples/mem_scheduler/try_schedule_modules.py

@@ -58,10 +58,15 @@ def __init__(self, config: MOSConfig, user_manager: UserManager | None = None):
                 f"User '{self.user_id}' does not exist or is inactive. Please create user first."
             )
 
-        # Lazy initialization marker
+        # Initialize mem_scheduler
         self._mem_scheduler_lock = Lock()
         self.enable_mem_scheduler = self.config.get("enable_mem_scheduler", False)
-        self._mem_scheduler: GeneralScheduler = None
+        if self.enable_mem_scheduler:
+            self._mem_scheduler = self._initialize_mem_scheduler()
+            self._mem_scheduler.mem_cubes = self.mem_cubes
+        else:
+            self._mem_scheduler: GeneralScheduler = None
+
         logger.info(f"MOS initialized for user: {self.user_id}")
 
     @property
@@ -93,14 +98,16 @@ def mem_scheduler(self, value: GeneralScheduler | None) -> None:
             else:
                 logger.debug("Memory scheduler cleared")
 
-    def _initialize_mem_scheduler(self):
+    def _initialize_mem_scheduler(self) -> GeneralScheduler:
         """Initialize the memory scheduler on first access."""
         if not self.config.enable_mem_scheduler:
             logger.debug("Memory scheduler is disabled in config")
             self._mem_scheduler = None
+            return self._mem_scheduler
         elif not hasattr(self.config, "mem_scheduler"):
             logger.error("Config of Memory scheduler is not available")
             self._mem_scheduler = None
+            return self._mem_scheduler
         else:
             logger.info("Initializing memory scheduler...")
             scheduler_config = self.config.mem_scheduler
@@ -111,13 +118,16 @@ def _initialize_mem_scheduler(self):
                     f"Memory reader of type {type(self.mem_reader).__name__} "
                     "missing required 'llm' attribute"
                 )
-                self._mem_scheduler.initialize_modules(chat_llm=self.chat_llm)
+                self._mem_scheduler.initialize_modules(
+                    chat_llm=self.chat_llm, process_llm=self.chat_llm
+                )
             else:
                 # Configure scheduler modules
                 self._mem_scheduler.initialize_modules(
                     chat_llm=self.chat_llm, process_llm=self.mem_reader.llm
                 )
             self._mem_scheduler.start()
+            return self._mem_scheduler
 
     def mem_scheduler_on(self) -> bool:
         if not self.config.enable_mem_scheduler or self._mem_scheduler is None:

src/memos/mem_os/core.py

@@ -36,7 +36,7 @@ def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         Args:
             messages: List of query messages to process
         """
-        logger.debug(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
+        logger.info(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
 
         # Process the query in a session turn
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
@@ -67,7 +67,7 @@ def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         Args:
           messages: List of answer messages to process
         """
-        logger.debug(f"Messages {messages} assigned to {ANSWER_LABEL} handler.")
+        logger.info(f"Messages {messages} assigned to {ANSWER_LABEL} handler.")
         # Process the query in a session turn
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
@@ -93,7 +93,7 @@ def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                     )
 
     def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
-        logger.debug(f"Messages {messages} assigned to {ADD_LABEL} handler.")
+        logger.info(f"Messages {messages} assigned to {ADD_LABEL} handler.")
         # Process the query in a session turn
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
@@ -155,32 +155,59 @@ def process_session_turn(
             logger.error("Not implemented!", exc_info=True)
             return
 
+        logger.info(f"Processing {len(queries)} queries.")
+
         working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
         text_working_memory: list[str] = [w_m.memory for w_m in working_memory]
         intent_result = self.monitor.detect_intent(
             q_list=query_history, text_working_memory=text_working_memory
         )
 
-        if intent_result["trigger_retrieval"]:
-            missing_evidences = intent_result["missing_evidences"]
-            num_evidence = len(missing_evidences)
-            k_per_evidence = max(1, top_k // max(1, num_evidence))
-            new_candidates = []
-            for item in missing_evidences:
-                logger.debug(f"missing_evidences: {item}")
-                results = self.retriever.search(
-                    query=item, mem_cube=mem_cube, top_k=k_per_evidence, method=self.search_method
-                )
-                logger.debug(f"search results for {missing_evidences}: {results}")
-                new_candidates.extend(results)
-
-            new_order_working_memory = self.retriever.replace_working_memory(
-                queries=queries,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-                original_memory=working_memory,
-                new_memory=new_candidates,
-                top_k=top_k,
+        time_trigger_flag = False
+        if self.monitor.timed_trigger(
+            last_time=self.monitor._last_query_consume_time,
+            interval_seconds=self.monitor.query_trigger_interval,
+        ):
+            time_trigger_flag = True
+            self._query_consume_time = True
+
+        if (not intent_result["trigger_retrieval"]) and (not time_trigger_flag):
+            logger.info(f"Query schedule not triggered. Intent_result: {intent_result}")
+            return
+        elif (not intent_result["trigger_retrieval"]) and time_trigger_flag:
+            logger.info("Query schedule is forced to trigger due to time ticker")
+            intent_result["trigger_retrieval"] = True
+            intent_result["missing_evidences"] = queries
+        else:
+            logger.info(
+                f"Query schedule is triggered, and missing_evidences: {intent_result['missing_evidences']}"
             )
-            logger.debug(f"size of new_order_working_memory: {len(new_order_working_memory)}")
+
+        missing_evidences = intent_result["missing_evidences"]
+        num_evidence = len(missing_evidences)
+        k_per_evidence = max(1, top_k // max(1, num_evidence))
+        new_candidates = []
+        for item in missing_evidences:
+            logger.debug(f"missing_evidences: {item}")
+            results = self.retriever.search(
+                query=item, mem_cube=mem_cube, top_k=k_per_evidence, method=self.search_method
+            )
+            logger.debug(f"search results for {missing_evidences}: {results}")
+            new_candidates.extend(results)
+
+        new_order_working_memory = self.retriever.replace_working_memory(
+            queries=queries,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+            original_memory=working_memory,
+            new_memory=new_candidates,
+            top_k=top_k,
+        )
+        logger.debug(f"size of new_order_working_memory: {len(new_order_working_memory)}")
+
+        self.monitor.update_memory_monitors(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )

src/memos/mem_scheduler/general_scheduler.py

@@ -22,7 +22,7 @@ class SchedulerDispatcher(BaseSchedulerModule):
     - Bulk handler registration
     """
 
-    def __init__(self, max_workers=3, enable_parallel_dispatch=False):
+    def __init__(self, max_workers=30, enable_parallel_dispatch=False):
         super().__init__()
         # Main dispatcher thread pool
         self.max_workers = max_workers
@@ -128,16 +128,13 @@ def dispatch(self, msg_list: list[ScheduleMessageItem]):
             else:
                 handler = self.handlers[label]
             # dispatch to different handler
-            logger.debug(f"Dispatch {len(msgs)} messages to {label} handler.")
+            logger.debug(f"Dispatch {len(msgs)} message(s) to {label} handler.")
             if self.enable_parallel_dispatch and self.dispatcher_executor is not None:
                 # Capture variables in lambda to avoid loop variable issues
-                # TODO check this
-                future = self.dispatcher_executor.submit(handler, msgs)
-                logger.debug(f"Dispatched {len(msgs)} messages as future task")
-                return future
+                self.dispatcher_executor.submit(handler, msgs)
+                logger.info(f"Dispatched {len(msgs)} message(s) as future task")
             else:
                 handler(msgs)
-                return None
 
     def join(self, timeout: float | None = None) -> bool:
         """Wait for all dispatched tasks to complete.
@@ -159,7 +156,7 @@ def shutdown(self) -> None:
         if self.dispatcher_executor is not None:
             self.dispatcher_executor.shutdown(wait=True)
         self._running = False
-        logger.info("Dispatcher has been shutdown")
+        logger.info("Dispatcher has been shutdown.")
 
     def __enter__(self):
         self._running = True

src/memos/mem_scheduler/modules/dispatcher.py

@@ -31,7 +31,8 @@ def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
 
         # hyper-parameters
         self.config: BaseSchedulerConfig = config
-        self.act_mem_update_interval = self.config.get("act_mem_update_interval", 300)
+        self.act_mem_update_interval = self.config.get("act_mem_update_interval", 30)
+        self.query_trigger_interval = self.config.get("query_trigger_interval", 10)
 
         # Partial Retention Strategy
         self.partial_retention_number = 2
@@ -46,6 +47,7 @@ def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
 
         # Lifecycle monitor
         self._last_activation_mem_update_time = datetime.min
+        self._last_query_consume_time = datetime.min
 
         self._process_llm = process_llm
 

src/memos/mem_scheduler/modules/monitor.py

@@ -44,7 +44,7 @@ def chat(self, query: str, user_id: str | None = None) -> str:
 
         chat_history = self.chat_history_manager[target_user_id]
 
-        topk_for_scheduler = 2
+        topk_for_scheduler = 5
 
         if self.config.enable_textual_memory and self.mem_cubes:
             memories_all = []
@@ -64,14 +64,10 @@ def chat(self, query: str, user_id: str | None = None) -> str:
                         content=query,
                         timestamp=datetime.now(),
                     )
-                    self.mem_scheduler.submit_messages(messages=[message_item])
 
-                self.mem_scheduler.monitor.register_memory_manager_if_not_exists(
-                    user_id=user_id,
-                    mem_cube_id=mem_cube_id,
-                    memory_monitors=self.mem_scheduler.monitor.working_memory_monitors,
-                    max_capacity=self.mem_scheduler.monitor.working_mem_monitor_capacity,
-                )
+                    # --- force to run mem_scheduler ---
+                    self.mem_scheduler.monitor.query_trigger_interval = 0
+                    self.mem_scheduler._query_message_consumer(messages=[message_item])
 
                 # from scheduler
                 scheduler_memories = self.mem_scheduler.monitor.get_monitor_memories(
@@ -80,13 +76,13 @@ def chat(self, query: str, user_id: str | None = None) -> str:
                     memory_type=MONITOR_WORKING_MEMORY_TYPE,
                     top_k=topk_for_scheduler,
                 )
+                print(f"Memories from the scheduler: {scheduler_memories}")
                 memories_all.extend(scheduler_memories)
 
                 # from mem_cube
-                memories = mem_cube.text_mem.search(
-                    query, top_k=self.config.top_k - topk_for_scheduler
-                )
+                memories = mem_cube.text_mem.search(query, top_k=self.config.top_k)
                 text_memories = [m.memory for m in memories]
+                print(f"Memories from search: {text_memories}")
                 memories_all.extend(text_memories)
 
                 memories_all = list(set(memories_all))