fix bugs: fix scheduler logging bug of working memory replacement

Author: tangg555

examples/mem_scheduler/schedule_chat_and_web.py

@@ -171,14 +171,8 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
     for item in questions:
         query = item["question"]
 
-        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,
-        )
+        response = mos.chat(query=query, user_id=user_id)
+        print(f"Query:\n {query}\n\nAnswer:\n {response}")
 
     show_web_logs(mos.mem_scheduler)
 

src/memos/mem_scheduler/base_scheduler.py

@@ -16,7 +16,6 @@
 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_HF_BACKEND,
     ACTIVATION_MEMORY_TYPE,
     ACTIVATION_MEMORY_VLLM_BACKEND,
     ADD_LABEL,
@@ -199,31 +198,25 @@ def update_activation_memory(
                     ]
                 )
             )
-            if self.act_mem_backend == ACTIVATION_MEMORY_HF_BACKEND:
-                # huggingface kv cache
-                original_cache_items: list[KVCacheItem] = act_mem.get_all()
-                pre_cache_item: KVCacheItem = original_cache_items[-1]
-                original_text_memories = pre_cache_item.records.text_memories
-                act_mem.delete_all()
-                cache_item: KVCacheItem = act_mem.extract(text_memory)
-                cache_item.records.text_memories = new_text_memories
-
-                act_mem.add(cache_item)
-                act_mem.dump(self.act_mem_dump_path)
-
-            elif self.act_mem_backend == ACTIVATION_MEMORY_VLLM_BACKEND:
-                # vllm kv cache
-                #  original_text_memories is empty
-                original_text_memories = []
-                self.log_activation_memory_update(
-                    original_text_memories=original_text_memories,
-                    new_text_memories=new_text_memories,
-                    user_id=user_id,
-                    mem_cube_id=mem_cube_id,
-                    mem_cube=mem_cube,
-                )
-            else:
-                raise NotImplementedError(self.act_mem_backend)
+
+            # huggingface or vllm kv cache
+            original_cache_items: list[KVCacheItem] = act_mem.get_all()
+            pre_cache_item: KVCacheItem = original_cache_items[-1]
+            original_text_memories = pre_cache_item.records.text_memories
+            act_mem.delete_all()
+            cache_item: KVCacheItem = act_mem.extract(text_memory)
+            cache_item.records.text_memories = new_text_memories
+
+            act_mem.add(cache_item)
+            act_mem.dump(self.act_mem_dump_path)
+
+            self.log_activation_memory_update(
+                original_text_memories=original_text_memories,
+                new_text_memories=new_text_memories,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+            )
 
         except Exception as e:
             logger.warning(f"MOS-based activation memory update failed: {e}")

src/memos/mem_scheduler/modules/retriever.py

@@ -27,7 +27,7 @@ def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
 
         # hyper-parameters
         self.filter_similarity_threshold = 0.75
-        self.filter_min_length_threshold = 5
+        self.filter_min_length_threshold = 6
 
         # log function callbacks
         self.log_working_memory_replacement = None
