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

memos/mem_scheduler/modules/dispatcher.py

@@ -4,7 +4,7 @@ from concurrent.futures import ThreadPoolExecutor
 
 from memos.log import get_logger
 from memos.mem_scheduler.modules.base import BaseSchedulerModule
-from memos.mem_scheduler.modules.schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 logger = get_logger(__name__)
@@ -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
@@ -73,6 +73,38 @@ class SchedulerDispatcher(BaseSchedulerModule):
     def _default_message_handler(self, messages: list[ScheduleMessageItem]) -> None:
         logger.debug(f"Using _default_message_handler to deal with messages: {messages}")
 
+    def group_messages_by_user_and_cube(
+        self, messages: list[ScheduleMessageItem]
+    ) -> dict[str, dict[str, list[ScheduleMessageItem]]]:
+        """
+        Groups messages into a nested dictionary structure first by user_id, then by mem_cube_id.
+
+        Args:
+            messages: List of ScheduleMessageItem objects to be grouped
+
+        Returns:
+            A nested dictionary with the structure:
+            {
+                "user_id_1": {
+                    "mem_cube_id_1": [msg1, msg2, ...],
+                    "mem_cube_id_2": [msg3, msg4, ...],
+                    ...
+                },
+                "user_id_2": {
+                    ...
+                },
+                ...
+            }
+            Where each msg is the original ScheduleMessageItem object
+        """
+        grouped_dict = defaultdict(lambda: defaultdict(list))
+
+        for msg in messages:
+            grouped_dict[msg.user_id][msg.mem_cube_id].append(msg)
+
+        # Convert defaultdict to regular dict for cleaner output
+        return {user_id: dict(cube_groups) for user_id, cube_groups in grouped_dict.items()}
+
     def dispatch(self, msg_list: list[ScheduleMessageItem]):
         """
         Dispatch a list of messages to their respective handlers.
@@ -96,8 +128,39 @@ class SchedulerDispatcher(BaseSchedulerModule):
             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
                 self.dispatcher_executor.submit(handler, msgs)
+                logger.info(f"Dispatched {len(msgs)} message(s) as future task")
             else:
-                handler(msgs)  # Direct serial execution
+                handler(msgs)
+
+    def join(self, timeout: float | None = None) -> bool:
+        """Wait for all dispatched tasks to complete.
+
+        Args:
+            timeout: Maximum time to wait in seconds. None means wait forever.
+
+        Returns:
+            bool: True if all tasks completed, False if timeout occurred.
+        """
+        if not self.enable_parallel_dispatch or self.dispatcher_executor is None:
+            return True  # 串行模式无需等待
+
+        self.dispatcher_executor.shutdown(wait=True, timeout=timeout)
+        return True
+
+    def shutdown(self) -> None:
+        """Gracefully shutdown the dispatcher."""
+        if self.dispatcher_executor is not None:
+            self.dispatcher_executor.shutdown(wait=True)
+        self._running = False
+        logger.info("Dispatcher has been shutdown.")
+
+    def __enter__(self):
+        self._running = True
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()

memos/mem_scheduler/modules/misc.py

@@ -0,0 +1,104 @@
+import json
+
+from contextlib import suppress
+from datetime import datetime
+from queue import Empty, Full, Queue
+from typing import TYPE_CHECKING, TypeVar
+
+from pydantic import field_serializer
+
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+T = TypeVar("T")
+
+BaseModelType = TypeVar("T", bound="BaseModel")
+
+
+class DictConversionMixin:
+    """
+    Provides conversion functionality between Pydantic models and dictionaries,
+    including datetime serialization handling.
+    """
+
+    @field_serializer("timestamp", check_fields=False)
+    def serialize_datetime(self, dt: datetime | None, _info) -> str | None:
+        """
+        Custom datetime serialization logic.
+        - Supports timezone-aware datetime objects
+        - Compatible with models without timestamp field (via check_fields=False)
+        """
+        if dt is None:
+            return None
+        return dt.isoformat()
+
+    def to_dict(self) -> dict:
+        """
+        Convert model instance to dictionary.
+        - Uses model_dump to ensure field consistency
+        - Prioritizes custom serializer for timestamp handling
+        """
+        dump_data = self.model_dump()
+        if hasattr(self, "timestamp") and self.timestamp is not None:
+            dump_data["timestamp"] = self.serialize_datetime(self.timestamp, None)
+        return dump_data
+
+    @classmethod
+    def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
+        """
+        Create model instance from dictionary.
+        - Automatically converts timestamp strings to datetime objects
+        """
+        data_copy = data.copy()  # Avoid modifying original dictionary
+        if "timestamp" in data_copy and isinstance(data_copy["timestamp"], str):
+            try:
+                data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])
+            except ValueError:
+                # Handle invalid time formats - adjust as needed (e.g., log warning or set to None)
+                data_copy["timestamp"] = None
+
+        return cls(**data_copy)
+
+    def __str__(self) -> str:
+        """
+        Convert to formatted JSON string.
+        - Used for user-friendly display in print() or str() calls
+        """
+        return json.dumps(
+            self.to_dict(),
+            indent=4,
+            ensure_ascii=False,
+            default=lambda o: str(o),  # Handle other non-serializable objects
+        )
+
+
+class AutoDroppingQueue(Queue[T]):
+    """A thread-safe queue that automatically drops the oldest item when full."""
+
+    def __init__(self, maxsize: int = 0):
+        super().__init__(maxsize=maxsize)
+
+    def put(self, item: T, block: bool = False, timeout: float | None = None) -> None:
+        """Put an item into the queue.
+
+        If the queue is full, the oldest item will be automatically removed to make space.
+        This operation is thread-safe.
+
+        Args:
+            item: The item to be put into the queue
+            block: Ignored (kept for compatibility with Queue interface)
+            timeout: Ignored (kept for compatibility with Queue interface)
+        """
+        try:
+            # First try non-blocking put
+            super().put(item, block=block, timeout=timeout)
+        except Full:
+            with suppress(Empty):
+                self.get_nowait()  # Remove oldest item
+            # Retry putting the new item
+            super().put(item, block=block, timeout=timeout)
+
+    def get_queue_content_without_pop(self) -> list[T]:
+        """Return a copy of the queue's contents without modifying it."""
+        return list(self.queue)

