MemOS-MCP

import queue import threading import time from abc import abstractmethod from queue import Queue from memos.configs.mem_scheduler import BaseSchedulerConfig from memos.llms.base import BaseLLM from memos.log import get_logger from memos.mem_scheduler.modules.dispatcher import SchedulerDispatcher from memos.mem_scheduler.modules.redis_service import RedisSchedulerModule from memos.mem_scheduler.modules.schemas import ( DEFAULT_CONSUME_INTERVAL_SECONDS, DEFAULT_THREAD__POOL_MAX_WORKERS, ScheduleLogForWebItem, ScheduleMessageItem, ) logger = get_logger(__name__) class BaseScheduler(RedisSchedulerModule): """Base class for all mem_scheduler.""" def __init__(self, config: BaseSchedulerConfig): """Initialize the scheduler with the given configuration.""" super().__init__() self.config = config self.max_workers = self.config.get( "thread_pool_max_workers", DEFAULT_THREAD__POOL_MAX_WORKERS ) self.retriever = None self.monitor = None self.enable_parallel_dispatch = self.config.get("enable_parallel_dispatch", False) self.dispatcher = SchedulerDispatcher( max_workers=self.max_workers, enable_parallel_dispatch=self.enable_parallel_dispatch ) # message queue self.memos_message_queue: Queue[ScheduleMessageItem] = Queue() self._web_log_message_queue: Queue[ScheduleLogForWebItem] = Queue() self._consumer_thread = None # Reference to our consumer thread self._running = False self._consume_interval = self.config.get( "consume_interval_seconds", DEFAULT_CONSUME_INTERVAL_SECONDS ) # others self._current_user_id: str | None = None @abstractmethod def initialize_modules(self, chat_llm: BaseLLM) -> None: """Initialize all necessary modules for the scheduler Args: chat_llm: The LLM instance to be used for chat interactions """ def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]): """Submit multiple messages to the message queue.""" if isinstance(messages, ScheduleMessageItem): messages = [messages] # transform single message to list for message in messages: self.memos_message_queue.put(message) logger.info(f"Submitted message: {message.label} - {message.content}") def _submit_web_logs(self, messages: ScheduleLogForWebItem | list[ScheduleLogForWebItem]): if isinstance(messages, ScheduleLogForWebItem): messages = [messages] # transform single message to list for message in messages: self._web_log_message_queue.put(message) logger.info( f"Submitted Scheduling log for web: {message.log_title} - {message.log_content}" ) logger.debug(f"{len(messages)} submitted. {self._web_log_message_queue.qsize()} in queue.") def get_web_log_messages(self) -> list[dict]: """ Retrieves all web log messages from the queue and returns them as a list of JSON-serializable dictionaries. Returns: List[dict]: A list of dictionaries representing ScheduleLogForWebItem objects, ready for JSON serialization. The list is ordered from oldest to newest. """ messages = [] # Process all items in the queue while not self._web_log_message_queue.empty(): item = self._web_log_message_queue.get() # Convert the ScheduleLogForWebItem to a dictionary and ensure datetime is serialized item_dict = item.to_dict() messages.append(item_dict) return messages def _message_consumer(self) -> None: """ Continuously checks the queue for messages and dispatches them. Runs in a dedicated thread to process messages at regular intervals. """ while self._running: # Use a running flag for graceful shutdown try: # Check if queue has messages (non-blocking) if not self.memos_message_queue.empty(): # Get all available messages at once messages = [] while not self.memos_message_queue.empty(): try: messages.append(self.memos_message_queue.get_nowait()) except queue.Empty: break if messages: try: self.dispatcher.dispatch(messages) except Exception as e: logger.error(f"Error dispatching messages: {e!s}") finally: # Mark all messages as processed for _ in messages: self.memos_message_queue.task_done() # Sleep briefly to prevent busy waiting time.sleep(self._consume_interval) # Adjust interval as needed except Exception as e: logger.error(f"Unexpected error in message consumer: {e!s}") time.sleep(self._consume_interval) # Prevent tight error loops def start(self) -> None: """ Start the message consumer thread. Initializes and starts a daemon thread that will periodically check for and process messages from the queue. """ if self._consumer_thread is not None and self._consumer_thread.is_alive(): logger.warning("Consumer thread is already running") return self._running = True self._consumer_thread = threading.Thread( target=self._message_consumer, daemon=True, # Allows program to exit even if thread is running name="MessageConsumerThread", ) self._consumer_thread.start() logger.info("Message consumer thread started") def stop(self) -> None: """Stop the consumer thread and clean up resources.""" if self._consumer_thread is None or not self._running: logger.warning("Consumer thread is not running") return self._running = False if self._consumer_thread.is_alive(): self._consumer_thread.join(timeout=5.0) # Wait up to 5 seconds if self._consumer_thread.is_alive(): logger.warning("Consumer thread did not stop gracefully") logger.info("Message consumer thread stopped")