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

memos/mem_scheduler/base_scheduler.py

@@ -15,24 +15,23 @@ from memos.mem_scheduler.modules.monitor import SchedulerMonitor
 from memos.mem_scheduler.modules.rabbitmq_service import RabbitMQSchedulerModule
 from memos.mem_scheduler.modules.redis_service import RedisSchedulerModule
 from memos.mem_scheduler.modules.retriever import SchedulerRetriever
-from memos.mem_scheduler.modules.schemas import (
-    ACTIVATION_MEMORY_TYPE,
-    ADD_LABEL,
+from memos.mem_scheduler.modules.scheduler_logger import SchedulerLoggerModule
+from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_ACT_MEM_DUMP_PATH,
     DEFAULT_CONSUME_INTERVAL_SECONDS,
     DEFAULT_THREAD__POOL_MAX_WORKERS,
-    LONG_TERM_MEMORY_TYPE,
-    NOT_INITIALIZED,
-    PARAMETER_MEMORY_TYPE,
-    QUERY_LABEL,
-    TEXT_MEMORY_TYPE,
-    USER_INPUT_TYPE,
-    WORKING_MEMORY_TYPE,
+    MemCubeID,
+    TreeTextMemory_SEARCH_METHOD,
+    UserID,
+)
+from memos.mem_scheduler.schemas.message_schemas import (
     ScheduleLogForWebItem,
     ScheduleMessageItem,
-    TreeTextMemory_SEARCH_METHOD,
 )
-from memos.mem_scheduler.utils import transform_name_to_key
+from memos.mem_scheduler.schemas.monitor_schemas import MemoryMonitorItem
+from memos.mem_scheduler.utils.filter_utils import (
+    transform_name_to_key,
+)
 from memos.memories.activation.kv import KVCacheMemory
 from memos.memories.activation.vllmkv import VLLMKVCacheItem, VLLMKVCacheMemory
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
@@ -42,7 +41,7 @@ from memos.templates.mem_scheduler_prompts import MEMORY_ASSEMBLY_TEMPLATE
 logger = get_logger(__name__)
 
 
-class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
+class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLoggerModule):
     """Base class for all mem_scheduler."""
 
     def __init__(self, config: BaseSchedulerConfig):
@@ -51,12 +50,11 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
         self.config = config
 
         # hyper-parameters
-        self.top_k = self.config.get("top_k", 5)
+        self.top_k = self.config.get("top_k", 10)
         self.context_window_size = self.config.get("context_window_size", 5)
         self.enable_act_memory_update = self.config.get("enable_act_memory_update", False)
         self.act_mem_dump_path = self.config.get("act_mem_dump_path", DEFAULT_ACT_MEM_DUMP_PATH)
         self.search_method = TreeTextMemory_SEARCH_METHOD
