Skip to content

Fix None issue in reorganizer queue #125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CaralHsi merged 2 commits into from
Jul 19, 2025
Merged

Fix None issue in reorganizer queue #125

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1c0701a
fix: revise None type in QueueMessage to op=end
J1awei-Yang Jul 18, 2025
4d93f44
Merge branch 'dev' into fix/reorganizer_queue
CaralHsi Jul 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions src/memos/memories/textual/tree_text_memory/organize/reorganizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
class QueueMessage:
def __init__(
self,
op: Literal["add", "remove", "merge", "update"],
op: Literal["add", "remove", "merge", "update", "end"],
# `str` for node and edge IDs, `GraphDBNode` and `GraphDBEdge` for actual objects
before_node: list[str] | list[GraphDBNode] | None = None,
before_edge: list[str] | list[GraphDBEdge] | None = None,
Expand All @@ -48,7 +48,7 @@ def __str__(self) -> str:
return f"QueueMessage(op={self.op}, before_node={self.before_node if self.before_node is None else len(self.before_node)}, after_node={self.after_node if self.after_node is None else len(self.after_node)})"

def __lt__(self, other: "QueueMessage") -> bool:
op_priority = {"add": 2, "remove": 2, "merge": 1}
op_priority = {"add": 2, "remove": 2, "merge": 1, "end": 0}

Copilot AI Jul 18, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The op_priority dictionary is missing the 'update' operation which is defined in the Literal type. This will cause a KeyError when comparing QueueMessage objects with op='update'.

Suggested change
op_priority = {"add": 2, "remove": 2, "merge": 1, "end": 0}
op_priority = {"add": 2, "remove": 2, "merge": 1, "update": 1, "end": 0}

Copilot uses AI. Check for mistakes.
return op_priority[self.op] < op_priority[other.op]


Expand Down Expand Up @@ -103,7 +103,7 @@ def wait_until_current_task_done(self):
def _run_message_consumer_loop(self):
while True:
message = self.queue.get()
if message is None:
if message.op == "end":
break

try:
Expand Down Expand Up @@ -140,7 +140,7 @@ def stop(self):
if not self.is_reorganize:
return

self.add_message(None)
self.add_message(QueueMessage(op="end"))
self.thread.join()
logger.info("Reorganize thread stopped.")
self._stop_scheduler = True
Expand All @@ -158,9 +158,6 @@ def handle_message(self, message: QueueMessage):

def handle_add(self, message: QueueMessage):
logger.debug(f"Handling add operation: {str(message)[:500]}")
assert message.before_node is None and message.before_edge is None, (
"Before node and edge should be None for `add` operation."
)
# ———————— 1. check for conflicts ————————
added_node = message.after_node[0]
conflicts = self.conflict.detect(added_node, scope=added_node.metadata.memory_type)
Expand Down Expand Up @@ -586,7 +583,7 @@ def _link_cluster_nodes(self, parent_node: GraphDBNode, child_nodes: list[GraphD

def _preprocess_message(self, message: QueueMessage) -> bool:
message = self._convert_id_to_node(message)
if None in message.after_node:
if message.after_node is None or None in message.after_node:
logger.debug(
f"Found non-existent node in after_node in message: {message}, skip this message."
)
Expand Down