MemoryOS 0.2.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

memos/mem_scheduler/monitors/dispatcher_monitor.py

@@ -0,0 +1,305 @@
+import threading
+import time
+
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from time import perf_counter
+
+from memos.configs.mem_scheduler import BaseSchedulerConfig
+from memos.log import get_logger
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
+from memos.mem_scheduler.general_modules.dispatcher import SchedulerDispatcher
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerDispatcherMonitor(BaseSchedulerModule):
+    """Monitors and manages scheduling operations with LLM integration."""
+
+    def __init__(self, config: BaseSchedulerConfig):
+        super().__init__()
+        self.config: BaseSchedulerConfig = config
+
+        self.check_interval = self.config.get("dispatcher_monitor_check_interval", 60)
+        self.max_failures = self.config.get("dispatcher_monitor_max_failures", 2)
+
+        # Registry of monitored thread pools
+        self._pools: dict[str, dict] = {}
+        self._pool_lock = threading.Lock()
+
+        # thread pool monitor
+        self._monitor_thread: threading.Thread | None = None
+        self._running = False
+        self._restart_in_progress = False
+
+        # modules with thread pool
+        self.dispatcher: SchedulerDispatcher | None = None
+        self.dispatcher_pool_name = "dispatcher"
+
+    def initialize(self, dispatcher: SchedulerDispatcher):
+        self.dispatcher = dispatcher
+        self.register_pool(
+            name=self.dispatcher_pool_name,
+            executor=self.dispatcher.dispatcher_executor,
+            max_workers=self.dispatcher.max_workers,
+            restart_on_failure=True,
+        )
+
+    def register_pool(
+        self,
+        name: str,
+        executor: ThreadPoolExecutor,
+        max_workers: int,
+        restart_on_failure: bool = True,
+    ) -> bool:
+        """
+        Register a thread pool for monitoring.
+
+        Args:
+            name: Unique identifier for the pool
+            executor: ThreadPoolExecutor instance to monitor
+            max_workers: Expected maximum worker count
+            restart_on_failure: Whether to restart if pool fails
+
+        Returns:
+            bool: True if registration succeeded, False if pool already registered
+        """
+        with self._pool_lock:
+            if name in self._pools:
+                logger.warning(f"Thread pool '{name}' is already registered")
+                return False
+
+            self._pools[name] = {
+                "executor": executor,
+                "max_workers": max_workers,
+                "restart": restart_on_failure,
+                "failure_count": 0,
+                "last_active": datetime.utcnow(),
+                "healthy": True,
+            }
+            logger.info(f"Registered thread pool '{name}' for monitoring")
+            return True
+
+    def unregister_pool(self, name: str) -> bool:
+        """
+        Remove a thread pool from monitoring.
+
+        Args:
+            name: Identifier of the pool to remove
+
+        Returns:
+            bool: True if removal succeeded, False if pool not found
+        """
+        with self._pool_lock:
+            if name not in self._pools:
+                logger.warning(f"Thread pool '{name}' not found in registry")
+                return False
+
+            del self._pools[name]
+            logger.info(f"Unregistered thread pool '{name}'")
+            return True
+
+    def _monitor_loop(self) -> None:
+        """Main monitoring loop that periodically checks all registered pools."""
+        logger.info(f"Starting monitor loop with {self.check_interval} second interval")
+
+        while self._running:
+            time.sleep(self.check_interval)
+            try:
+                self._check_pools_health()
+            except Exception as e:
+                logger.error(f"Error during health check: {e!s}", exc_info=True)
+
+        logger.debug("Monitor loop exiting")
+
+    def start(self) -> bool:
+        """
+        Start the monitoring thread.
+
+        Returns:
+            bool: True if monitor started successfully, False if already running
+        """
+        if self._running:
+            logger.warning("Dispatcher Monitor is already running")
+            return False
+
+        self._running = True
+        self._monitor_thread = threading.Thread(
+            target=self._monitor_loop, name="threadpool_monitor", daemon=True
+        )
+        self._monitor_thread.start()
+        logger.info("Dispatcher Monitor  monitor started")
+        return True
+
+    def stop(self) -> None:
+        """
+        Stop the monitoring thread and clean up all managed thread pools.
+        Ensures proper shutdown of all monitored executors.
+        """
+        if not self._running:
+            return
+
+        # Stop the monitoring loop
+        self._running = False
+        if self._monitor_thread and self._monitor_thread.is_alive():
+            self._monitor_thread.join(timeout=5)
+
+        # Shutdown all registered pools
+        with self._pool_lock:
+            for name, pool_info in self._pools.items():
+                executor = pool_info["executor"]
+                if not executor._shutdown:  # pylint: disable=protected-access
+                    try:
+                        logger.info(f"Shutting down thread pool '{name}'")
+                        executor.shutdown(wait=True, cancel_futures=True)
+                        logger.info(f"Successfully shut down thread pool '{name}'")
+                    except Exception as e:
+                        logger.error(f"Error shutting down pool '{name}': {e!s}", exc_info=True)
+
+        # Clear the pool registry
+        self._pools.clear()
+        logger.info("Thread pool monitor and all pools stopped")
+
+    def _check_pools_health(self) -> None:
+        """Check health of all registered thread pools."""
+        for name, pool_info in list(self._pools.items()):
+            is_healthy, reason = self._check_pool_health(
+                pool_info=pool_info,
+                stuck_max_interval=4,
+            )
+            logger.info(f"Pool '{name}'. is_healthy: {is_healthy}. pool_info: {pool_info}")
+            with self._pool_lock:
+                if is_healthy:
+                    pool_info["failure_count"] = 0
+                    pool_info["healthy"] = True
+                    return
+                else:
+                    pool_info["failure_count"] += 1
+                    pool_info["healthy"] = False
+                    logger.warning(
+                        f"Pool '{name}' unhealthy ({pool_info['failure_count']}/{self.max_failures}): {reason}"
+                    )
+
+            if (
+                pool_info["failure_count"] >= self.max_failures
+                and pool_info["restart"]
+                and not self._restart_in_progress
+            ):
+                self._restart_pool(name, pool_info)
+
+    def _check_pool_health(self, pool_info: dict, stuck_max_interval=4) -> tuple[bool, str]:
+        """
+        Check health of a single thread pool.
+
+        Args:
+            pool_info: Dictionary containing pool configuration
+
+        Returns:
+            Tuple: (is_healthy, reason) where reason explains failure if not healthy
+        """
+        executor = pool_info["executor"]
+
+        # Check if executor is shutdown
+        if executor._shutdown:  # pylint: disable=protected-access
+            return False, "Executor is shutdown"
+
+        # Check thread activity
+        active_threads = sum(
+            1
+            for t in threading.enumerate()
+            if t.name.startswith(executor._thread_name_prefix)  # pylint: disable=protected-access
+        )
+
+        # Check if no threads are active but should be
+        if active_threads == 0 and pool_info["max_workers"] > 0:
+            return False, "No active worker threads"
+
+        # Check if threads are stuck (no activity for 2 intervals)
+        time_delta = (datetime.utcnow() - pool_info["last_active"]).total_seconds()
+        if time_delta >= self.check_interval * stuck_max_interval:
+            return False, "No recent activity"
+
+        # If we got here, pool appears healthy
+        pool_info["last_active"] = datetime.utcnow()
+        return True, ""
+
+    def _restart_pool(self, name: str, pool_info: dict) -> None:
+        """
+        Attempt to restart a failed thread pool.
+
+        Args:
+            name: Name of the pool to restart
+            pool_info: Dictionary containing pool configuration
+        """
+        if self._restart_in_progress:
+            return
+
+        self._restart_in_progress = True
+        logger.warning(f"Attempting to restart thread pool '{name}'")
+
+        try:
+            old_executor = pool_info["executor"]
+            self.dispatcher.shutdown()
+
+            # Create new executor with same parameters
+            new_executor = ThreadPoolExecutor(
+                max_workers=pool_info["max_workers"],
+                thread_name_prefix=self.dispatcher.thread_name_prefix,  # pylint: disable=protected-access
+            )
+            self.unregister_pool(name=self.dispatcher_pool_name)
+            self.dispatcher.dispatcher_executor = new_executor
+            self.register_pool(
+                name=self.dispatcher_pool_name,
+                executor=self.dispatcher.dispatcher_executor,
+                max_workers=self.dispatcher.max_workers,
+                restart_on_failure=True,
+            )
+
+            # Replace in registry
+            start_time = perf_counter()
+            with self._pool_lock:
+                pool_info["executor"] = new_executor
+                pool_info["failure_count"] = 0
+                pool_info["healthy"] = True
+                pool_info["last_active"] = datetime.utcnow()
+
+                elapsed_time = perf_counter() - start_time
+                if elapsed_time > 1:
+                    logger.warning(f"Long lock wait: {elapsed_time:.3f}s")
+
+            # Shutdown old executor
+            try:
+                old_executor.shutdown(wait=False)
+            except Exception as e:
+                logger.error(f"Error shutting down old executor: {e!s}", exc_info=True)
+
+            logger.info(f"Successfully restarted thread pool '{name}'")
+        except Exception as e:
+            logger.error(f"Failed to restart pool '{name}': {e!s}", exc_info=True)
+        finally:
+            self._restart_in_progress = False
+
+    def get_status(self, name: str | None = None) -> dict:
+        """
+        Get status of monitored pools.
+
+        Args:
+            name: Optional specific pool name to check
+
+        Returns:
+            Dictionary of status information
+        """
+        with self._pool_lock:
+            if name:
+                return {name: self._pools.get(name, {}).copy()}
+            return {k: v.copy() for k, v in self._pools.items()}
+
+    def __enter__(self):
+        """Context manager entry point."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit point."""
+        self.stop()

