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

memos/mem_scheduler/general_scheduler.py

@@ -9,11 +9,13 @@ from memos.mem_scheduler.schemas.general_schemas import (
     ANSWER_LABEL,
     DEFAULT_MAX_QUERY_KEY_WORDS,
     QUERY_LABEL,
+    WORKING_MEMORY_TYPE,
     MemCubeID,
     UserID,
 )
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 from memos.mem_scheduler.schemas.monitor_schemas import QueryMonitorItem
+from memos.mem_scheduler.utils.filter_utils import is_all_chinese, is_all_english
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
 
 
@@ -35,11 +37,12 @@ class GeneralScheduler(BaseScheduler):
 
     # for evaluation
     def search_for_eval(
-        self,
-        query: str,
-        user_id: UserID | str,
-        top_k: int,
-    ) -> list[str]:
+        self, query: str, user_id: UserID | str, top_k: int, scheduler_flag: bool = True
+    ) -> (list[str], bool):
+        self.monitor.register_query_monitor_if_not_exists(
+            user_id=user_id, mem_cube_id=self.current_mem_cube_id
+        )
+
         query_keywords = self.monitor.extract_query_keywords(query=query)
         logger.info(f'Extract keywords "{query_keywords}" from query "{query}"')
 
@@ -48,35 +51,61 @@ class GeneralScheduler(BaseScheduler):
             keywords=query_keywords,
             max_keywords=DEFAULT_MAX_QUERY_KEY_WORDS,
         )
-        self.monitor.query_monitors.put(item=item)
-        logger.debug(
-            f"Queries in monitor are {self.monitor.query_monitors.get_queries_with_timesort()}."
-        )
+        query_monitor = self.monitor.query_monitors[user_id][self.current_mem_cube_id]
+        query_monitor.put(item=item)
+        logger.debug(f"Queries in monitor are {query_monitor.get_queries_with_timesort()}.")
 
         queries = [query]
 
         # recall
-        cur_working_memory, new_candidates = self.process_session_turn(
-            queries=queries,
-            user_id=user_id,
-            mem_cube_id=self.current_mem_cube_id,
-            mem_cube=self.current_mem_cube,
-            top_k=self.top_k,
-        )
-        logger.info(f"Processed {queries} and get {len(new_candidates)} new candidate memories.")
-
-        # rerank
-        new_order_working_memory = self.replace_working_memory(
-            user_id=user_id,
-            mem_cube_id=self.current_mem_cube_id,
-            mem_cube=self.current_mem_cube,
-            original_memory=cur_working_memory,
-            new_memory=new_candidates,
+        mem_cube = self.current_mem_cube
+        text_mem_base = mem_cube.text_mem
+
+        cur_working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
+        text_working_memory: list[str] = [w_m.memory for w_m in cur_working_memory]
+        intent_result = self.monitor.detect_intent(
+            q_list=queries, text_working_memory=text_working_memory
         )
-        new_order_working_memory = new_order_working_memory[:top_k]
-        logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
 
-        return [m.memory for m in new_order_working_memory]
+        if not scheduler_flag:
+            return text_working_memory, intent_result["trigger_retrieval"]
+        else:
+            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.info(f"missing_evidences: {item}")
+                    results: list[TextualMemoryItem] = self.retriever.search(
+                        query=item,
+                        mem_cube=mem_cube,
+                        top_k=k_per_evidence,
+                        method=self.search_method,
+                    )
+                    logger.info(
+                        f"search results for {missing_evidences}: {[one.memory for one in results]}"
+                    )
+                    new_candidates.extend(results)
+                print(
+                    f"missing_evidences: {missing_evidences} and get {len(new_candidates)} new candidate memories."
+                )
+            else:
+                new_candidates = []
+                print(f"intent_result: {intent_result}. not triggered")
+
+            # rerank
+            new_order_working_memory = self.replace_working_memory(
+                user_id=user_id,
+                mem_cube_id=self.current_mem_cube_id,
+                mem_cube=self.current_mem_cube,
+                original_memory=cur_working_memory,
+                new_memory=new_candidates,
+            )
+            new_order_working_memory = new_order_working_memory[:top_k]
+            logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
+
+            return [m.memory for m in new_order_working_memory], intent_result["trigger_retrieval"]
 
     def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
@@ -105,18 +134,42 @@ class GeneralScheduler(BaseScheduler):
 
                 # update query monitors
                 for msg in messages:
+                    self.monitor.register_query_monitor_if_not_exists(
+                        user_id=user_id, mem_cube_id=mem_cube_id
+                    )
+
                     query = msg.content
                     query_keywords = self.monitor.extract_query_keywords(query=query)
                     logger.info(f'Extract keywords "{query_keywords}" from query "{query}"')
 
+                    if len(query_keywords) == 0:
+                        stripped_query = query.strip()
+                        # Determine measurement method based on language
+                        if is_all_english(stripped_query):
+                            words = stripped_query.split()  # Word count for English
+                        elif is_all_chinese(stripped_query):
+                            words = stripped_query  # Character count for Chinese
+                        else:
+                            logger.debug(
+                                f"Mixed-language memory, using character count: {stripped_query[:50]}..."
+                            )
+                            words = stripped_query  # Default to character count
+
+                        query_keywords = list(set(words[:20]))
+                        logger.error(
+                            f"Keyword extraction failed for query. Using fallback keywords: {query_keywords[:10]}... (truncated)"
+                        )
+
                     item = QueryMonitorItem(
                         query_text=query,
                         keywords=query_keywords,
                         max_keywords=DEFAULT_MAX_QUERY_KEY_WORDS,
                     )
-                    self.monitor.query_monitors.put(item=item)
+
+                    self.monitor.query_monitors[user_id][mem_cube_id].put(item=item)
                 logger.debug(
-                    f"Queries in monitor are {self.monitor.query_monitors.get_queries_with_timesort()}."
+                    f"Queries in monitor are "
+                    f"{self.monitor.query_monitors[user_id][mem_cube_id].get_queries_with_timesort()}."
                 )
 
                 queries = [msg.content for msg in messages]
@@ -143,6 +196,20 @@ class GeneralScheduler(BaseScheduler):
                 )
                 logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
 
