Fix None issue in reorganizer queue (#125)

fix: revise None type in QueueMessage to op=end

Co-authored-by: CaralHsi <[email protected]>

Author: J1awei-Yang

src/memos/memories/textual/tree_text_memory/organize/reorganizer.py

@@ -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,
@@ -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}
         return op_priority[self.op] < op_priority[other.op]
 
 
@@ -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:
@@ -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
@@ -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)
@@ -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."
             )