Skip to content

Commit d01d759

Browse files
committed
fix bugs: fix scheduler logging bug of working memory replacement
1 parent 49af60e commit d01d759

File tree

6 files changed

+42
-65
lines changed

6 files changed

+42
-65
lines changed

examples/mem_scheduler/schedule_chat_and_web.py

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -171,14 +171,8 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
171171
for item in questions:
172172
query = item["question"]
173173

174-
mos.mem_scheduler.process_session_turn(
175-
queries=[query],
176-
user_id=user_id,
177-
mem_cube_id=mem_cube_id,
178-
mem_cube=mem_cube,
179-
top_k=10,
180-
query_history=None,
181-
)
174+
response = mos.chat(query=query, user_id=user_id)
175+
print(f"Query:\n {query}\n\nAnswer:\n {response}")
182176

183177
show_web_logs(mos.mem_scheduler)
184178

src/memos/mem_scheduler/base_scheduler.py

Lines changed: 19 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
from memos.mem_scheduler.modules.redis_service import RedisSchedulerModule
1717
from memos.mem_scheduler.modules.retriever import SchedulerRetriever
1818
from memos.mem_scheduler.modules.schemas import (
19-
ACTIVATION_MEMORY_HF_BACKEND,
2019
ACTIVATION_MEMORY_TYPE,
2120
ACTIVATION_MEMORY_VLLM_BACKEND,
2221
ADD_LABEL,
@@ -199,31 +198,25 @@ def update_activation_memory(
199198
]
200199
)
201200
)
202-
if self.act_mem_backend == ACTIVATION_MEMORY_HF_BACKEND:
203-
# huggingface kv cache
204-
original_cache_items: list[KVCacheItem] = act_mem.get_all()
205-
pre_cache_item: KVCacheItem = original_cache_items[-1]
206-
original_text_memories = pre_cache_item.records.text_memories
207-
act_mem.delete_all()
208-
cache_item: KVCacheItem = act_mem.extract(text_memory)
209-
cache_item.records.text_memories = new_text_memories
210-
211-
act_mem.add(cache_item)
212-
act_mem.dump(self.act_mem_dump_path)
213-
214-
elif self.act_mem_backend == ACTIVATION_MEMORY_VLLM_BACKEND:
215-
# vllm kv cache
216-
# original_text_memories is empty
217-
original_text_memories = []
218-
self.log_activation_memory_update(
219-
original_text_memories=original_text_memories,
220-
new_text_memories=new_text_memories,
221-
user_id=user_id,
222-
mem_cube_id=mem_cube_id,
223-
mem_cube=mem_cube,
224-
)
225-
else:
226-
raise NotImplementedError(self.act_mem_backend)
201+
202+
# huggingface or vllm kv cache
203+
original_cache_items: list[KVCacheItem] = act_mem.get_all()
204+
pre_cache_item: KVCacheItem = original_cache_items[-1]
205+
original_text_memories = pre_cache_item.records.text_memories
206+
act_mem.delete_all()
207+
cache_item: KVCacheItem = act_mem.extract(text_memory)
208+
cache_item.records.text_memories = new_text_memories
209+
210+
act_mem.add(cache_item)
211+
act_mem.dump(self.act_mem_dump_path)
212+
213+
self.log_activation_memory_update(
214+
original_text_memories=original_text_memories,
215+
new_text_memories=new_text_memories,
216+
user_id=user_id,
217+
mem_cube_id=mem_cube_id,
218+
mem_cube=mem_cube,
219+
)
227220

228221
except Exception as e:
229222
logger.warning(f"MOS-based activation memory update failed: {e}")

src/memos/mem_scheduler/modules/retriever.py

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
2727

2828
# hyper-parameters
2929
self.filter_similarity_threshold = 0.75
30-
self.filter_min_length_threshold = 5
30+
self.filter_min_length_threshold = 6
3131

