fix bugs: fix bugs of filter and updating activation memories

Author: tangg555

examples/mem_scheduler/try_schedule_modules.py

@@ -0,0 +1,196 @@
+import shutil
+import sys
+
+from pathlib import Path
+from queue import Queue
+from typing import TYPE_CHECKING
+
+from memos.configs.mem_cube import GeneralMemCubeConfig
+from memos.configs.mem_os import MOSConfig
+from memos.configs.mem_scheduler import AuthConfig
+from memos.log import get_logger
+from memos.mem_cube.general import GeneralMemCube
+from memos.mem_scheduler.general_scheduler import GeneralScheduler
+from memos.mem_scheduler.modules.schemas import NOT_APPLICABLE_TYPE
+from memos.mem_scheduler.mos_for_test_scheduler import MOSForTestScheduler
+
+
+if TYPE_CHECKING:
+    from memos.mem_scheduler.modules.schemas import (
+        ScheduleLogForWebItem,
+    )
+
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent
+sys.path.insert(0, str(BASE_DIR))  # Enable execution from any working directory
+
+logger = get_logger(__name__)
+
+
+def init_task():
+    conversations = [
+        {
+            "role": "user",
+            "content": "I have two dogs - Max (golden retriever) and Bella (pug). We live in Seattle.",
+        },
+        {"role": "assistant", "content": "Great! Any special care for them?"},
+        {
+            "role": "user",
+            "content": "Max needs joint supplements. Actually, we're moving to Chicago next month.",
+        },
+        {
+            "role": "user",
+            "content": "Correction: Bella is 6, not 5. And she's allergic to chicken.",
+        },
+        {
+            "role": "user",
+            "content": "My partner's cat Whiskers visits weekends. Bella chases her sometimes.",
+        },
+    ]
+
+    questions = [
+        # 1. Basic factual recall (simple)
+        {
+            "question": "What breed is Max?",
+            "category": "Pet",
+            "expected": "golden retriever",
+            "difficulty": "easy",
+        },
+        # 2. Temporal context (medium)
+        {
+            "question": "Where will I live next month?",
+            "category": "Location",
+            "expected": "Chicago",
+            "difficulty": "medium",
+        },
+        # 3. Information correction (hard)
+        {
+            "question": "How old is Bella really?",
+            "category": "Pet",
+            "expected": "6",
+            "difficulty": "hard",
+            "hint": "User corrected the age later",
+        },
+        # 4. Relationship inference (harder)
+        {
+            "question": "Why might Whiskers be nervous around my pets?",
+            "category": "Behavior",
+            "expected": "Bella chases her sometimes",
+            "difficulty": "harder",
+        },
+        # 5. Combined medical info (hardest)
+        {
+            "question": "Which pets have health considerations?",
+            "category": "Health",
+            "expected": "Max needs joint supplements, Bella is allergic to chicken",
+            "difficulty": "hardest",
+            "requires": ["combining multiple facts", "ignoring outdated info"],
+        },
+    ]
+    return conversations, questions
+
+
+def show_web_logs(mem_scheduler: GeneralScheduler):
+    """Display all web log entries from the scheduler's log queue.
+
+    Args:
+        mem_scheduler: The scheduler instance containing web logs to display
+    """
+    if mem_scheduler._web_log_message_queue.empty():
+        print("Web log queue is currently empty.")
+        return
+
+    print("\n" + "=" * 50 + " WEB LOGS " + "=" * 50)
+
+    # Create a temporary queue to preserve the original queue contents
+    temp_queue = Queue()
+    log_count = 0
+
+    while not mem_scheduler._web_log_message_queue.empty():
+        log_item: ScheduleLogForWebItem = mem_scheduler._web_log_message_queue.get()
+        temp_queue.put(log_item)
+        log_count += 1
+
+        # Print log entry details
+        print(f"\nLog Entry #{log_count}:")
+        print(f'- "{log_item.label}" log: {log_item}')
+
+        print("-" * 50)
+
+    # Restore items back to the original queue
+    while not temp_queue.empty():
+        mem_scheduler._web_log_message_queue.put(temp_queue.get())
+
+    print(f"\nTotal {log_count} web log entries displayed.")
+    print("=" * 110 + "\n")
+
+
+if __name__ == "__main__":
+    # set up data
+    conversations, questions = init_task()
+
+    # set configs
+    mos_config = MOSConfig.from_yaml_file(
+        f"{BASE_DIR}/examples/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml"
+    )
+
+    mem_cube_config = GeneralMemCubeConfig.from_yaml_file(
+        f"{BASE_DIR}/examples/data/config/mem_scheduler/mem_cube_config.yaml"
+    )
+
+    # default local graphdb uri
+    if AuthConfig.default_config_exists():
+        auth_config = AuthConfig.from_local_yaml()
+
+        mos_config.mem_reader.config.llm.config.api_key = auth_config.openai.api_key
+        mos_config.mem_reader.config.llm.config.api_base = auth_config.openai.base_url
+
+        mem_cube_config.text_mem.config.graph_db.config.uri = auth_config.graph_db.uri
+
+    # Initialization
+    mos = MOSForTestScheduler(mos_config)
+
+    user_id = "user_1"
+    mos.create_user(user_id)
+
+    mem_cube_id = "mem_cube_5"
+    mem_cube_name_or_path = f"{BASE_DIR}/outputs/mem_scheduler/{user_id}/{mem_cube_id}"
+
+    if Path(mem_cube_name_or_path).exists():
+        shutil.rmtree(mem_cube_name_or_path)
+        print(f"{mem_cube_name_or_path} is not empty, and has been removed.")
+
+    mem_cube = GeneralMemCube(mem_cube_config)
+    mem_cube.dump(mem_cube_name_or_path)
+    mos.register_mem_cube(
+        mem_cube_name_or_path=mem_cube_name_or_path, mem_cube_id=mem_cube_id, user_id=user_id
+    )
+
+    mos.add(conversations, user_id=user_id, mem_cube_id=mem_cube_id)
+
+    for item in questions:
+        query = item["question"]
+
+        # test process_session_turn
+        mos.mem_scheduler.process_session_turn(
+            queries=[query],
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+            top_k=10,
+            query_history=None,
+        )
+
+        # test activation memory update
+        mos.mem_scheduler.update_activation_memory_periodically(
+            interval_seconds=0,
+            label=NOT_APPLICABLE_TYPE,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )
+
+    show_web_logs(mos.mem_scheduler)
+
+    mos.mem_scheduler.stop()