memos/mem_scheduler/modules/monitor.py

@@ -1,11 +1,27 @@
-import json
-
+from datetime import datetime
 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.utils import extract_json_dict
+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,
+    UserID,
+)
+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
 
 
@@ -13,33 +29,224 @@ logger = get_logger(__name__)
 
 
 class SchedulerMonitor(BaseSchedulerModule):
-    def __init__(self, chat_llm, activation_mem_size=5):
+    """Monitors and manages scheduling operations with LLM integration."""
+
+    def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
         super().__init__()
-        self.statistics = {}
-        self.intent_history: list[str] = []
-        self.activation_mem_size = activation_mem_size
-        self.activation_memory_freq_list = [
-            {"memory": None, "count": 0} for _ in range(self.activation_mem_size)
-        ]
-
-        self._chat_llm = chat_llm
-
-    def update_stats(self, mem_cube):
-        self.statistics["activation_mem_size"] = self.activation_mem_size
-        mem_cube_info = self.get_mem_cube_info(mem_cube)
-        self.statistics.update(mem_cube_info)
-
-    def get_mem_cube_info(self, mem_cube: GeneralMemCube):
-        mem_cube_info = {}
-
-        text_mem = mem_cube.text_mem
-        if isinstance(text_mem, TreeTextMemory):
-            memory_size_dict = text_mem.memory_manager.memory_size
-            mem_cube_info["text_mem"] = memory_size_dict
+
+        # hyper-parameters
+        self.config: BaseSchedulerConfig = config
+        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
+
+        # attributes
+        # recording query_messages
+        self.query_monitors: QueryMonitorQueue[QueryMonitorItem] = QueryMonitorQueue(
+            maxsize=self.config.context_window_size
+        )
+
+        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_query_consume_time = datetime.min
+
+        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_memory_manager_if_not_exists(
+        self,
+        user_id: str,
+        mem_cube_id: 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.
+
+        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.
+
+        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
+
+        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]):
+            # 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
+            )
+
+            # Safely register the new manager in the nested dictionary structure
+            memory_monitors.setdefault(user_id, {})[mem_cube_id] = monitor_manager
+            logger.info(
+                f"Registered new MemoryMonitorManager for user_id={user_id},"
+                f" mem_cube_id={mem_cube_id} with max_capacity={max_capacity}"
+            )
         else:
-            logger.error("Not Implemented")
+            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_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
+        assert isinstance(text_mem_base, TreeTextMemory)
+        self.working_mem_monitor_capacity = min(
+            DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT,
+            (
+                text_mem_base.memory_manager.memory_size["WorkingMemory"]
+                + self.partial_retention_number
+            ),
+        )
 