@@ -81,6 +81,14 @@ def filter_similar_memories(
             logging.warning("Received empty memories list - nothing to filter")
             return []
 
+        for idx in range(len(text_memories)):
+            if not isinstance(text_memories[idx], str):
+                logger.error(
+                    f"{text_memories[idx]} in memories is not a string,"
+                    f" and now has been transformed to be a string."
+                )
+                text_memories[idx] = str(text_memories[idx])
+
         try:
             # Step 1: Vectorize texts using TF-IDF
             vectorizer = TfidfVectorizer()
@@ -148,7 +156,7 @@ def filter_too_short_memories(
                 removal_indices.append(idx)
 
         if removal_indices:
-            logging.info(
+            logger.warning(
                 f"Removed {len(removal_indices)} short memories "
                 f"(shorter than {min_length_threshold} characters). "
                 f"Sample removed: {text_memories[removal_indices[0]][:50]}..."
@@ -178,21 +186,6 @@ def replace_working_memory(
             combined_text_memory = [normalize_name(text=m.memory) for m in combined_memory]
 
             # apply filters
-            # TODO: need to verify
-            """
-            Log Entry #8:
-- log: item_id='7b92fffa-7cb8-403e-9396-78f1d11e4f6e' user_id='user_1' mem_cube_id='mem_cube_5' label='query' from_memory_type='UserMemory' to_memory_type='WorkingMemory' log_content='The user is planning to move to Chicago next month, although the exact date of the move is unclear.' current_memory_sizes={'long_term_memory_size': 209, 'user_memory_size': 605, 'working_memory_size': 20, 'transformed_act_memory_size': -1} memory_capacities={'long_term_memory_capacity': 10000, 'user_memory_capacity': 10000, 'working_memory_capacity': 20, 'transformed_act_memory_capacity': -1} timestamp=datetime.datetime(2025, 7, 14, 21, 28, 0, 496940)
---------------------------------------------------
-
-Log Entry #9:
-- log: item_id='33acc8a6-9a77-4d9d-a7a0-970e5b74b3df' user_id='user_1' mem_cube_id='mem_cube_5' label='query' from_memory_type='UserMemory' to_memory_type='WorkingMemory' log_content='The user is planning to move to Chicago next month, which reflects a significant change in their living situation.' current_memory_sizes={'long_term_memory_size': 209, 'user_memory_size': 605, 'working_memory_size': 20, 'transformed_act_memory_size': -1} memory_capacities={'long_term_memory_capacity': 10000, 'user_memory_capacity': 10000, 'working_memory_capacity': 20, 'transformed_act_memory_capacity': -1} timestamp=datetime.datetime(2025, 7, 14, 21, 28, 0, 497328)
---------------------------------------------------
-
-Log Entry #10:
-- log: item_id='fa63eed9-e113-4a55-bb3f-4c2064d62d9d' user_id='user_1' mem_cube_id='mem_cube_5' label='query' from_memory_type='UserMemory' to_memory_type='WorkingMemory' log_content='The user is planning to move to Chicago in the upcoming month, indicating a significant change in their living situation.' current_memory_sizes={'long_term_memory_size': 209, 'user_memory_size': 605, 'working_memory_size': 20, 'transformed_act_memory_size': -1} memory_capacities={'long_term_memory_capacity': 10000, 'user_memory_capacity': 10000, 'working_memory_capacity': 20, 'transformed_act_memory_capacity': -1} timestamp=datetime.datetime(2025, 7, 14, 21, 28, 0, 497694)
---------------------------------------------------
-
-            """
             filtered_combined_text_memory = self.filter_similar_memories(
                 text_memories=combined_text_memory,
                 similarity_threshold=self.filter_similarity_threshold,
@@ -221,24 +214,23 @@ def replace_working_memory(
 
             memories_with_new_order = []
             for text in text_memories_with_new_order:
-                text = normalize_name(text=text)
+                normalized_text = normalize_name(text=text)
                 if text in memory_map:
-                    memories_with_new_order.append(memory_map[text])
+                    memories_with_new_order.append(memory_map[normalized_text])
                 else:
                     logger.warning(
                         f"Memory text not found in memory map. text: {text}; keys of memory_map: {memory_map.keys()}"
                     )
 
-            text_mem_base.replace_working_memory(memories_with_new_order[:top_k])
-            memories_with_new_order = memories_with_new_order[:top_k]
+            text_mem_base.replace_working_memory(memories_with_new_order)
             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=memories_with_new_order,
                 user_id=user_id,
                 mem_cube_id=mem_cube_id,
-                new_memory=new_memory,
                 mem_cube=mem_cube,
             )
         else:

src/memos/mem_scheduler/modules/schemas.py

@@ -40,8 +40,8 @@
 WORKING_MEMORY_TYPE = "WorkingMemory"
 TEXT_MEMORY_TYPE = "TextMemory"
 ACTIVATION_MEMORY_TYPE = "ActivationMemory"
+PARAMETER_MEMORY_TYPE = "ParameterMemory"
 USER_INPUT_TYPE = "UserInput"
-NotAvailabel = "NotAvailabel"
 
 # monitors
 MONITOR_WORKING_MEMORY_TYPE = "MonitorWorkingMemoryType"

src/memos/mem_scheduler/utils.py

@@ -38,6 +38,8 @@ def normalize_name(text):
     # Optional: Collapse multiple whitespaces into single space
     normalized = " ".join(normalized.split())
 
+    normalized = normalized.lower()
+
     return normalized
 
 

tests/mem_scheduler/test_retriever.py

@@ -74,16 +74,12 @@ def test_filter_similar_memories_no_duplicates(self):
     def test_filter_similar_memories_with_duplicates(self):
         """Test filter_similar_memories with duplicate memories."""
         memories = [
-            "This is a memory about dogs and cats",
-            "This is a memory about dogs and cats and birds",
-            "This is a completely different memory",
-            "This is a memory about dogs and cats",  # Exact duplicate
-            "This is a memory about DOGS and CATS",  # Near duplicate with different case
+            "The user is planning to move to Chicago next month, although the exact date of the move is unclear.",
+            "The user is planning to move to Chicago next month, which reflects a significant change in their living situation.",
+            "The user is planning to move to Chicago in the upcoming month, indicating a significant change in their living situation.",
         ]
-
-        result = self.retriever.filter_similar_memories(memories, similarity_threshold=0.8)
+        result = self.retriever.filter_similar_memories(memories, similarity_threshold=0.75)
         self.assertLess(len(result), len(memories))
-        self.assertIn("This is a completely different memory", result)
 
         # Verify logging was called for removed items
         self.assertGreater(self.mock_logger_info.call_count, 0)