Skip to content

Commit ec1e867

Browse files
committed
feat: update filter mem
1 parent 7bb26a9 commit ec1e867

File tree

1 file changed

+28
-5
lines changed

1 file changed

+28
-5
lines changed

src/memos/mem_os/product.py

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -704,16 +704,39 @@ def run_async_in_thread():
704704
thread.start()
705705

706706
def _filter_memories_by_threshold(
707-
self, memories: list[TextualMemoryItem], threshold: float = 0.30, min_num: int = 3
707+
self,
708+
memories: list[TextualMemoryItem],
709+
threshold: float = 0.30,
710+
min_num: int = 3,
711+
memory_type: Literal["OuterMemory"] = "OuterMemory"
708712
) -> list[TextualMemoryItem]:
709713
"""
710-
Filter memories by threshold.
714+
Filter memories by threshold and type, at least min_num memories for Non-OuterMemory.
715+
Args:
716+
memories: list[TextualMemoryItem],
717+
threshold: float,
718+
min_num: int,
719+
memory_type: Literal["OuterMemory"],
720+
Returns:
721+
list[TextualMemoryItem]
711722
"""
712723
sorted_memories = sorted(memories, key=lambda m: m.metadata.relativity, reverse=True)
713-
filtered = [m for m in sorted_memories if m.metadata.relativity >= threshold]
724+
filtered_person = [m for m in memories if m.metadata.memory_type != memory_type]
725+
filtered_outer = [m for m in memories if m.metadata.memory_type == memory_type]
726+
filtered = []
727+
per_memory_count = 0
728+
for m in sorted_memories:
729+
if m.metadata.relativity >= threshold:
730+
if m.metadata.memory_type != memory_type:
731+
per_memory_count += 1
732+
filtered.append(m)
714733
if len(filtered) < min_num:
715-
filtered = sorted_memories[:min_num]
716-
return filtered
734+
filtered = filtered_person[:min_num] + filtered_outer[:min_num]
735+
else:
736+
if len(per_memory_count) < min_num:
737+
filtered += filtered_person[per_memory_count: min_num]
738+
filtered_memory = sorted(filtered, key=lambda m: m.metadata.relativity, reverse=True)
739+
return filtered_memory
717740

718741
def register_mem_cube(
719742
self,

0 commit comments

Comments
 (0)