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

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

@@ -3,7 +3,7 @@ from pathlib import Path
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
-from memos.mem_scheduler.modules.schemas import BASE_DIR
+from memos.mem_scheduler.schemas.general_schemas import BASE_DIR
 from memos.templates.mem_scheduler_prompts import PROMPT_MAPPING
 
 
@@ -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:
@@ -75,9 +74,9 @@ class BaseSchedulerModule:
     @property
     def mem_cube(self) -> GeneralMemCube:
         """The memory cube associated with this MemChat."""
-        return self._current_mem_cube
+        return self.current_mem_cube
 
     @mem_cube.setter
     def mem_cube(self, value: GeneralMemCube) -> None:
         """The memory cube associated with this MemChat."""
-        self._current_mem_cube = value
+        self.current_mem_cube = value

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

@@ -1,10 +1,12 @@
+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.modules.schemas import ScheduleMessageItem
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 logger = get_logger(__name__)
@@ -22,24 +24,31 @@ 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
+
         # 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,32 +128,29 @@ 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)} 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
-                # TODO check this
                 future = self.dispatcher_executor.submit(handler, msgs)
-                logger.debug(f"Dispatched {len(msgs)} messages as future task")
-                return future
+                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)
-                return None
 
     def join(self, timeout: float | None = None) -> bool:
         """Wait for all dispatched tasks to complete.
@@ -151,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/general_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 → general_modules}/rabbitmq_service.py

@@ -4,13 +4,13 @@ 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
+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
 
 
 logger = get_logger(__name__)
@@ -38,7 +38,9 @@ class RabbitMQSchedulerModule(BaseSchedulerModule):
 
         # 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_message_cache = AutoDroppingQueue(
+            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
@@ -69,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
@@ -214,12 +216,12 @@ class RabbitMQSchedulerModule(BaseSchedulerModule):
     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}")
+            print(f"Received message: {body.decode()}\n")
+            self.rabbitmq_message_cache.put({"properties": properties, "body": body})
+            print(f"message delivery_tag: {method.delivery_tag}\n")
             channel.basic_ack(delivery_tag=method.delivery_tag)
         except Exception as e:
-            logger.error(f"Message handling failed: {e}")
+            logger.error(f"Message handling failed: {e}", exc_info=True)
 
     def wait_for_connection_ready(self):
         start_time = time.time()

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/general_modules/retriever.py

@@ -0,0 +1,199 @@
+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.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 (
+    filter_similar_memories,
+    filter_too_short_memories,
+    transform_name_to_key,
+)
+from memos.mem_scheduler.utils.misc_utils import (
+    extract_json_dict,
+)
+from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerRetriever(BaseSchedulerModule):
+    def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
+        super().__init__()
+
+        # hyper-parameters
+        self.filter_similarity_threshold = 0.75
+        self.filter_min_length_threshold = 6
+
+        self.config: BaseSchedulerConfig = config
+        self.process_llm = process_llm
+
+    def search(
+        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.
+
+        Args:
+            query: The search query string
+            top_k: Number of top results to return
+            method: Search method to use
+
+        Returns:
+            Search results or None if not implemented
+        """
+        text_mem_base = mem_cube.text_mem
+        try:
+            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", mode=mode, info=info
+                )
+                results_user = text_mem_base.search(
+                    query=query, top_k=top_k, memory_type="UserMemory", mode=mode, info=info
+                )
+                results = results_long_term + results_user
+            else:
+                raise NotImplementedError(str(type(text_mem_base)))
+        except Exception as e:
+            logger.error(f"Fail to search. The exeption is {e}.", exc_info=True)
+            results = []
+        return results
+
+    def rerank_memories(
+        self,
+        queries: list[str],
+        original_memories: list[str],
+        top_k: int,
+    ) -> (list[str], bool):
+        """
+        Rerank memories based on relevance to given queries using LLM.
+
+        Args:
+            queries: List of query strings to determine relevance
+            original_memories: List of memory strings to be reranked
+            top_k: Number of top memories to return after reranking
+
+        Returns:
+            List of reranked memory strings (length <= top_k)
+
+        Note:
+            If LLM reranking fails, falls back to original order (truncated to top_k)
+        """
+        success_flag = False
+
+        logger.info(f"Starting memory reranking for {len(original_memories)} memories")
+
+        # Build LLM prompt for memory reranking
+        prompt = self.build_prompt(
+            "memory_reranking",
+            queries=[f"[0] {queries[0]}"],
+            current_order=[f"[{i}] {mem}" for i, mem in enumerate(original_memories)],
+        )
+        logger.debug(f"Generated reranking prompt: {prompt[:200]}...")  # Log first 200 chars
+
+        # Get LLM response
+        response = self.process_llm.generate([{"role": "user", "content": prompt}])
+        logger.debug(f"Received LLM response: {response[:200]}...")  # Log first 200 chars
+
+        try:
+            # Parse JSON response
+            response = extract_json_dict(response)
+            new_order = response["new_order"][:top_k]
+            text_memories_with_new_order = [original_memories[idx] for idx in new_order]
+            logger.info(
+                f"Successfully reranked memories. Returning top {len(text_memories_with_new_order)} items;"
+                f"Ranking reasoning: {response['reasoning']}"
+            )
+            success_flag = True
+        except Exception as e:
+            logger.error(
+                f"Failed to rerank memories with LLM. Exception: {e}. Raw response: {response} ",
+                exc_info=True,
+            )
+            text_memories_with_new_order = original_memories[:top_k]
+            success_flag = False
+        return text_memories_with_new_order, success_flag
+
+    def process_and_rerank_memories(
+        self,
+        queries: list[str],
+        original_memory: list[TextualMemoryItem],
+        new_memory: list[TextualMemoryItem],
+        top_k: int = 10,
+    ) -> list[TextualMemoryItem] | None:
+        """
+        Process and rerank memory items by combining original and new memories,
+        applying filters, and then reranking based on relevance to queries.
+
+        Args:
+            queries: List of query strings to rerank memories against
+            original_memory: List of original TextualMemoryItem objects
+            new_memory: List of new TextualMemoryItem objects to merge
+            top_k: Maximum number of memories to return after reranking
+
+        Returns:
+            List of reranked TextualMemoryItem objects, or None if processing fails
+        """
+        # Combine original and new memories into a single list
+        combined_memory = original_memory + new_memory
+
+        # Create a mapping from normalized text to memory objects
+        memory_map = {
+            transform_name_to_key(name=mem_obj.memory): mem_obj for mem_obj in combined_memory
+        }
+
+        # Extract normalized text representations from all memory items
+        combined_text_memory = [m.memory for m in combined_memory]
+
+        # Apply similarity filter to remove overly similar memories
+        filtered_combined_text_memory = filter_similar_memories(
+            text_memories=combined_text_memory,
+            similarity_threshold=self.filter_similarity_threshold,
+        )
+
+        # Apply length filter to remove memories that are too short
+        filtered_combined_text_memory = filter_too_short_memories(
+            text_memories=filtered_combined_text_memory,
+            min_length_threshold=self.filter_min_length_threshold,
+        )
+
+        # Ensure uniqueness of memory texts using dictionary keys (preserves order)
+        unique_memory = list(dict.fromkeys(filtered_combined_text_memory))
+
+        # Rerank the filtered memories based on relevance to the queries
+        text_memories_with_new_order, success_flag = self.rerank_memories(
+            queries=queries,
+            original_memories=unique_memory,
+            top_k=top_k,
+        )
+
+        # Map reranked text entries back to their original memory objects
+        memories_with_new_order = []
+        for text in text_memories_with_new_order:
+            normalized_text = transform_name_to_key(name=text)
+            if normalized_text in memory_map:  # Ensure correct key matching
+                memories_with_new_order.append(memory_map[normalized_text])
+            else:
+                logger.warning(
+                    f"Memory text not found in memory map. text: {text};\n"
+                    f"Keys of memory_map: {memory_map.keys()}"
+                )
+
+        return memories_with_new_order, success_flag