memos/mem_scheduler/{modules/monitor.py → monitors/general_monitor.py}

@@ -1,29 +1,35 @@
 from datetime import datetime
+from threading import Lock
 from typing import Any
 
 from memos.configs.mem_scheduler import BaseSchedulerConfig
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
-from memos.mem_scheduler.modules.base import BaseSchedulerModule
-from memos.mem_scheduler.modules.misc import AutoDroppingQueue as Queue
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
+from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT,
+    DEFAULT_WEIGHT_VECTOR_FOR_RANKING,
     DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT,
     MONITOR_ACTIVATION_MEMORY_TYPE,
     MONITOR_WORKING_MEMORY_TYPE,
     MemCubeID,
-    MemoryMonitorManager,
     UserID,
 )
-from memos.mem_scheduler.utils import extract_json_dict
-from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
+from memos.mem_scheduler.schemas.monitor_schemas import (
+    MemoryMonitorItem,
+    MemoryMonitorManager,
+    QueryMonitorItem,
+    QueryMonitorQueue,
+)
+from memos.mem_scheduler.utils.misc_utils import extract_json_dict
+from memos.memories.textual.tree import TreeTextMemory
 
 
 logger = get_logger(__name__)
 
 
-class SchedulerMonitor(BaseSchedulerModule):
+class SchedulerGeneralMonitor(BaseSchedulerModule):
     """Monitors and manages scheduling operations with LLM integration."""
 
     def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
