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

memos/mem_scheduler/base_scheduler.py

@@ -9,13 +9,14 @@ from memos.configs.mem_scheduler import AuthConfig, 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.dispatcher import SchedulerDispatcher
-from memos.mem_scheduler.modules.misc import AutoDroppingQueue as Queue
-from memos.mem_scheduler.modules.monitor import SchedulerMonitor
-from memos.mem_scheduler.modules.rabbitmq_service import RabbitMQSchedulerModule
-from memos.mem_scheduler.modules.redis_service import RedisSchedulerModule
-from memos.mem_scheduler.modules.retriever import SchedulerRetriever
-from memos.mem_scheduler.modules.scheduler_logger import SchedulerLoggerModule
+from memos.mem_scheduler.general_modules.dispatcher import SchedulerDispatcher
+from memos.mem_scheduler.general_modules.misc import AutoDroppingQueue as Queue
+from memos.mem_scheduler.general_modules.rabbitmq_service import RabbitMQSchedulerModule
+from memos.mem_scheduler.general_modules.redis_service import RedisSchedulerModule
+from memos.mem_scheduler.general_modules.retriever import SchedulerRetriever
+from memos.mem_scheduler.general_modules.scheduler_logger import SchedulerLoggerModule
+from memos.mem_scheduler.monitors.dispatcher_monitor import SchedulerDispatcherMonitor
+from memos.mem_scheduler.monitors.general_monitor import SchedulerGeneralMonitor
 from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_ACT_MEM_DUMP_PATH,
     DEFAULT_CONSUME_INTERVAL_SECONDS,
@@ -52,19 +53,20 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         # hyper-parameters
         self.top_k = self.config.get("top_k", 10)
         self.context_window_size = self.config.get("context_window_size", 5)
-        self.enable_act_memory_update = self.config.get("enable_act_memory_update", False)
+        self.enable_activation_memory = self.config.get("enable_activation_memory", False)
         self.act_mem_dump_path = self.config.get("act_mem_dump_path", DEFAULT_ACT_MEM_DUMP_PATH)
         self.search_method = TreeTextMemory_SEARCH_METHOD
         self.enable_parallel_dispatch = self.config.get("enable_parallel_dispatch", False)
