refactor: improve architecture and configurations for memory scheduler #142

tangg555 · 2025-07-22T14:04:53Z

Description

Summary: Refactor memory scheduler with improved architecture and configurations

Reorganized scheduler modules into more fine-grained schemas (general, message, monitor)
Reorganized scheduler utils into a categorized folder (filter_utils and misc_utils)
Enhanced configuration parameters (increased context_window_size to 10, shorter trigger time intervals)
Improved memory monitoring with better scoring and retention logic
Added comprehensive logging for memory transformations
Refactored retriever with better filtering and reranking
Updated examples to use new schemas and show web logs
Added validation and type safety throughout the scheduler
Fix small bugs in dispatcher and retriever modules

Fix: #5

Docs Issue/PR: (docs-issue-or-pr-link)

Reviewer: @fridayL @Ki-Seki @CaralHsi

Checklist:

I have performed a self-review of my own code | 我已自行检查了自己的代码
I have commented my code in hard-to-understand areas | 我已在难以理解的地方对代码进行了注释
I have added tests that prove my fix is effective or that my feature works | 我已添加测试以证明我的修复有效或功能正常
I have created related documentation issue/PR in MemOS-Docs (if applicable) | 我已在 MemOS-Docs 中创建了相关的文档 issue/PR（如果适用）
I have linked the issue to this PR (if applicable) | 我已将 issue 链接到此 PR（如果适用）
I have mentioned the person who will review this PR | 我已提及将审查此 PR 的人

…dispatch

…ors and web logs, and refactor the code structure

…mples; some bugs due to refactoring have been fixed

Ki-Seki · 2025-07-23T02:29:28Z

@tangg555

It is recommended to improve the PR description, especially the checklist. The current PR modifications may involve updates to the MemOS-Docs documentation, so it is advised to add links to related issues/PRs in the PR description. A well-filled example of a PR title & description can be found in PR #120.

建议完善PR描述，尤其是其中的检查清单。当前的PR修改可能涉及MemOS-Docs文档的更新，建议在PR描述中添加相关issue/PR的链接。一个填写完善的PR title&description示例如 PR #120.

Copilot

Pull Request Overview

This PR refactors the memory scheduler with improved architecture and configurations. The refactoring reorganizes scheduler modules into clear schemas (general, message, monitor), enhances configuration parameters, and improves memory monitoring with better scoring and retention logic.

Key changes include:

Split large schema file into specialized modules for better maintainability
Enhanced memory monitoring with keyword-based scoring and importance ranking
Improved filter utilities with better similarity detection and length validation
Updated examples to demonstrate new schema organization and web logging capabilities

Reviewed Changes

Copilot reviewed 35 out of 37 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
src/memos/mem_scheduler/schemas/	New schema organization split into general, message, and monitor modules
src/memos/mem_scheduler/utils/	Refactored utilities split into filter and misc utilities
src/memos/mem_scheduler/modules/	Updated modules to use new schema imports and enhanced functionality
tests/mem_scheduler/	Updated test imports and removed obsolete test methods
examples/	Updated examples to use new schemas and demonstrate web logs

Copilot · 2025-07-23T02:34:53Z

src/memos/mem_scheduler/utils/misc_utils.py

+
+def parse_yaml(yaml_file):
+    yaml_path = Path(yaml_file)
+    yaml_path = Path(yaml_file)


Duplicate line detected - yaml_path = Path(yaml_file) is assigned twice consecutively (lines 18-19). Remove the duplicate assignment.

Suggested change

yaml_path = Path(yaml_file)

Copilot · 2025-07-23T02:34:54Z

src/memos/mem_scheduler/modules/retriever.py

    def search(
        self, query: str, mem_cube: GeneralMemCube, top_k: int, method=TreeTextMemory_SEARCH_METHOD
-    ):
+    ) -> list[TextualMemoryItem]:


The return type annotation is inconsistent with the actual return statement. The method can return both list[TextualMemoryItem] and an empty list [] in the exception case, but the type annotation only specifies the success case.

Copilot · 2025-07-23T02:34:54Z

src/memos/mem_scheduler/modules/retriever.py

-            logger.error("memory_base is not supported")

-        return memories_with_new_order
+        return memories_with_new_order, success_flag


The return type annotation specifies list[TextualMemoryItem] | None but the method always returns a tuple (list[TextualMemoryItem], bool). Update the return type annotation to match the actual return value.

Copilot · 2025-07-23T02:34:54Z

src/memos/mem_scheduler/schemas/monitor_schemas.py

+    def generate_mapping_key(cls, v, values):  # noqa: N805
+        if v is None and "memory_text" in values:
+            return transform_name_to_key(values["memory_text"])


The field validator accesses values dictionary without proper validation. Use info.data instead of values for Pydantic v2 compatibility and add proper error handling for missing keys.

Suggested change

def generate_mapping_key(cls, v, values): # noqa: N805

if v is None and "memory_text" in values:

return transform_name_to_key(values["memory_text"])

def generate_mapping_key(cls, v, info): # noqa: N805

memory_text = info.data.get("memory_text")

if v is None and memory_text is not None:

return transform_name_to_key(memory_text)

Copilot · 2025-07-23T02:34:54Z

src/memos/mem_scheduler/schemas/monitor_schemas.py

+            f"Removed: {len(memories_to_remove)} (excluding top {partial_retention_number} by recording_count)"
+        )
+
+        return self.memories