@@ -31,33 +37,79 @@ class SchedulerMonitor(BaseSchedulerModule):
 
         # 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
-        self.working_mem_monitor_capacity = DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT
-        self.activation_mem_monitor_capacity = DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT
+        self.working_mem_monitor_capacity = self.config.get(
+            "working_mem_monitor_capacity", DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT
+        )
+        self.activation_mem_monitor_capacity = self.config.get(
+            "activation_mem_monitor_capacity", DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT
+        )
 
         # attributes
-        self.query_history = Queue(maxsize=self.config.context_window_size)
-        self.intent_history = Queue(maxsize=self.config.context_window_size)
+        # recording query_messages
+        self.query_monitors: dict[UserID, dict[MemCubeID, QueryMonitorQueue[QueryMonitorItem]]] = {}
+
         self.working_memory_monitors: dict[UserID, dict[MemCubeID, MemoryMonitorManager]] = {}
         self.activation_memory_monitors: dict[UserID, dict[MemCubeID, MemoryMonitorManager]] = {}
 
         # Lifecycle monitor
-        self._last_activation_mem_update_time = datetime.min
+        self.last_activation_mem_update_time = datetime.min
+        self.last_query_consume_time = datetime.min
 
+        self._register_lock = Lock()
         self._process_llm = process_llm
 