3232
# log function callbacks
3333
self.log_working_memory_replacement = None
@@ -81,6 +81,14 @@ def filter_similar_memories(
8181
logging.warning("Received empty memories list - nothing to filter")
8282
return []
8383

84+
for idx in range(len(text_memories)):
85+
if not isinstance(text_memories[idx], str):
86+
logger.error(
87+
f"{text_memories[idx]} in memories is not a string,"
88+
f" and now has been transformed to be a string."
89+
)
90+
text_memories[idx] = str(text_memories[idx])
91+
8492
try:
8593
# Step 1: Vectorize texts using TF-IDF
8694
vectorizer = TfidfVectorizer()
@@ -148,7 +156,7 @@ def filter_too_short_memories(
148156
removal_indices.append(idx)
149157

150158
if removal_indices:
151-
logging.info(
159+
logger.warning(
152160
f"Removed {len(removal_indices)} short memories "
153161
f"(shorter than {min_length_threshold} characters). "
154162
f"Sample removed: {text_memories[removal_indices[0]][:50]}..."
@@ -178,21 +186,6 @@ def replace_working_memory(
178186
combined_text_memory = [normalize_name(text=m.memory) for m in combined_memory]
179187

180188
# apply filters
181-
# TODO: need to verify
182-
"""
183-
Log Entry #8:
184-
- 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)
185-
--------------------------------------------------
186-
187-
Log Entry #9:
188-
- 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)
189-
--------------------------------------------------
190-
191-
Log Entry #10:
192-
- 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)
193-
--------------------------------------------------
194-
195-
"""
196189
filtered_combined_text_memory = self.filter_similar_memories(
197190
text_memories=combined_text_memory,
198191
similarity_threshold=self.filter_similarity_threshold,
@@ -221,24 +214,23 @@ def replace_working_memory(
221214

222215
memories_with_new_order = []
223216
for text in text_memories_with_new_order:
224-
text = normalize_name(text=text)
217+
normalized_text = normalize_name(text=text)
225218
if text in memory_map:
226-
memories_with_new_order.append(memory_map[text])
219+
memories_with_new_order.append(memory_map[normalized_text])
227220
else:
228221
logger.warning(
229222
f"Memory text not found in memory map. text: {text}; keys of memory_map: {memory_map.keys()}"
230223
)
231224

232-
text_mem_base.replace_working_memory(memories_with_new_order[:top_k])
233-
memories_with_new_order = memories_with_new_order[:top_k]
225+
text_mem_base.replace_working_memory(memories_with_new_order)
234226
logger.info(
235227
f"The working memory has been replaced with {len(memories_with_new_order)} new memories."
236228
)
237229
self.log_working_memory_replacement(
238230
original_memory=original_memory,
231+
new_memory=memories_with_new_order,
239232
user_id=user_id,
240233
mem_cube_id=mem_cube_id,
241-
new_memory=new_memory,
242234
mem_cube=mem_cube,
243235
)
244236
else:

src/memos/mem_scheduler/modules/schemas.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,8 @@
4040
WORKING_MEMORY_TYPE = "WorkingMemory"
4141
TEXT_MEMORY_TYPE = "TextMemory"
4242
ACTIVATION_MEMORY_TYPE = "ActivationMemory"
43+
PARAMETER_MEMORY_TYPE = "ParameterMemory"
4344
USER_INPUT_TYPE = "UserInput"
44-
NotAvailabel = "NotAvailabel"
4545

4646
# monitors
4747
MONITOR_WORKING_MEMORY_TYPE = "MonitorWorkingMemoryType"

src/memos/mem_scheduler/utils.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ def normalize_name(text):
3838
# Optional: Collapse multiple whitespaces into single space
3939
normalized = " ".join(normalized.split())
4040

41+
normalized = normalized.lower()
42+
4143
return normalized
4244

4345

tests/mem_scheduler/test_retriever.py

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,16 +74,12 @@ def test_filter_similar_memories_no_duplicates(self):
7474
def test_filter_similar_memories_with_duplicates(self):
7575
"""Test filter_similar_memories with duplicate memories."""
7676
memories = [
77-
"This is a memory about dogs and cats",
78-
"This is a memory about dogs and cats and birds",
79-
"This is a completely different memory",
80-
"This is a memory about dogs and cats", # Exact duplicate
81-
"This is a memory about DOGS and CATS", # Near duplicate with different case
77+
"The user is planning to move to Chicago next month, although the exact date of the move is unclear.",
78+
"The user is planning to move to Chicago next month, which reflects a significant change in their living situation.",
79+
"The user is planning to move to Chicago in the upcoming month, indicating a significant change in their living situation.",
8280
]
83-
84-
result = self.retriever.filter_similar_memories(memories, similarity_threshold=0.8)
81+
result = self.retriever.filter_similar_memories(memories, similarity_threshold=0.75)
8582
self.assertLess(len(result), len(memories))
86-
self.assertIn("This is a completely different memory", result)
8783

8884
# Verify logging was called for removed items
8985
self.assertGreater(self.mock_logger_info.call_count, 0)

0 commit comments

Comments
 (0)