+                # update activation memories
+                logger.info(
+                    f"Activation memory update {'enabled' if self.enable_activation_memory else 'disabled'} "
+                    f"(interval: {self.monitor.act_mem_update_interval}s)"
+                )
+                if self.enable_activation_memory:
+                    self.update_activation_memory_periodically(
+                        interval_seconds=self.monitor.act_mem_update_interval,
+                        label=QUERY_LABEL,
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=messages[0].mem_cube,
+                    )
+
     def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
         Process and handle answer trigger messages from the queue.
@@ -165,26 +232,6 @@ class GeneralScheduler(BaseScheduler):
                 # for status update
                 self._set_current_context_from_message(msg=messages[0])
 
-                # update activation memories
-                if self.enable_act_memory_update:
-                    if (
-                        len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories)
-                        == 0
-                    ):
-                        self.initialize_working_memory_monitors(
-                            user_id=user_id,
-                            mem_cube_id=mem_cube_id,
-                            mem_cube=messages[0].mem_cube,
-                        )
-
-                    self.update_activation_memory_periodically(
-                        interval_seconds=self.monitor.act_mem_update_interval,
-                        label=ANSWER_LABEL,
-                        user_id=user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=messages[0].mem_cube,
-                    )
-
     def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         logger.info(f"Messages {messages} assigned to {ADD_LABEL} handler.")
         # Process the query in a session turn
@@ -215,6 +262,9 @@ class GeneralScheduler(BaseScheduler):
                             mem_type = mem_item.metadata.memory_type
                             mem_content = mem_item.memory
 
+                            if mem_type == WORKING_MEMORY_TYPE:
+                                continue
+
                             self.log_adding_memory(
                                 memory=mem_content,
                                 memory_type=mem_type,
@@ -224,15 +274,6 @@ class GeneralScheduler(BaseScheduler):
                                 log_func_callback=self._submit_web_logs,
                             )
 
-                    # update activation memories
-                    if self.enable_act_memory_update:
-                        self.update_activation_memory_periodically(
-                            interval_seconds=self.monitor.act_mem_update_interval,
-                            label=ADD_LABEL,
-                            user_id=user_id,
-                            mem_cube_id=mem_cube_id,
-                            mem_cube=messages[0].mem_cube,
-                        )
         except Exception as e:
             logger.error(f"Error: {e}", exc_info=True)
 
@@ -289,18 +330,20 @@ class GeneralScheduler(BaseScheduler):
         new_candidates = []
         for item in missing_evidences:
             logger.info(f"missing_evidences: {item}")
+            info = {
+                "user_id": user_id,
+                "session_id": "",
+            }
+
             results: list[TextualMemoryItem] = self.retriever.search(
-                query=item, mem_cube=mem_cube, top_k=k_per_evidence, method=self.search_method
+                query=item,
+                mem_cube=mem_cube,
+                top_k=k_per_evidence,
+                method=self.search_method,
+                info=info,
             )
             logger.info(
                 f"search results for {missing_evidences}: {[one.memory for one in results]}"
             )
             new_candidates.extend(results)
-
-        if len(new_candidates) == 0:
-            logger.warning(
-                f"As new_candidates is empty, new_candidates is set same to working_memory.\n"
-                f"time_trigger_flag: {time_trigger_flag}; intent_result: {intent_result}"
-            )
-            new_candidates = cur_working_memory
         return cur_working_memory, new_candidates

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()