+    def extract_query_keywords(self, query: str) -> list:
+        """Extracts core keywords from a user query based on specific semantic rules."""
+        prompt_name = "query_keywords_extraction"
+        prompt = self.build_prompt(
+            template_name=prompt_name,
+            query=query,
+        )
+        llm_response = self._process_llm.generate([{"role": "user", "content": prompt}])
+        try:
+            # Parse JSON output from LLM response
+            keywords = extract_json_dict(llm_response)
+            assert isinstance(keywords, list)
+        except Exception as e:
+            logger.error(
+                f"Failed to parse keywords from LLM response: {llm_response}. Error: {e!s}"
+            )
+            keywords = [query]
+        return keywords
+
+    def register_query_monitor_if_not_exists(
+        self,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
+    ) -> None:
+        # First check (lock-free, fast path)
+        if user_id in self.query_monitors and mem_cube_id in self.query_monitors[user_id]:
+            return
+
+        # Second check (with lock, ensures uniqueness)
+        with self._register_lock:
+            if user_id not in self.query_monitors:
+                self.query_monitors[user_id] = {}
+            if mem_cube_id not in self.query_monitors[user_id]:
+                self.query_monitors[user_id][mem_cube_id] = QueryMonitorQueue(
+                    maxsize=self.config.context_window_size
+                )
+
     def register_memory_manager_if_not_exists(
         self,
-        user_id: str,
-        mem_cube_id: str,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
         memory_monitors: dict[UserID, dict[MemCubeID, MemoryMonitorManager]],
         max_capacity: int,
     ) -> None:
         """
         Register a new MemoryMonitorManager for the given user and memory cube if it doesn't exist.
+        Thread-safe implementation using double-checked locking pattern.
 
         Checks if a MemoryMonitorManager already exists for the specified user_id and mem_cube_id.
         If not, creates a new MemoryMonitorManager with appropriate capacity settings and registers it.
@@ -65,14 +117,34 @@ class SchedulerMonitor(BaseSchedulerModule):
         Args:
             user_id: The ID of the user to associate with the memory manager
             mem_cube_id: The ID of the memory cube to monitor
+            memory_monitors: Dictionary storing existing memory monitor managers
+            max_capacity: Maximum capacity for the new memory monitor manager
+            lock: Threading lock to ensure safe concurrent access
 
         Note:
             This function will update the loose_max_working_memory_capacity based on the current
             WorkingMemory size plus partial retention number before creating a new manager.
         """
-        # Check if a MemoryMonitorManager already exists for the current user_id and mem_cube_id
-        # If doesn't exist, create and register a new one
-        if (user_id not in memory_monitors) or (mem_cube_id not in memory_monitors[user_id]):
+        # First check (lock-free, fast path)
+        # Quickly verify existence without lock overhead
+        if user_id in memory_monitors and mem_cube_id in memory_monitors[user_id]:
+            logger.info(
+                f"MemoryMonitorManager already exists for user_id={user_id}, "
+                f"mem_cube_id={mem_cube_id} in the provided memory_monitors dictionary"
+            )
+            return
+
+        # Second check (with lock, ensures uniqueness)
+        # Acquire lock before modification and verify again to prevent race conditions
+        with self._register_lock:
+            # Re-check after acquiring lock, as another thread might have created it
+            if user_id in memory_monitors and mem_cube_id in memory_monitors[user_id]:
+                logger.info(
+                    f"MemoryMonitorManager already exists for user_id={user_id}, "
+                    f"mem_cube_id={mem_cube_id} in the provided memory_monitors dictionary"
+                )
+                return
+
             # Initialize MemoryMonitorManager with user ID, memory cube ID, and max capacity
             monitor_manager = MemoryMonitorManager(
                 user_id=user_id, mem_cube_id=mem_cube_id, max_capacity=max_capacity
@@ -84,19 +156,16 @@ class SchedulerMonitor(BaseSchedulerModule):
                 f"Registered new MemoryMonitorManager for user_id={user_id},"
                 f" mem_cube_id={mem_cube_id} with max_capacity={max_capacity}"
             )
-        else:
-            logger.info(
-                f"MemoryMonitorManager already exists for user_id={user_id}, "
-                f"mem_cube_id={mem_cube_id} in the provided memory_monitors dictionary"
-            )
 
-    def update_memory_monitors(self, user_id: str, mem_cube_id: str, mem_cube: GeneralMemCube):
+    def update_working_memory_monitors(
+        self,
+        new_working_memory_monitors: list[MemoryMonitorItem],
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+    ):
         text_mem_base: TreeTextMemory = mem_cube.text_mem
-
-        if not isinstance(text_mem_base, TreeTextMemory):
-            logger.error("Not Implemented")
-            return
-
+        assert isinstance(text_mem_base, TreeTextMemory)
         self.working_mem_monitor_capacity = min(
             DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT,
             (
@@ -105,17 +174,6 @@ class SchedulerMonitor(BaseSchedulerModule):
             ),
         )
 
-        self.update_working_memory_monitors(
-            user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
-        )
-
-        self.update_activation_memory_monitors(
-            user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
-        )
-
-    def update_working_memory_monitors(
-        self, user_id: str, mem_cube_id: str, mem_cube: GeneralMemCube
-    ):
         # register monitors
         self.register_memory_manager_if_not_exists(
             user_id=user_id,
@@ -124,14 +182,8 @@ class SchedulerMonitor(BaseSchedulerModule):
             max_capacity=self.working_mem_monitor_capacity,
         )
 
-        # === update working memory monitors ===
-        # Retrieve current working memory content
-        text_mem_base: TreeTextMemory = mem_cube.text_mem
-        working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
-        text_working_memory: list[str] = [w_m.memory for w_m in working_memory]
-
         self.working_memory_monitors[user_id][mem_cube_id].update_memories(
-            text_working_memories=text_working_memory,
+            new_memory_monitors=new_working_memory_monitors,
             partial_retention_number=self.partial_retention_number,
         )
 
