MemOS-MCP

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.modules.schemas import ScheduleMessageItem logger = get_logger(__name__) class SchedulerDispatcher(BaseSchedulerModule): """ Thread pool-based message dispatcher that routes messages to dedicated handlers based on their labels. Features: - Dedicated thread pool per message label - Batch message processing - Graceful shutdown - Bulk handler registration """ def __init__(self, max_workers=3, enable_parallel_dispatch=False): 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 if self.enable_parallel_dispatch: self.dispatcher_executor = ThreadPoolExecutor( max_workers=self.max_workers, thread_name_prefix="dispatcher" ) 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 def register_handler(self, label: str, handler: Callable[[list[ScheduleMessageItem]], None]): """ Register a handler function for a specific message label. Args: label: Message label to handle handler: Callable that processes messages of this label """ self.handlers[label] = handler def register_handlers( self, handlers: dict[str, Callable[[list[ScheduleMessageItem]], None]] ) -> None: """ Bulk register multiple handlers from a dictionary. Args: handlers: Dictionary mapping labels to handler functions Format: {label: handler_callable} """ for label, handler in handlers.items(): if not isinstance(label, str): logger.error(f"Invalid label type: {type(label)}. Expected str.") continue if not callable(handler): logger.error(f"Handler for label '{label}' is not callable.") continue self.register_handler(label=label, handler=handler) logger.info(f"Registered {len(handlers)} handlers in bulk") def _default_message_handler(self, messages: list[ScheduleMessageItem]) -> None: logger.debug(f"Using _default_message_handler to deal with messages: {messages}") def dispatch(self, msg_list: list[ScheduleMessageItem]): """ Dispatch a list of messages to their respective handlers. Args: msg_list: List of ScheduleMessageItem objects to process """ # Group messages by their labels 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] # dispatch to different handler logger.debug(f"Dispatch {len(msgs)} messages to {label} handler.") if self.enable_parallel_dispatch and self.dispatcher_executor is not None: self.dispatcher_executor.submit(handler, msgs) else: handler(msgs) # Direct serial execution