The method update_memories has inconsistent return type. The return type annotation specifies MemoryMonitorItem but the method returns self.memories which is list[MemoryMonitorItem]. Either update the return type annotation or change the return statement.

Copilot · 2025-07-23T02:34:55Z

src/memos/mem_scheduler/modules/misc.py

+        data_copy = data.copy()  # Avoid modifying original dictionary
+        if "timestamp" in data_copy and isinstance(data_copy["timestamp"], str):
+            try:
+                data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])
+            except ValueError:
+                # Handle invalid time formats - adjust as needed (e.g., log warning or set to None)
+                data_copy["timestamp"] = None
+
+        return cls(**data_copy)


[nitpick] Creating a copy of the entire dictionary for every conversion may be inefficient for large dictionaries. Consider modifying only the timestamp field in-place or using a more targeted approach.

Suggested change

data_copy = data.copy() # Avoid modifying original dictionary

if "timestamp" in data_copy and isinstance(data_copy["timestamp"], str):

try:

data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])

except ValueError:

# Handle invalid time formats - adjust as needed (e.g., log warning or set to None)

data_copy["timestamp"] = None

return cls(**data_copy)

if "timestamp" in data and isinstance(data["timestamp"], str):

try:

data["timestamp"] = datetime.fromisoformat(data["timestamp"])

except ValueError:

# Handle invalid time formats - adjust as needed (e.g., log warning or set to None)

data["timestamp"] = None

return cls(**data)

Copilot · 2025-07-23T02:34:55Z

src/memos/mem_scheduler/modules/rabbitmq_service.py

+            print(f"Received message: {body.decode()}\n")
+            self.rabbitmq_message_cache.put({"properties": properties, "body": body})
+            print(f"message delivery_tag: {method.delivery_tag}\n")


Using print() statements for logging in production code is not recommended. Use the logger instance that's already available in the class for consistent logging behavior.

Suggested change

print(f"Received message: {body.decode()}\n")

self.rabbitmq_message_cache.put({"properties": properties, "body": body})

print(f"message delivery_tag: {method.delivery_tag}\n")

logger.debug(f"Received message: {body.decode()}")

self.rabbitmq_message_cache.put({"properties": properties, "body": body})

logger.debug(f"Message delivery_tag: {method.delivery_tag}")

Copilot · 2025-07-23T02:34:55Z

src/memos/mem_scheduler/general_scheduler.py

+
+        if (not intent_result["trigger_retrieval"]) and (not time_trigger_flag):
+            logger.info(f"Query schedule not triggered. Intent_result: {intent_result}")
+            return


The method process_session_turn has inconsistent return behavior. It can return None (line 222) or a tuple (line 251), but the return type annotation specifies only the tuple. Update the return type to include None or ensure consistent return values.

Ki-Seki · 2025-07-23T02:35:41Z

Looks good overall. Might want to incorporate Copilot's suggested changes @tangg555.

Also, I suggest having @fridayL and @J1awei-Yang review this as well.

fridayL

Too many changes, they're all closely related to the demo, need merge and deploy online .

MemTensor#142) * fix bugs: modify mos_for_test_scheduler.py and fix bugs of scheduler dispatch * feat & fix bugs & refactor: mem scheudler add more functions of monitors and web logs, and refactor the code structure * new feat & fix bugs: mem_scheduler update the logging feature and examples; some bugs due to refactoring have been fixed * fix ruff problem * fix UTC problem

tangg555 added 4 commits July 22, 2025 21:51

fix bugs: modify mos_for_test_scheduler.py and fix bugs of scheduler …

8cf59ea

…dispatch

feat & fix bugs & refactor: mem scheudler add more functions of monit…

9ed747e

…ors and web logs, and refactor the code structure

new feat & fix bugs: mem_scheduler update the logging feature and exa…

76ca88b

…mples; some bugs due to refactoring have been fixed

fix ruff problem

2c0c25b

tangg555 requested a review from fridayL July 22, 2025 14:04

tangg555 self-assigned this Jul 22, 2025

tangg555 added bug Something isn't working enhancement New feature or request [MemScheduler] labels Jul 22, 2025

fix UTC problem

8ded2cf

fridayL requested review from J1awei-Yang and Ki-Seki July 23, 2025 02:04

Ki-Seki changed the title ~~Dev~~ refactor: improve architecture and configurations for memory scheduler Jul 23, 2025

Ki-Seki requested a review from Copilot July 23, 2025 02:33

Copilot AI reviewed Jul 23, 2025

View reviewed changes

fridayL approved these changes Jul 23, 2025

View reviewed changes

fridayL merged commit c133eff into MemTensor:dev Jul 23, 2025
40 checks passed

-    def generate_mapping_key(cls, v, values):  # noqa: N805
-        if v is None and "memory_text" in values:
-            return transform_name_to_key(values["memory_text"])
+    def generate_mapping_key(cls, v, info):  # noqa: N805
+        memory_text = info.data.get("memory_text")
+        if v is None and memory_text is not None:
+            return transform_name_to_key(memory_text)

refactor: improve architecture and configurations for memory scheduler #142

refactor: improve architecture and configurations for memory scheduler #142

Uh oh!

Conversation

tangg555 commented Jul 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist:

Uh oh!

Ki-Seki commented Jul 23, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Ki-Seki commented Jul 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fridayL left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tangg555 commented Jul 22, 2025 •

edited

Loading

Ki-Seki commented Jul 23, 2025 •

edited

Loading