src/memos/mem_scheduler/base_scheduler.py

@@ -23,6 +23,7 @@
     DEFAULT_THREAD__POOL_MAX_WORKERS,
     LONG_TERM_MEMORY_TYPE,
     NOT_INITIALIZED,
+    PARAMETER_MEMORY_TYPE,
     QUERY_LABEL,
     TEXT_MEMORY_TYPE,
     USER_INPUT_TYPE,
@@ -31,7 +32,7 @@
     ScheduleMessageItem,
     TreeTextMemory_SEARCH_METHOD,
 )
-from memos.mem_scheduler.utils import normalize_name
+from memos.mem_scheduler.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
@@ -166,9 +167,10 @@ def _validate_message(self, message: ScheduleMessageItem, label: str):
     def update_activation_memory(
         self,
         new_memories: list[str | TextualMemoryItem],
+        label: str,
+        user_id: str,
+        mem_cube_id: str,
         mem_cube: GeneralMemCube,
-        user_id: str | None = None,
-        mem_cube_id: str | None = None,
     ) -> None:
         """
         Update activation memory by extracting KVCacheItems from new_memory (list of str),
@@ -220,6 +222,7 @@ def update_activation_memory(
             self.log_activation_memory_update(
                 original_text_memories=original_text_memories,
                 new_text_memories=new_text_memories,
+                label=label,
                 user_id=user_id,
                 mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
@@ -231,6 +234,7 @@ def update_activation_memory(
     def update_activation_memory_periodically(
         self,
         interval_seconds: int,
+        label: str,
         user_id: str,
         mem_cube_id: str,
         mem_cube: GeneralMemCube,
@@ -248,14 +252,21 @@ def update_activation_memory_periodically(
             )
 
             new_activation_memories = [
-                m.memory_text for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id]
+                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"
             )
 
-            self.update_activation_memory(new_memories=new_activation_memories, 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()
 
@@ -290,7 +301,7 @@ def _submit_web_logs(self, messages: ScheduleLogForWebItem | list[ScheduleLogFor
         for message in messages:
             self._web_log_message_queue.put(message)
             logger.info(f"Submitted Scheduling log for web: {message.log_content}")
-            logger.info(f"Submitted Scheduling log for web: {message.log_content}")
+
             if self.is_rabbitmq_connected():
                 logger.info("Submitted Scheduling log to rabbitmq")
                 self.rabbitmq_publish_message(message=message.to_dict())
@@ -300,6 +311,7 @@ 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,
@@ -318,16 +330,25 @@ def log_activation_memory_update(
 
         # recording messages
         for mem in added_memories:
-            log_message = self.create_autofilled_log_item(
+            log_message_a = self.create_autofilled_log_item(
                 log_content=mem,
-                label=QUERY_LABEL,
-                from_memory_type=WORKING_MEMORY_TYPE,
+                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,
             )
-            self._submit_web_logs(messages=log_message)
+            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."
@@ -343,7 +364,7 @@ def log_working_memory_replacement(
     ):
         """Log changes when working memory is replaced."""
         memory_type_map = {
-            normalize_name(text=m.memory): m.metadata.memory_type
+            transform_name_to_key(name=m.memory): m.metadata.memory_type
             for m in original_memory + new_memory
         }
 
@@ -359,7 +380,7 @@ def log_working_memory_replacement(
 
         # recording messages
         for mem in added_memories:
-            normalized_mem = normalize_name(text=mem)
+            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

src/memos/mem_scheduler/general_scheduler.py

@@ -86,6 +86,7 @@ def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                 if self.enable_act_memory_update:
                     self.update_activation_memory_periodically(
                         interval_seconds=self.monitor.act_mem_update_interval,
+                        label=ANSWER_LABEL,
                         user_id=user_id,
                         mem_cube_id=mem_cube_id,
                         mem_cube=messages[0].mem_cube,
@@ -121,6 +122,7 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                 if self.enable_act_memory_update:
                     self.update_activation_memory_periodically(
                         interval_seconds=self.monitor.act_mem_update_interval,
+                        label=ADD_LABEL,
                         user_id=user_id,
                         mem_cube_id=mem_cube_id,
                         mem_cube=messages[0].mem_cube,

src/memos/mem_scheduler/modules/monitor.py

@@ -148,13 +148,13 @@ def update_activation_memory_monitors(
         # === 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],
+            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 for m in top_k_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(