@@ -149,25 +201,22 @@ class SchedulerMonitor(BaseSchedulerModule):
         # Sort by importance_score in descending order and take top k
         top_k_memories = sorted(
             self.working_memory_monitors[user_id][mem_cube_id].memories,
-            key=lambda m: m.get_score(),
+            key=lambda m: m.get_importance_score(weight_vector=DEFAULT_WEIGHT_VECTOR_FOR_RANKING),
             reverse=True,
         )[: self.activation_mem_monitor_capacity]
 
-        # Extract just the text from these memories
-        text_top_k_memories = [m.memory_text for m in top_k_memories]
-
         # Update the activation memory monitors with these important memories
         self.activation_memory_monitors[user_id][mem_cube_id].update_memories(
-            text_working_memories=text_top_k_memories,
+            new_memory_monitors=top_k_memories,
             partial_retention_number=self.partial_retention_number,
         )
 
     def timed_trigger(self, last_time: datetime, interval_seconds: float) -> bool:
-        now = datetime.now()
+        now = datetime.utcnow()
         elapsed = (now - last_time).total_seconds()
         if elapsed >= interval_seconds:
             return True
-        logger.debug(f"Time trigger not ready, {elapsed:.1f}s elapsed (needs {interval_seconds}s)")
+        logger.info(f"Time trigger not ready, {elapsed:.1f}s elapsed (needs {interval_seconds}s)")
         return False
 
     def get_monitor_memories(
@@ -206,10 +255,10 @@ class SchedulerMonitor(BaseSchedulerModule):
             )
             return []
 
-        manager = monitor_dict[user_id][mem_cube_id]
+        manager: MemoryMonitorManager = monitor_dict[user_id][mem_cube_id]
         # Sort memories by recording_count in descending order and return top_k items
-        sorted_memories = sorted(manager.memories, key=lambda m: m.recording_count, reverse=True)
-        sorted_text_memories = [m.memory_text for m in sorted_memories[:top_k]]
+        sorted_memory_monitors = manager.get_sorted_mem_monitors(reverse=True)
+        sorted_text_memories = [m.memory_text for m in sorted_memory_monitors[:top_k]]
         return sorted_text_memories
 
     def get_monitors_info(self, user_id: str, mem_cube_id: str) -> dict[str, Any]:

memos/mem_scheduler/mos_for_test_scheduler.py

@@ -3,12 +3,12 @@ from datetime import datetime
 from memos.configs.mem_os import MOSConfig
 from memos.log import get_logger
 from memos.mem_os.main import MOS
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.schemas.general_schemas import (
     ANSWER_LABEL,
     MONITOR_WORKING_MEMORY_TYPE,
     QUERY_LABEL,
-    ScheduleMessageItem,
 )
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 logger = get_logger(__name__)
@@ -54,24 +54,20 @@ class MOSForTestScheduler(MOS):
                 if not mem_cube.text_mem:
                     continue
 
-                # submit message to scheduler
-                if self.enable_mem_scheduler and self.mem_scheduler is not None:
-                    message_item = ScheduleMessageItem(
-                        user_id=target_user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=mem_cube,
-                        label=QUERY_LABEL,
-                        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,
+                message_item = ScheduleMessageItem(
+                    user_id=target_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,
+                    mem_cube=mem_cube,
+                    label=QUERY_LABEL,
+                    content=query,
+                    timestamp=datetime.now(),
                 )
+                cur_working_memories = [m.memory for m in mem_cube.text_mem.get_working_memory()]
+                print(f"Working memories before schedule: {cur_working_memories}")
+
+                # --- 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,21 @@ class MOSForTestScheduler(MOS):
                     memory_type=MONITOR_WORKING_MEMORY_TYPE,
                     top_k=topk_for_scheduler,
                 )
+                print(f"Working memories after schedule: {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
+                    query,
+                    top_k=self.config.top_k - topk_for_scheduler,
+                    info={
+                        "user_id": target_user_id,
+                        "session_id": self.session_id,
+                        "chat_history": chat_history.chat_history,
+                    },
                 )
                 text_memories = [m.memory for m in memories]
+                print(f"Search results with new working memories: {text_memories}")
                 memories_all.extend(text_memories)
 
                 memories_all = list(set(memories_all))
@@ -139,5 +143,4 @@ class MOSForTestScheduler(MOS):
                     timestamp=datetime.now(),
                 )
                 self.mem_scheduler.submit_messages(messages=[message_item])
-
         return response