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

memos/mem_scheduler/modules/monitor.py

@@ -1,45 +1,241 @@
-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.modules.misc import AutoDroppingQueue as Queue
+from memos.mem_scheduler.modules.schemas import (
+    DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT,
+    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 TreeTextMemory
+from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
 
 
 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", 300)
+
+        # 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
+        self.query_history = Queue(maxsize=self.config.context_window_size)
+        self.intent_history = Queue(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._process_llm = process_llm
+
+    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.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):
+        text_mem_base: TreeTextMemory = mem_cube.text_mem
+
+        if not isinstance(text_mem_base, TreeTextMemory):
             logger.error("Not Implemented")
+            return
+
+        self.working_mem_monitor_capacity = min(
+            DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT,
+            (
+                text_mem_base.memory_manager.memory_size["WorkingMemory"]
+                + self.partial_retention_number
+            ),
+        )
+
+        self.update_working_memory_monitors(
+            user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
+        )
 
-        return mem_cube_info
+        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,
+            mem_cube_id=mem_cube_id,
+            memory_monitors=self.working_memory_monitors,
+            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,
+            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_score(),
+            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,
+            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 = 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]]
+        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 +251,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

memos/mem_scheduler/modules/rabbitmq_service.py

@@ -0,0 +1,317 @@
+import json
+import ssl
+import threading
+import time
+
+from pathlib import Path
+from queue import Queue
+
+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.schemas import DIRECT_EXCHANGE_TYPE, FANOUT_EXCHANGE_TYPE
+
+
+logger = get_logger(__name__)
+
+
+class RabbitMQSchedulerModule(BaseSchedulerModule):
+    @require_python_package(
+        import_name="pika",
+        install_command="pip install pika",
+        install_link="https://pika.readthedocs.io/en/stable/index.html",
+    )
+    def __init__(self):
+        """
+        Initialize RabbitMQ connection settings.
+        """
+        super().__init__()
+
+        # RabbitMQ settings
+        self.rabbitmq_config: RabbitMQConfig | None = None
+        self.rabbit_queue_name = "memos-scheduler"
+        self.rabbitmq_exchange_name = "memos-fanout"
+        self.rabbitmq_exchange_type = FANOUT_EXCHANGE_TYPE
+        self.rabbitmq_connection = None
+        self.rabbitmq_channel = None
+
+        # fixed params
+        self.rabbitmq_message_cache_max_size = 10  # Max 10 messages
+        self.rabbitmq_message_cache = Queue(maxsize=self.rabbitmq_message_cache_max_size)
+        self.rabbitmq_connection_attempts = 3  # Max retry attempts on connection failure
+        self.rabbitmq_retry_delay = 5  # Delay (seconds) between retries
+        self.rabbitmq_heartbeat = 60  # Heartbeat interval (seconds) for connectio
+        self.rabbitmq_conn_max_waiting_seconds = 30
+        self.rabbitmq_conn_sleep_seconds = 1
+
+        # Thread management
+        self._rabbitmq_io_loop_thread = None  # For IOLoop execution
+        self._rabbitmq_stop_flag = False  # Graceful shutdown flag
+        self._rabbitmq_lock = threading.Lock()  # Ensure thread safety
+
+    def is_rabbitmq_connected(self) -> bool:
+        """Check if RabbitMQ connection is alive"""
+        return (
+            self.rabbitmq_connection
+            and self.rabbitmq_connection.is_open
+            and self.rabbitmq_channel
+            and self.rabbitmq_channel.is_open
+        )
+
+    def initialize_rabbitmq(
+        self, config: dict | None | RabbitMQConfig = None, config_path: str | Path | None = None
+    ):
+        """
+        Establish connection to RabbitMQ using pika.
+        """
+        from pika.adapters.select_connection import SelectConnection
+
+        if config is None:
+            if config_path is None and AuthConfig.default_config_exists():
+                auth_config = AuthConfig.from_local_yaml()
+            elif Path(config_path).exists():
+                auth_config = AuthConfig.from_local_yaml(config_path=config_path)
+            else:
+                logger.error("Fail to initialize auth_config")
+                return
+            self.rabbitmq_config = auth_config.rabbitmq
+        elif isinstance(config, RabbitMQConfig):
+            self.rabbitmq_config = config
+        elif isinstance(config, dict):
+            self.rabbitmq_config = AuthConfig.from_dict(config).rabbitmq
+        else:
+            logger.error("Not implemented")
+
+            # Start connection process
+        parameters = self.get_rabbitmq_connection_param()
+        self.rabbitmq_connection = SelectConnection(
+            parameters,
+            on_open_callback=self.on_rabbitmq_connection_open,
+            on_open_error_callback=self.on_rabbitmq_connection_error,
+            on_close_callback=self.on_rabbitmq_connection_closed,
+        )
+
+        # Start IOLoop in dedicated thread
+        self._io_loop_thread = threading.Thread(
+            target=self.rabbitmq_connection.ioloop.start, daemon=True
+        )
+        self._io_loop_thread.start()
+        logger.info("RabbitMQ connection process started")
+
+    def get_rabbitmq_queue_size(self) -> int:
+        """Get the current number of messages in the queue.
+
+        Returns:
+            int: Number of messages in the queue.
+                 Returns -1 if there's an error or no active connection.
+        """
+        if self.rabbitmq_exchange_type != DIRECT_EXCHANGE_TYPE:
+            logger.warning("Queue size can only be checked for direct exchanges")
+            return None
+
+        with self._rabbitmq_lock:
+            if not self.is_rabbitmq_connected():
+                logger.warning("No active connection to check queue size")
+                return -1
+
+            # Declare queue passively (only checks existence, doesn't create)
+            # Using passive=True prevents accidental queue creation
+            result = self.rabbitmq_channel.queue_declare(
+                queue=self.rabbit_queue_name,
+                durable=True,  # Match the original queue durability setting
+                passive=True,  # Only check queue existence, don't create
+            )
+
+            if result is None:
+                return 0
+            # Return the message count from the queue declaration result
+            return result.method.message_count
+
+    def get_rabbitmq_connection_param(self):
+        import pika
+
+        credentials = pika.PlainCredentials(
+            username=self.rabbitmq_config.user_name,
+            password=self.rabbitmq_config.password,
+            erase_on_connect=self.rabbitmq_config.erase_on_connect,
+        )
+        if self.rabbitmq_config.port == 5671:
+            context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+            context.check_hostname = False
+            context.verify_mode = False
+            return pika.ConnectionParameters(
+                host=self.rabbitmq_config.host_name,
+                port=self.rabbitmq_config.port,
+                virtual_host=self.rabbitmq_config.virtual_host,
+                credentials=credentials,
+                ssl_options=pika.SSLOptions(context),
+                connection_attempts=self.rabbitmq_connection_attempts,
+                retry_delay=self.rabbitmq_retry_delay,
+                heartbeat=self.rabbitmq_heartbeat,
+            )
+        else:
+            return pika.ConnectionParameters(
+                host=self.rabbitmq_config.host_name,
+                port=self.rabbitmq_config.port,
+                virtual_host=self.rabbitmq_config.virtual_host,
+                credentials=credentials,
+                connection_attempts=self.rabbitmq_connection_attempts,
+                retry_delay=self.rabbitmq_retry_delay,
+                heartbeat=self.rabbitmq_heartbeat,
+            )
+
+    # Connection lifecycle callbacks
+    def on_rabbitmq_connection_open(self, connection):
+        """Called when connection is established."""
+        logger.debug("Connection opened")
+        connection.channel(on_open_callback=self.on_rabbitmq_channel_open)
+
+    def on_rabbitmq_connection_error(self, connection, error):
+        """Called if connection fails to open."""
+        logger.error(f"Connection failed: {error}")
+        self.rabbit_reconnect()
+
+    def on_rabbitmq_connection_closed(self, connection, reason):
+        """Called when connection closes."""
+        logger.warning(f"Connection closed: {reason}")
+        if not self._rabbitmq_stop_flag:
+            self.rabbit_reconnect()
+
+    # Channel lifecycle callbacks
+    def on_rabbitmq_channel_open(self, channel):
+        """Called when channel is ready."""
+        self.rabbitmq_channel = channel
+        logger.debug("Channel opened")
+
+        # Setup exchange and queue
+        channel.exchange_declare(
+            exchange=self.rabbitmq_exchange_name,
+            exchange_type=self.rabbitmq_exchange_type,
+            durable=True,
+            callback=self.on_rabbitmq_exchange_declared,
+        )
+
+    def on_rabbitmq_exchange_declared(self, frame):
+        """Called when exchange is ready."""
+        self.rabbitmq_channel.queue_declare(
+            queue=self.rabbit_queue_name, durable=True, callback=self.on_rabbitmq_queue_declared
+        )
+
+    def on_rabbitmq_queue_declared(self, frame):
+        """Called when queue is ready."""
+        self.rabbitmq_channel.queue_bind(
+            exchange=self.rabbitmq_exchange_name,
+            queue=self.rabbit_queue_name,
+            routing_key=self.rabbit_queue_name,
+            callback=self.on_rabbitmq_bind_ok,
+        )
+
+    def on_rabbitmq_bind_ok(self, frame):
+        """Final setup step when bind is complete."""
+        logger.info("RabbitMQ setup completed")
+
+    def on_rabbitmq_message(self, channel, method, properties, body):
+        """Handle incoming messages. Only for test."""
+        try:
+            print(f"Received message: {body.decode()}")
+            self.rabbitmq_message_cache.put_nowait({"properties": properties, "body": body})
+            print(f"message delivery_tag: {method.delivery_tag}")
+            channel.basic_ack(delivery_tag=method.delivery_tag)
+        except Exception as e:
+            logger.error(f"Message handling failed: {e}")
+
+    def wait_for_connection_ready(self):
+        start_time = time.time()
+        while not self.is_rabbitmq_connected():
+            delta_time = time.time() - start_time
+            if delta_time > self.rabbitmq_conn_max_waiting_seconds:
+                logger.error("Failed to start consuming: Connection timeout")
+                return False
+            self.rabbit_reconnect()
+            time.sleep(self.rabbitmq_conn_sleep_seconds)  # Reduced frequency of checks
+
+    # Message handling
+    def rabbitmq_start_consuming(self):
+        """Start consuming messages asynchronously."""
+        self.wait_for_connection_ready()
+
+        self.rabbitmq_channel.basic_consume(
+            queue=self.rabbit_queue_name,
+            on_message_callback=self.on_rabbitmq_message,
+            auto_ack=False,
+        )
+        logger.info("Started rabbitmq consuming messages")
+
+    def rabbitmq_publish_message(self, message: dict):
+        """
+        Publish a message to RabbitMQ.
+        """
+        import pika
+
+        with self._rabbitmq_lock:
+            if not self.is_rabbitmq_connected():
+                logger.error("Cannot publish - no active connection")
+                return False
+
+            try:
+                self.rabbitmq_channel.basic_publish(
+                    exchange=self.rabbitmq_exchange_name,
+                    routing_key=self.rabbit_queue_name,
+                    body=json.dumps(message),
+                    properties=pika.BasicProperties(
+                        delivery_mode=2,  # Persistent
+                    ),
+                    mandatory=True,
+                )
+                logger.debug(f"Published message: {message}")
+                return True
+            except Exception as e:
+                logger.error(f"Failed to publish message: {e}")
+                self.rabbit_reconnect()
+                return False
+
+    # Connection management
+    def rabbit_reconnect(self):
+        """Schedule reconnection attempt."""
+        logger.info("Attempting to reconnect...")
+        if self.rabbitmq_connection and not self.rabbitmq_connection.is_closed:
+            self.rabbitmq_connection.ioloop.stop()
+
+        # Reset connection state
+        self.rabbitmq_channel = None
+        self.initialize_rabbitmq()
+
+    def rabbitmq_close(self):
+        """Gracefully shutdown connection."""
+        with self._rabbitmq_lock:
+            self._rabbitmq_stop_flag = True
+
+            # Close channel if open
+            if self.rabbitmq_channel and self.rabbitmq_channel.is_open:
+                try:
+                    self.rabbitmq_channel.close()
+                except Exception as e:
+                    logger.warning(f"Error closing channel: {e}")
+
+            # Close connection if open
+            if self.rabbitmq_connection:
+                if self.rabbitmq_connection.is_open:
+                    try:
+                        self.rabbitmq_connection.close()
+                    except Exception as e:
+                        logger.warning(f"Error closing connection: {e}")
+
+                # Stop IOLoop if running
+                try:
+                    self.rabbitmq_connection.ioloop.stop()
+                except Exception as e:
+                    logger.warning(f"Error stopping IOLoop: {e}")
+
+            # Wait for IOLoop thread to finish
+            if self._io_loop_thread and self._io_loop_thread.is_alive():
+                self._io_loop_thread.join(timeout=5)
+                if self._io_loop_thread.is_alive():
+                    logger.warning("IOLoop thread did not terminate cleanly")
+
+        logger.info("RabbitMQ connection closed")