-
         self.enable_parallel_dispatch = self.config.get("enable_parallel_dispatch", False)
         self.max_workers = self.config.get(
             "thread_pool_max_workers", DEFAULT_THREAD__POOL_MAX_WORKERS
@@ -85,7 +83,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
 
         # other attributes
         self._context_lock = threading.Lock()
-        self._current_user_id: str | None = None
+        self.current_user_id: UserID | str | None = None
         self.auth_config_path: str | Path | None = self.config.get("auth_config_path", None)
         self.auth_config = None
         self.rabbitmq_config = None
@@ -99,7 +97,6 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
         self.process_llm = process_llm
         self.monitor = SchedulerMonitor(process_llm=self.process_llm, config=self.config)
         self.retriever = SchedulerRetriever(process_llm=self.process_llm, config=self.config)
-        self.retriever.log_working_memory_replacement = self.log_working_memory_replacement
 
         # initialize with auth_cofig
         if self.auth_config_path is not None and Path(self.auth_config_path).exists():
@@ -118,58 +115,161 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
     @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
         self.retriever.mem_cube = value
 
     def _set_current_context_from_message(self, msg: ScheduleMessageItem) -> None:
         """Update current user/cube context from the incoming message (thread-safe)."""
         with self._context_lock:
-            self._current_user_id = msg.user_id
-            self._current_mem_cube_id = msg.mem_cube_id
-            self._current_mem_cube = msg.mem_cube
+            self.current_user_id = msg.user_id
+            self.current_mem_cube_id = msg.mem_cube_id
+            self.current_mem_cube = msg.mem_cube
 
-    def _validate_messages(self, messages: list[ScheduleMessageItem], label: str):
-        """Validate if all messages match the expected label.
+    def transform_memories_to_monitors(
+        self, memories: list[TextualMemoryItem]
+    ) -> list[MemoryMonitorItem]:
+        """
+        Convert a list of TextualMemoryItem objects into MemoryMonitorItem objects
+        with importance scores based on keyword matching.
 
         Args:
-            messages: List of message items to validate.
-            label: Expected message label (e.g., QUERY_LABEL/ANSWER_LABEL).
+            memories: List of TextualMemoryItem objects to be transformed.
 
         Returns:
-            bool: True if all messages passed validation, False if any failed.
+            List of MemoryMonitorItem objects with computed importance scores.
         """
-        for message in messages:
-            if not self._validate_message(message, label):
-                return False
-                logger.error("Message batch contains invalid labels, aborting processing")
-        return True
+        query_keywords = self.monitor.query_monitors.get_keywords_collections()
+        logger.debug(
+            f"Processing {len(memories)} memories with {len(query_keywords)} query keywords"
+        )
 
-    def _validate_message(self, message: ScheduleMessageItem, label: str):
-        """Validate if the message matches the expected label.
+        result = []
+        mem_length = len(memories)
+        for idx, mem in enumerate(memories):
+            text_mem = mem.memory
+            mem_key = transform_name_to_key(name=text_mem)
+
+            # Calculate importance score based on keyword matches
+            keywords_score = 0
+            if query_keywords and text_mem:
+                for keyword, count in query_keywords.items():
+                    keyword_count = text_mem.count(keyword)
+                    if keyword_count > 0:
+                        keywords_score += keyword_count * count
+                        logger.debug(
+                            f"Matched keyword '{keyword}' {keyword_count} times, added {keywords_score} to keywords_score"
+                        )
+
+            # rank score
+            sorting_score = mem_length - idx
+
+            mem_monitor = MemoryMonitorItem(
+                memory_text=text_mem,
+                tree_memory_item=mem,
+                tree_memory_item_mapping_key=mem_key,
+                sorting_score=sorting_score,
+                keywords_score=keywords_score,
+                recording_count=1,
+            )
+            result.append(mem_monitor)
 
-        Args:
-            message: Incoming message item to validate.
-            label: Expected message label (e.g., QUERY_LABEL/ANSWER_LABEL).
+        logger.debug(f"Transformed {len(result)} memories to monitors")
+        return result
 
-        Returns:
-            bool: True if validation passed, False otherwise.
-        """
-        if message.label != label:
-            logger.error(f"Handler validation failed: expected={label}, actual={message.label}")
-            return False
-        return True
+    def replace_working_memory(
+        self,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
+        mem_cube: GeneralMemCube,
+        original_memory: list[TextualMemoryItem],
+        new_memory: list[TextualMemoryItem],
+    ) -> None | list[TextualMemoryItem]:
+        """Replace working memory with new memories after reranking."""
+        text_mem_base = mem_cube.text_mem
+        if isinstance(text_mem_base, TreeTextMemory):
+            text_mem_base: TreeTextMemory = text_mem_base
+
+            # process rerank memories with llm
+            query_history = self.monitor.query_monitors.get_queries_with_timesort()
+            memories_with_new_order, rerank_success_flag = (
+                self.retriever.process_and_rerank_memories(
+                    queries=query_history,
+                    original_memory=original_memory,
+                    new_memory=new_memory,
+                    top_k=self.top_k,
+                )
+            )
+
+            # update working memory monitors
+            new_working_memory_monitors = self.transform_memories_to_monitors(
+                memories=memories_with_new_order
+            )
+
+            if not rerank_success_flag:
+                for one in new_working_memory_monitors:
+                    one.sorting_score = 0
+
+            self.monitor.update_working_memory_monitors(
+                new_working_memory_monitors=new_working_memory_monitors,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+            )
+
+            mem_monitors: list[MemoryMonitorItem] = self.monitor.working_memory_monitors[user_id][
+                mem_cube_id
+            ].get_sorted_mem_monitors(reverse=True)
+            new_working_memories = [mem_monitor.tree_memory_item for mem_monitor in mem_monitors]
+
+            text_mem_base.replace_working_memory(memories=new_working_memories)
+
+            logger.info(
+                f"The working memory has been replaced with {len(memories_with_new_order)} new memories."
+            )
+            self.log_working_memory_replacement(
+                original_memory=original_memory,
+                new_memory=new_working_memories,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+                log_func_callback=self._submit_web_logs,
+            )
+        else:
+            logger.error("memory_base is not supported")
+            memories_with_new_order = new_memory
+
+        return memories_with_new_order
+
+    def initialize_working_memory_monitors(
+        self,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
+        mem_cube: GeneralMemCube,
+    ):
+        text_mem_base: TreeTextMemory = mem_cube.text_mem
+        working_memories = text_mem_base.get_working_memory()
+
+        working_memory_monitors = self.transform_memories_to_monitors(
+            memories=working_memories,
+        )
+        self.monitor.update_working_memory_monitors(
+            new_working_memory_monitors=working_memory_monitors,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )
 
     def update_activation_memory(
         self,
         new_memories: list[str | TextualMemoryItem],
         label: str,
-        user_id: str,
-        mem_cube_id: str,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
         mem_cube: GeneralMemCube,
     ) -> None:
         """
@@ -195,7 +295,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
                 logger.error("Not Implemented.")
                 return
 
-            text_memory = MEMORY_ASSEMBLY_TEMPLATE.format(
+            new_text_memory = MEMORY_ASSEMBLY_TEMPLATE.format(
                 memory_text="".join(
                     [
                         f"{i + 1}. {sentence.strip()}\n"
@@ -211,9 +311,18 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
             if len(original_cache_items) > 0:
                 pre_cache_item: VLLMKVCacheItem = original_cache_items[-1]
                 original_text_memories = pre_cache_item.records.text_memories
+                original_composed_text_memory = pre_cache_item.records.composed_text_memory
+                if original_composed_text_memory == new_text_memory:
+                    logger.warning(
+                        "Skipping memory update - new composition matches existing cache: %s",
+                        new_text_memory[:50] + "..."
+                        if len(new_text_memory) > 50
+                        else new_text_memory,
+                    )
+                    return
                 act_mem.delete_all()
 
-            cache_item = act_mem.extract(text_memory)
+            cache_item = act_mem.extract(new_text_memory)
             cache_item.records.text_memories = new_text_memories
 
             act_mem.add([cache_item])
@@ -226,6 +335,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
                 user_id=user_id,
                 mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
+                log_func_callback=self._submit_web_logs,
             )
 
         except Exception as e:
@@ -235,50 +345,70 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
         self,
         interval_seconds: int,
         label: str,
-        user_id: str,
-        mem_cube_id: str,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
         mem_cube: GeneralMemCube,
     ):
-        new_activation_memories = []
-
-        if self.monitor.timed_trigger(
-            last_time=self.monitor._last_activation_mem_update_time,
-            interval_seconds=interval_seconds,
-        ):
-            logger.info(f"Updating activation memory for user {user_id} and mem_cube {mem_cube_id}")
+        try:
+            if (
+                self.monitor.last_activation_mem_update_time == datetime.min
+                or self.monitor.timed_trigger(
+                    last_time=self.monitor.last_activation_mem_update_time,
+                    interval_seconds=interval_seconds,
+                )
+            ):
+                logger.info(
+                    f"Updating activation memory for user {user_id} and mem_cube {mem_cube_id}"
+                )
 
-            self.monitor.update_memory_monitors(
-                user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
-            )
+                if (
+                    user_id not in self.monitor.working_memory_monitors
+                    or mem_cube_id not in self.monitor.working_memory_monitors[user_id]
+                    or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0
+                ):
+                    logger.warning(
+                        "No memories found in working_memory_monitors, initializing from current working_memories"
+                    )
+                    self.initialize_working_memory_monitors(
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=mem_cube,
+                    )
+
+                self.monitor.update_activation_memory_monitors(
+                    user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
+                )
 
-            new_activation_memories = [
-                m.memory_text
-                for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id].memories
-            ]
+                new_activation_memories = [
+                    m.memory_text
+                    for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id].memories
+                ]
 
-            logger.info(
-                f"Collected {len(new_activation_memories)} new memory entries for processing"
-            )
+                logger.info(
+                    f"Collected {len(new_activation_memories)} new memory entries for processing"
+                )
 
-            self.update_activation_memory(
-                new_memories=new_activation_memories,
-                label=label,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
+                self.update_activation_memory(
+                    new_memories=new_activation_memories,
+                    label=label,
+                    user_id=user_id,
+                    mem_cube_id=mem_cube_id,
+                    mem_cube=mem_cube,
+                )
 
-            self.monitor._last_activation_mem_update_time = datetime.now()
+                self.monitor.last_activation_mem_update_time = datetime.now()
 
-            logger.debug(
-                f"Activation memory update completed at {self.monitor._last_activation_mem_update_time}"
-            )
-        else:
-            logger.info(
-                f"Skipping update - {interval_seconds} second interval not yet reached. "
-                f"Last update time is {self.monitor._last_activation_mem_update_time} and now is"
-                f"{datetime.now()}"
-            )
+                logger.debug(
+                    f"Activation memory update completed at {self.monitor.last_activation_mem_update_time}"
+                )
+            else:
+                logger.info(
+                    f"Skipping update - {interval_seconds} second interval not yet reached. "
+                    f"Last update time is {self.monitor.last_activation_mem_update_time} and now is"
+                    f"{datetime.now()}"
+                )
+        except Exception as e:
+            logger.error(f"Error: {e}", exc_info=True)
 
     def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
         """Submit multiple messages to the message queue."""
@@ -289,7 +419,9 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
             self.memos_message_queue.put(message)
             logger.info(f"Submitted message: {message.label} - {message.content}")
 
-    def _submit_web_logs(self, messages: ScheduleLogForWebItem | list[ScheduleLogForWebItem]):
+    def _submit_web_logs(
+        self, messages: ScheduleLogForWebItem | list[ScheduleLogForWebItem]
+    ) -> None:
         """Submit log messages to the web log queue and optionally to RabbitMQ.
 
         Args:
@@ -300,176 +432,14 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule):
 
         for message in messages:
             self._web_log_message_queue.put(message)
-            logger.info(f"Submitted Scheduling log for web: {message.log_content}")
+            message_info = message.debug_info()
+            logger.debug(f"Submitted Scheduling log for web: {message_info}")
 
             if self.is_rabbitmq_connected():
-                logger.info("Submitted Scheduling log to rabbitmq")
+                logger.info(f"Submitted Scheduling log to rabbitmq: {message_info}")
                 self.rabbitmq_publish_message(message=message.to_dict())
         logger.debug(f"{len(messages)} submitted. {self._web_log_message_queue.qsize()} in queue.")
 
-    def log_activation_memory_update(
-        self,
-        original_text_memories: list[str],
-        new_text_memories: list[str],
-        label: str,
-        user_id: str,
-        mem_cube_id: str,
-        mem_cube: GeneralMemCube,
-    ):
-        """Log changes when activation memory is updated.
-
-        Args:
-            original_text_memories: List of original memory texts
-            new_text_memories: List of new memory texts
-        """
-        original_set = set(original_text_memories)
-        new_set = set(new_text_memories)
-
-        # Identify changes
-        added_memories = list(new_set - original_set)  # Present in new but not original
-
-        # recording messages
-        for mem in added_memories:
-            log_message_a = self.create_autofilled_log_item(
-                log_content=mem,
-                label=label,
-                from_memory_type=TEXT_MEMORY_TYPE,
-                to_memory_type=ACTIVATION_MEMORY_TYPE,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
-            log_message_b = self.create_autofilled_log_item(
-                log_content=mem,
-                label=label,
-                from_memory_type=ACTIVATION_MEMORY_TYPE,
-                to_memory_type=PARAMETER_MEMORY_TYPE,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
-            self._submit_web_logs(messages=[log_message_a, log_message_b])
-            logger.info(
-                f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
-                f"transformed to {WORKING_MEMORY_TYPE} memories."
-            )
-
-    def log_working_memory_replacement(
-        self,
-        original_memory: list[TextualMemoryItem],
-        new_memory: list[TextualMemoryItem],
-        user_id: str,
-        mem_cube_id: str,
-        mem_cube: GeneralMemCube,
-    ):
-        """Log changes when working memory is replaced."""
-        memory_type_map = {
-            transform_name_to_key(name=m.memory): m.metadata.memory_type
-            for m in original_memory + new_memory
-        }
-
-        original_text_memories = [m.memory for m in original_memory]
-        new_text_memories = [m.memory for m in new_memory]
-
-        # Convert to sets for efficient difference operations
-        original_set = set(original_text_memories)
-        new_set = set(new_text_memories)
-
-        # Identify changes
-        added_memories = list(new_set - original_set)  # Present in new but not original
-
-        # recording messages
-        for mem in added_memories:
-            normalized_mem = transform_name_to_key(name=mem)
-            if normalized_mem not in memory_type_map:
-                logger.error(f"Memory text not found in type mapping: {mem[:50]}...")
-            # Get the memory type from the map, default to LONG_TERM_MEMORY_TYPE if not found
-            mem_type = memory_type_map.get(normalized_mem, LONG_TERM_MEMORY_TYPE)
-
-            if mem_type == WORKING_MEMORY_TYPE:
-                logger.warning(f"Memory already in working memory: {mem[:50]}...")
-                continue
-
-            log_message = self.create_autofilled_log_item(
-                log_content=mem,
-                label=QUERY_LABEL,
-                from_memory_type=mem_type,
-                to_memory_type=WORKING_MEMORY_TYPE,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
-            self._submit_web_logs(messages=log_message)
-            logger.info(
-                f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
-                f"transformed to {WORKING_MEMORY_TYPE} memories."
-            )
-
-    def log_adding_user_inputs(
-        self,
-        user_inputs: list[str],
-        user_id: str,
-        mem_cube_id: str,
-        mem_cube: GeneralMemCube,
-    ):
-        """Log changes when working memory is replaced."""
-
-        # recording messages
-        for input_str in user_inputs:
-            log_message = self.create_autofilled_log_item(
-                log_content=input_str,
-                label=ADD_LABEL,
-                from_memory_type=USER_INPUT_TYPE,
-                to_memory_type=TEXT_MEMORY_TYPE,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
-            self._submit_web_logs(messages=log_message)
-            logger.info(
-                f"{len(user_inputs)} {USER_INPUT_TYPE} memorie(s) "
-                f"transformed to {TEXT_MEMORY_TYPE} memories."
-            )
-
-    def create_autofilled_log_item(
-        self,
-        log_content: str,
-        label: str,
-        from_memory_type: str,
-        to_memory_type: str,
-        user_id: str,
-        mem_cube_id: str,
-        mem_cube: GeneralMemCube,
-    ) -> ScheduleLogForWebItem:
-        text_mem_base: TreeTextMemory = mem_cube.text_mem
-        current_memory_sizes = text_mem_base.get_current_memory_size()
-        current_memory_sizes = {
-            "long_term_memory_size": current_memory_sizes["LongTermMemory"],
-            "user_memory_size": current_memory_sizes["UserMemory"],
-            "working_memory_size": current_memory_sizes["WorkingMemory"],
-            "transformed_act_memory_size": NOT_INITIALIZED,
-            "parameter_memory_size": NOT_INITIALIZED,
-        }
-        memory_capacities = {
-            "long_term_memory_capacity": text_mem_base.memory_manager.memory_size["LongTermMemory"],
-            "user_memory_capacity": text_mem_base.memory_manager.memory_size["UserMemory"],
-            "working_memory_capacity": text_mem_base.memory_manager.memory_size["WorkingMemory"],
-            "transformed_act_memory_capacity": NOT_INITIALIZED,
-            "parameter_memory_capacity": NOT_INITIALIZED,
-        }
-
-        log_message = ScheduleLogForWebItem(
-            user_id=user_id,
-            mem_cube_id=mem_cube_id,
-            label=label,
-            from_memory_type=from_memory_type,
-            to_memory_type=to_memory_type,
-            log_content=log_content,
-            current_memory_sizes=current_memory_sizes,
-            memory_capacities=memory_capacities,
-        )
-        return log_message
-
     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.