-        self.max_workers = self.config.get(
+        self.thread_pool_max_workers = self.config.get(
             "thread_pool_max_workers", DEFAULT_THREAD__POOL_MAX_WORKERS
         )
 
         self.retriever: SchedulerRetriever | None = None
-        self.monitor: SchedulerMonitor | None = None
-
+        self.monitor: SchedulerGeneralMonitor | None = None
+        self.dispatcher_monitor: SchedulerDispatcherMonitor | None = None
         self.dispatcher = SchedulerDispatcher(
-            max_workers=self.max_workers, enable_parallel_dispatch=self.enable_parallel_dispatch
+            max_workers=self.thread_pool_max_workers,
+            enable_parallel_dispatch=self.enable_parallel_dispatch,
         )
 
         # internal message queue
@@ -84,6 +86,8 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         # other attributes
         self._context_lock = threading.Lock()
         self.current_user_id: UserID | str | None = None
+        self.current_mem_cube_id: MemCubeID | str | None = None
+        self.current_mem_cube: GeneralMemCube | None = None
         self.auth_config_path: str | Path | None = self.config.get("auth_config_path", None)
         self.auth_config = None
         self.rabbitmq_config = None
@@ -95,14 +99,19 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         # initialize submodules
         self.chat_llm = chat_llm
         self.process_llm = process_llm
-        self.monitor = SchedulerMonitor(process_llm=self.process_llm, config=self.config)
+        self.monitor = SchedulerGeneralMonitor(process_llm=self.process_llm, config=self.config)
+        self.dispatcher_monitor = SchedulerDispatcherMonitor(config=self.config)
         self.retriever = SchedulerRetriever(process_llm=self.process_llm, config=self.config)
 
+        if self.enable_parallel_dispatch:
+            self.dispatcher_monitor.initialize(dispatcher=self.dispatcher)
+            self.dispatcher_monitor.start()
+
         # initialize with auth_cofig
         if self.auth_config_path is not None and Path(self.auth_config_path).exists():
-            self.auth_config = AuthConfig.from_local_yaml(config_path=self.auth_config_path)
+            self.auth_config = AuthConfig.from_local_config(config_path=self.auth_config_path)
         elif AuthConfig.default_config_exists():
-            self.auth_config = AuthConfig.from_local_yaml()
+            self.auth_config = AuthConfig.from_local_config()
         else:
             self.auth_config = None
 
@@ -130,8 +139,8 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             self.current_mem_cube_id = msg.mem_cube_id
             self.current_mem_cube = msg.mem_cube
 
-    def transform_memories_to_monitors(
-        self, memories: list[TextualMemoryItem]
+    def transform_working_memories_to_monitors(
+        self, query_keywords, memories: list[TextualMemoryItem]
     ) -> list[MemoryMonitorItem]:
         """
         Convert a list of TextualMemoryItem objects into MemoryMonitorItem objects
@@ -143,10 +152,6 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         Returns:
             List of MemoryMonitorItem objects with computed importance scores.
         """
-        query_keywords = self.monitor.query_monitors.get_keywords_collections()
-        logger.debug(
-            f"Processing {len(memories)} memories with {len(query_keywords)} query keywords"
-        )
 
         result = []
         mem_length = len(memories)
@@ -178,7 +183,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             )
             result.append(mem_monitor)
 
-        logger.debug(f"Transformed {len(result)} memories to monitors")
+        logger.info(f"Transformed {len(result)} memories to monitors")
         return result
 
     def replace_working_memory(
@@ -195,7 +200,8 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             text_mem_base: TreeTextMemory = text_mem_base
 
             # process rerank memories with llm
-            query_history = self.monitor.query_monitors.get_queries_with_timesort()
+            query_monitor = self.monitor.query_monitors[user_id][mem_cube_id]
+            query_history = query_monitor.get_queries_with_timesort()
             memories_with_new_order, rerank_success_flag = (
                 self.retriever.process_and_rerank_memories(
                     queries=query_history,
@@ -206,14 +212,20 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             )
 
             # update working memory monitors
-            new_working_memory_monitors = self.transform_memories_to_monitors(
-                memories=memories_with_new_order
+            query_keywords = query_monitor.get_keywords_collections()
+            logger.info(
+                f"Processing {len(memories_with_new_order)} memories with {len(query_keywords)} query keywords"
+            )
+            new_working_memory_monitors = self.transform_working_memories_to_monitors(
+                query_keywords=query_keywords,
+                memories=memories_with_new_order,
             )
 
             if not rerank_success_flag:
                 for one in new_working_memory_monitors:
                     one.sorting_score = 0
 
+            logger.info(f"update {len(new_working_memory_monitors)} working_memory_monitors")
             self.monitor.update_working_memory_monitors(
                 new_working_memory_monitors=new_working_memory_monitors,
                 user_id=user_id,
@@ -245,25 +257,6 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
 
         return memories_with_new_order
 
-    def initialize_working_memory_monitors(
-        self,
-        user_id: UserID | str,
-        mem_cube_id: MemCubeID | str,
-        mem_cube: GeneralMemCube,
-    ):
-        text_mem_base: TreeTextMemory = mem_cube.text_mem
-        working_memories = text_mem_base.get_working_memory()
-
-        working_memory_monitors = self.transform_memories_to_monitors(
-            memories=working_memories,
-        )
-        self.monitor.update_working_memory_monitors(
-            new_working_memory_monitors=working_memory_monitors,
-            user_id=user_id,
-            mem_cube_id=mem_cube_id,
-            mem_cube=mem_cube,
-        )
-
     def update_activation_memory(
         self,
         new_memories: list[str | TextualMemoryItem],
@@ -324,6 +317,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
 
             cache_item = act_mem.extract(new_text_memory)
             cache_item.records.text_memories = new_text_memories
+            cache_item.records.timestamp = datetime.utcnow()
 
             act_mem.add([cache_item])
             act_mem.dump(self.act_mem_dump_path)
@@ -367,13 +361,9 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                     or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0
                 ):
                     logger.warning(
-                        "No memories found in working_memory_monitors, initializing from current working_memories"
-                    )
-                    self.initialize_working_memory_monitors(
-                        user_id=user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=mem_cube,
+                        "No memories found in working_memory_monitors, activation memory update is skipped"
                     )
+                    return
 
                 self.monitor.update_activation_memory_monitors(
                     user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
@@ -387,6 +377,11 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                 logger.info(
                     f"Collected {len(new_activation_memories)} new memory entries for processing"
                 )
+                # Print the content of each new activation memory
+                for i, memory in enumerate(new_activation_memories[:5], 1):
+                    logger.info(
+                        f"Part of New Activation Memorires | {i}/{len(new_activation_memories)}: {memory[:20]}"
+                    )
 
                 self.update_activation_memory(
                     new_memories=new_activation_memories,
@@ -396,19 +391,20 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                     mem_cube=mem_cube,
                 )
 
-                self.monitor.last_activation_mem_update_time = datetime.now()
+                self.monitor.last_activation_mem_update_time = datetime.utcnow()
 
                 logger.debug(
                     f"Activation memory update completed at {self.monitor.last_activation_mem_update_time}"
                 )
+
             else:
                 logger.info(
                     f"Skipping update - {interval_seconds} second interval not yet reached. "
                     f"Last update time is {self.monitor.last_activation_mem_update_time} and now is"
-                    f"{datetime.now()}"
+                    f"{datetime.utcnow()}"
                 )
         except Exception as e:
-            logger.error(f"Error: {e}", exc_info=True)
+            logger.error(f"Error in update_activation_memory_periodically: {e}", exc_info=True)
 
     def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
         """Submit multiple messages to the message queue."""
@@ -506,7 +502,9 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
 
         # Initialize dispatcher resources
         if self.enable_parallel_dispatch:
-            logger.info(f"Initializing dispatcher thread pool with {self.max_workers} workers")
+            logger.info(
+                f"Initializing dispatcher thread pool with {self.thread_pool_max_workers} workers"
+            )
 
         # Start consumer thread
         self._running = True
@@ -541,10 +539,15 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                 logger.info("Consumer thread stopped")
 
         # Shutdown dispatcher
-        if hasattr(self, "dispatcher") and self.dispatcher:
+        if self.dispatcher:
             logger.info("Shutting down dispatcher...")
             self.dispatcher.shutdown()
 
+        # Shutdown dispatcher_monitor
+        if self.dispatcher_monitor:
+            logger.info("Shutting down monitor...")
+            self.dispatcher_monitor.stop()
+
         # Clean up queues
         self._cleanup_queues()
         logger.info("Memory Scheduler stopped completely")

memos/mem_scheduler/{modules → general_modules}/base.py

@@ -17,8 +17,7 @@ class BaseSchedulerModule:
 
         self._chat_llm = None
         self._process_llm = None
-        self.current_mem_cube_id: str | None = None
-        self.current_mem_cube: GeneralMemCube | None = None
+
         self.mem_cubes: dict[str, GeneralMemCube] = {}
 
     def load_template(self, template_name: str) -> str:

memos/mem_scheduler/{modules → general_modules}/dispatcher.py

@@ -1,9 +1,11 @@
+import concurrent
+
 from collections import defaultdict
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 
 from memos.log import get_logger
-from memos.mem_scheduler.modules.base import BaseSchedulerModule
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
@@ -26,20 +28,27 @@ class SchedulerDispatcher(BaseSchedulerModule):
         super().__init__()
         # Main dispatcher thread pool
         self.max_workers = max_workers
+
         # Only initialize thread pool if in parallel mode
         self.enable_parallel_dispatch = enable_parallel_dispatch
+        self.thread_name_prefix = "dispatcher"
         if self.enable_parallel_dispatch:
             self.dispatcher_executor = ThreadPoolExecutor(
-                max_workers=self.max_workers, thread_name_prefix="dispatcher"
+                max_workers=self.max_workers, thread_name_prefix=self.thread_name_prefix
             )
         else:
             self.dispatcher_executor = None
         logger.info(f"enable_parallel_dispatch is set to {self.enable_parallel_dispatch}")
+
         # Registered message handlers
         self.handlers: dict[str, Callable] = {}
+
         # Dispatcher running state
         self._running = False
 
+        # Set to track active futures for monitoring purposes
+        self._futures = set()
+
     def register_handler(self, label: str, handler: Callable[[list[ScheduleMessageItem]], None]):
         """
         Register a handler function for a specific message label.
@@ -105,6 +114,13 @@ class SchedulerDispatcher(BaseSchedulerModule):
         # Convert defaultdict to regular dict for cleaner output
         return {user_id: dict(cube_groups) for user_id, cube_groups in grouped_dict.items()}
 
+    def _handle_future_result(self, future):
+        self._futures.remove(future)
+        try:
+            future.result()  # this will throw exception
+        except Exception as e:
+            logger.error(f"Handler execution failed: {e!s}", exc_info=True)
+
     def dispatch(self, msg_list: list[ScheduleMessageItem]):
         """
         Dispatch a list of messages to their respective handlers.
@@ -112,26 +128,26 @@ class SchedulerDispatcher(BaseSchedulerModule):
         Args:
             msg_list: List of ScheduleMessageItem objects to process
         """
+        if not msg_list:
+            logger.debug("Received empty message list, skipping dispatch")
+            return
 
-        # Group messages by their labels
+        # Group messages by their labels, and organize messages by label
         label_groups = defaultdict(list)
-
-        # Organize messages by label
         for message in msg_list:
             label_groups[message.label].append(message)
 
         # Process each label group
         for label, msgs in label_groups.items():
-            if label not in self.handlers:
-                logger.error(f"No handler registered for label: {label}")
-                handler = self._default_message_handler
-            else:
-                handler = self.handlers[label]
+            handler = self.handlers.get(label, self._default_message_handler)
+
             # dispatch to different 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
-                self.dispatcher_executor.submit(handler, msgs)
+                future = self.dispatcher_executor.submit(handler, msgs)
+                self._futures.add(future)
+                future.add_done_callback(self._handle_future_result)
                 logger.info(f"Dispatched {len(msgs)} message(s) as future task")
             else:
                 handler(msgs)
@@ -148,15 +164,38 @@ class SchedulerDispatcher(BaseSchedulerModule):
         if not self.enable_parallel_dispatch or self.dispatcher_executor is None:
             return True  # 串行模式无需等待
 
-        self.dispatcher_executor.shutdown(wait=True, timeout=timeout)
-        return True
+        done, not_done = concurrent.futures.wait(
+            self._futures, timeout=timeout, return_when=concurrent.futures.ALL_COMPLETED
+        )
+
+        # Check for exceptions in completed tasks
+        for future in done:
+            try:
+                future.result()
+            except Exception:
+                logger.error("Handler failed during shutdown", exc_info=True)
+
+        return len(not_done) == 0
 
     def shutdown(self) -> None:
         """Gracefully shutdown the dispatcher."""
+        self._running = False
+
         if self.dispatcher_executor is not None:
+            # Cancel pending tasks
+            cancelled = 0
+            for future in self._futures:
+                if future.cancel():
+                    cancelled += 1
+            logger.info(f"Cancelled {cancelled}/{len(self._futures)} pending tasks")
+
+        # Shutdown executor
+        try:
             self.dispatcher_executor.shutdown(wait=True)
-        self._running = False
-        logger.info("Dispatcher has been shutdown.")
+        except Exception as e:
+            logger.error(f"Executor shutdown error: {e}", exc_info=True)
+        finally:
+            self._futures.clear()
 
     def __enter__(self):
         self._running = True

memos/mem_scheduler/{modules → general_modules}/rabbitmq_service.py

@@ -8,8 +8,8 @@ from pathlib import Path
 from memos.configs.mem_scheduler import AuthConfig, RabbitMQConfig
 from memos.dependency import require_python_package
 from memos.log import get_logger
-from memos.mem_scheduler.modules.base import BaseSchedulerModule
-from memos.mem_scheduler.modules.misc import AutoDroppingQueue
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
+from memos.mem_scheduler.general_modules.misc import AutoDroppingQueue
 from memos.mem_scheduler.schemas.general_schemas import DIRECT_EXCHANGE_TYPE, FANOUT_EXCHANGE_TYPE
 
 
@@ -71,9 +71,9 @@ class RabbitMQSchedulerModule(BaseSchedulerModule):
 
         if config is None:
             if config_path is None and AuthConfig.default_config_exists():
-                auth_config = AuthConfig.from_local_yaml()
+                auth_config = AuthConfig.from_local_config()
             elif Path(config_path).exists():
-                auth_config = AuthConfig.from_local_yaml(config_path=config_path)
+                auth_config = AuthConfig.from_local_config(config_path=config_path)
             else:
                 logger.error("Fail to initialize auth_config")
                 return

memos/mem_scheduler/{modules → general_modules}/redis_service.py

@@ -6,7 +6,7 @@ from typing import Any
 
 from memos.dependency import require_python_package
 from memos.log import get_logger
-from memos.mem_scheduler.modules.base import BaseSchedulerModule
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 
 
 logger = get_logger(__name__)

memos/mem_scheduler/{modules → general_modules}/retriever.py

@@ -2,8 +2,9 @@ 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.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.schemas.general_schemas import (
+    TreeTextMemory_FINE_SEARCH_METHOD,
     TreeTextMemory_SEARCH_METHOD,
 )
 from memos.mem_scheduler.utils.filter_utils import (
@@ -32,7 +33,12 @@ class SchedulerRetriever(BaseSchedulerModule):
         self.process_llm = process_llm
 
     def search(
-        self, query: str, mem_cube: GeneralMemCube, top_k: int, method=TreeTextMemory_SEARCH_METHOD
+        self,
+        query: str,
+        mem_cube: GeneralMemCube,
+        top_k: int,
+        method: str = TreeTextMemory_SEARCH_METHOD,
+        info: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """Search in text memory with the given query.
 
@@ -46,13 +52,21 @@ class SchedulerRetriever(BaseSchedulerModule):
         """
         text_mem_base = mem_cube.text_mem
         try:
-            if method == TreeTextMemory_SEARCH_METHOD:
+            if method in [TreeTextMemory_SEARCH_METHOD, TreeTextMemory_FINE_SEARCH_METHOD]:
                 assert isinstance(text_mem_base, TreeTextMemory)
+                if info is None:
+                    logger.warning(
+                        "Please input 'info' when use tree.search so that "
+                        "the database would store the consume history."
+                    )
+                    info = {"user_id": "", "session_id": ""}
+
+                mode = "fast" if method == TreeTextMemory_SEARCH_METHOD else "fine"
                 results_long_term = text_mem_base.search(
-                    query=query, top_k=top_k, memory_type="LongTermMemory"
+                    query=query, top_k=top_k, memory_type="LongTermMemory", mode=mode, info=info
                 )
                 results_user = text_mem_base.search(
-                    query=query, top_k=top_k, memory_type="UserMemory"
+                    query=query, top_k=top_k, memory_type="UserMemory", mode=mode, info=info
                 )
                 results = results_long_term + results_user
             else:

memos/mem_scheduler/{modules → general_modules}/scheduler_logger.py

@@ -2,7 +2,7 @@ from collections.abc import Callable
 
 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.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.schemas.general_schemas import (
     ACTIVATION_MEMORY_TYPE,
     ADD_LABEL,
@@ -180,6 +180,11 @@ class SchedulerLoggerModule(BaseSchedulerModule):
                 mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
             )
+            logger.info(
+                f"{len(added_memories)} {TEXT_MEMORY_TYPE} memorie(s) "
+                f"transformed to {ACTIVATION_MEMORY_TYPE} memories."
+            )
+
             log_message_b = self.create_autofilled_log_item(
                 log_content=mem,
                 label=label,
@@ -189,12 +194,13 @@ class SchedulerLoggerModule(BaseSchedulerModule):
                 mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
             )
-            log_func_callback([log_message_a, log_message_b])
             logger.info(
-                f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
-                f"transformed to {WORKING_MEMORY_TYPE} memories."
+                f"{len(added_memories)} {ACTIVATION_MEMORY_TYPE} memorie(s) "
+                f"transformed to {PARAMETER_MEMORY_TYPE} memories."
             )
 
+            log_func_callback([log_message_a, log_message_b])
+
     @log_exceptions(logger=logger)
     def log_adding_memory(
         self,