-        return mem_cube_info
+        # register monitors
+        self.register_memory_manager_if_not_exists(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            memory_monitors=self.working_memory_monitors,
+            max_capacity=self.working_mem_monitor_capacity,
+        )
+
+        self.working_memory_monitors[user_id][mem_cube_id].update_memories(
+            new_memory_monitors=new_working_memory_monitors,
+            partial_retention_number=self.partial_retention_number,
+        )
+
+    def update_activation_memory_monitors(
+        self, user_id: str, mem_cube_id: str, mem_cube: GeneralMemCube
+    ):
+        self.register_memory_manager_if_not_exists(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            memory_monitors=self.activation_memory_monitors,
+            max_capacity=self.activation_mem_monitor_capacity,
+        )
+
+        # === update activation memory monitors ===
+        # 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_importance_score(weight_vector=DEFAULT_WEIGHT_VECTOR_FOR_RANKING),
+            reverse=True,
+        )[: self.activation_mem_monitor_capacity]
+
+        # Update the activation memory monitors with these important memories
+        self.activation_memory_monitors[user_id][mem_cube_id].update_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()
+        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)")
+        return False
+
+    def get_monitor_memories(
+        self,
+        user_id: str,
+        mem_cube_id: str,
+        memory_type: str = MONITOR_WORKING_MEMORY_TYPE,
+        top_k: int = 10,
+    ) -> list[str]:
+        """Retrieves memory items managed by the scheduler, sorted by recording count.
+
+        Args:
+            user_id: Unique identifier of the user
+            mem_cube_id: Unique identifier of the memory cube
+            memory_type: Type of memory to retrieve (MONITOR_WORKING_MEMORY_TYPE or
+                       MONITOR_ACTIVATION_MEMORY_TYPE)
+            top_k: Maximum number of memory items to return (default: 10)
+
+        Returns:
+            List of memory texts, sorted by recording count in descending order.
+            Returns empty list if no MemoryMonitorManager exists for the given parameters.
+        """
+        # Select the appropriate monitor dictionary based on memory_type
+        if memory_type == MONITOR_WORKING_MEMORY_TYPE:
+            monitor_dict = self.working_memory_monitors
+        elif memory_type == MONITOR_ACTIVATION_MEMORY_TYPE:
+            monitor_dict = self.activation_memory_monitors
+        else:
+            logger.warning(f"Invalid memory type: {memory_type}")
+            return []
+
+        if user_id not in monitor_dict or mem_cube_id not in monitor_dict[user_id]:
+            logger.warning(
+                f"MemoryMonitorManager not found for user {user_id}, "
+                f"mem_cube {mem_cube_id}, type {memory_type}"
+            )
+            return []
+
+        manager: MemoryMonitorManager = monitor_dict[user_id][mem_cube_id]
+        # Sort memories by recording_count in descending order and return top_k items
+        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]:
+        """Retrieves monitoring information for a specific memory cube."""
+        if (
+            user_id not in self.working_memory_monitors
+            or mem_cube_id not in self.working_memory_monitors[user_id]
+        ):
+            logger.warning(
+                f"MemoryMonitorManager not found for user {user_id}, mem_cube {mem_cube_id}"
+            )
+            return {}
+
+        info_dict = {}
+        for manager in [
+            self.working_memory_monitors[user_id][mem_cube_id],
+            self.activation_memory_monitors[user_id][mem_cube_id],
+        ]:
+            info_dict[str(type(manager))] = {
+                "user_id": user_id,
+                "mem_cube_id": mem_cube_id,
+                "memory_count": manager.memory_size,
+                "max_capacity": manager.max_capacity,
+                "top_memories": self.get_scheduler_working_memories(user_id, mem_cube_id, top_k=1),
+            }
+        return info_dict
 
     def detect_intent(
         self,
@@ -55,28 +262,11 @@ class SchedulerMonitor(BaseSchedulerModule):
             q_list=q_list,
             working_memory_list=text_working_memory,
         )
-        response = self._chat_llm.generate([{"role": "user", "content": prompt}])
-        response = extract_json_dict(response)
-        return response
-
-    def update_freq(
-        self,
-        answer: str,
-        activation_memory_freq_list: list[dict],
-        prompt_name="freq_detecting",
-    ) -> list[dict]:
-        """
-        Use LLM to detect which memories in activation_memory_freq_list appear in the answer,
-        increment their count by 1, and return the updated list.
-        """
-        prompt = self.build_prompt(
-            template_name=prompt_name,
-            answer=answer,
-            activation_memory_freq_list=activation_memory_freq_list,
-        )
-        response = self._chat_llm.generate([{"role": "user", "content": prompt}])
+        response = self._process_llm.generate([{"role": "user", "content": prompt}])
         try:
-            result = json.loads(response)
+            response = extract_json_dict(response)
+            assert ("trigger_retrieval" in response) and ("missing_evidences" in response)
         except Exception:
-            result = activation_memory_freq_list
-        return result
+            logger.error(f"Fail to extract json dict from response: {response}")
+            response = {"trigger_retrieval": False, "missing_evidences": q_list}
+        return response