feat: return ids when add memory (#119)

* feat: return ids when add memory

* feat: add example for how to use tree.add

* tests: add tests for tree memory add

* tests: add tests for tree memory add

Author: CaralHsi

examples/core_memories/tree_textual_memory.py

@@ -186,7 +186,9 @@ def embed_memory_item(memory: str) -> list[float]:
 memory = reader.get_memory(scene_data, type="chat", info={"user_id": "1234", "session_id": "2222"})
 
 for m_list in memory:
-    my_tree_textual_memory.add(m_list)
+    added_ids = my_tree_textual_memory.add(m_list)
+    for i, id in enumerate(added_ids):
+        print(f"{i}'th added result is:" + my_tree_textual_memory.get(id).memory)
     my_tree_textual_memory.memory_manager.wait_reorganizer()
 
 time.sleep(60)
@@ -217,7 +219,7 @@ def embed_memory_item(memory: str) -> list[float]:
 doc_memory = reader.get_memory(doc_paths, "doc", info={"user_id": "1111", "session_id": "2222"})
 
 for m_list in doc_memory:
-    my_tree_textual_memory.add(m_list)
+    added_ids = my_tree_textual_memory.add(m_list)
     my_tree_textual_memory.memory_manager.wait_reorganizer()
 
 results = my_tree_textual_memory.search(

src/memos/memories/textual/tree.py

@@ -57,7 +57,7 @@ def __init__(self, config: TreeTextMemoryConfig):
         else:
             logger.info("No internet retriever configured")
 
-    def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> None:
+    def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> list[str]:
         """Add memories.
         Args:
             memories: List of TextualMemoryItem objects or dictionaries to add.
@@ -67,7 +67,7 @@ def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> None:
             plan = plan_memory_operations(memory_items, metadata, self.graph_store)
             execute_plan(memory_items, metadata, plan, self.graph_store)
         """
-        self.memory_manager.add(memories)
+        return self.memory_manager.add(memories)
 
     def replace_working_memory(self, memories: list[TextualMemoryItem]) -> None:
         self.memory_manager.replace_working_memory(memories)

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

@@ -49,15 +49,18 @@ def __init__(
         )
         self._merged_threshold = merged_threshold
 
-    def add(self, memories: list[TextualMemoryItem]) -> None:
+    def add(self, memories: list[TextualMemoryItem]) -> list[str]:
         """
         Add new memories in parallel to different memory types (WorkingMemory, LongTermMemory, UserMemory).
         """
+        added_ids: list[str] = []
+
         with ThreadPoolExecutor(max_workers=8) as executor:
-            futures = [executor.submit(self._process_memory, memory) for memory in memories]
+            futures = {executor.submit(self._process_memory, m): m for m in memories}
             for future in as_completed(futures):
                 try:
-                    future.result()
+                    ids = future.result()
+                    added_ids.extend(ids)
                 except Exception as e:
                     logger.exception("Memory processing error: ", exc_info=e)
 
@@ -72,6 +75,7 @@ def add(self, memories: list[TextualMemoryItem]) -> None:
         )
 
         self._refresh_memory_size()
+        return added_ids
 
     def replace_working_memory(self, memories: list[TextualMemoryItem]) -> None:
         """
@@ -113,17 +117,23 @@ def _process_memory(self, memory: TextualMemoryItem):
         Process and add memory to different memory types (WorkingMemory, LongTermMemory, UserMemory).
         This method runs asynchronously to process each memory item.
         """
+        ids = []
+
         # Add to WorkingMemory
-        self._add_memory_to_db(memory, "WorkingMemory")
+        working_id = self._add_memory_to_db(memory, "WorkingMemory")
+        ids.append(working_id)
 
         # Add to LongTermMemory and UserMemory
         if memory.metadata.memory_type in ["LongTermMemory", "UserMemory"]:
-            self._add_to_graph_memory(
+            added_id = self._add_to_graph_memory(
                 memory=memory,
                 memory_type=memory.metadata.memory_type,
             )
+            ids.append(added_id)
+
+        return ids
 
-    def _add_memory_to_db(self, memory: TextualMemoryItem, memory_type: str):
+    def _add_memory_to_db(self, memory: TextualMemoryItem, memory_type: str) -> str:
         """
         Add a single memory item to the graph store, with FIFO logic for WorkingMemory.
         """
@@ -135,6 +145,7 @@ def _add_memory_to_db(self, memory: TextualMemoryItem, memory_type: str):
 
         # Insert node into graph
         self.graph_store.add_node(working_memory.id, working_memory.memory, metadata)
+        return working_memory.id
 
     def _add_to_graph_memory(self, memory: TextualMemoryItem, memory_type: str):
         """
@@ -159,7 +170,7 @@ def _add_to_graph_memory(self, memory: TextualMemoryItem, memory_type: str):
         )
 
         if similar_nodes and similar_nodes[0]["score"] > self._merged_threshold:
-            self._merge(memory, similar_nodes)
+            return self._merge(memory, similar_nodes)
         else:
             node_id = str(uuid.uuid4())
             # Step 2: Add new node to graph
@@ -172,8 +183,9 @@ def _add_to_graph_memory(self, memory: TextualMemoryItem, memory_type: str):
                     after_node=[node_id],
                 )
             )
+            return node_id
 
-    def _merge(self, source_node: TextualMemoryItem, similar_nodes: list[dict]) -> None:
+    def _merge(self, source_node: TextualMemoryItem, similar_nodes: list[dict]) -> str:
         """
         TODO: Add node traceability support by optionally preserving source nodes and linking them with MERGED_FROM edges.
 
@@ -200,7 +212,9 @@ def _merge(self, source_node: TextualMemoryItem, similar_nodes: list[dict]) -> N
         merged_background = f"{original_meta.background}\n⟵MERGED⟶\n{source_meta.background}"
         merged_embedding = self.embedder.embed([merged_text])[0]
 
-        merged_confidence = float((original_meta.confidence + source_meta.confidence) / 2)
+        original_conf = original_meta.confidence or 0.0
+        source_conf = source_meta.confidence or 0.0
+        merged_confidence = float((original_conf + source_conf) / 2)
         merged_usage = list(set((original_meta.usage or []) + (source_meta.usage or [])))
 
         # Create new merged node
@@ -243,6 +257,7 @@ def _merge(self, source_node: TextualMemoryItem, similar_nodes: list[dict]) -> N
                 after_node=[merged_id],
             )
         )
+        return merged_id
 
     def _inherit_edges(self, from_id: str, to_id: str) -> None:
         """

tests/memories/textual/test_tree.py

@@ -138,3 +138,27 @@ def test_drop_creates_backup_and_cleans(mock_tree_text_memory):
     mock_tree_text_memory.dump.assert_called_once()
     mock_tree_text_memory._cleanup_old_backups.assert_called_once()
     mock_tree_text_memory.graph_store.drop_database.assert_called_once()
+
+
+def test_add_returns_ids(mock_tree_text_memory):
+    # Mock the memory_manager.add to return specific IDs
+    dummy_ids = ["id1", "id2"]
+    mock_tree_text_memory.memory_manager.add = MagicMock(return_value=dummy_ids)
+
+    mock_items = [
+        TextualMemoryItem(
+            id=str(uuid.uuid4()),
+            memory="Memory 1",
+            metadata=TreeNodeTextualMemoryMetadata(updated_at=None),
+        ),
+        TextualMemoryItem(
+            id=str(uuid.uuid4()),
+            memory="Memory 2",
+            metadata=TreeNodeTextualMemoryMetadata(updated_at=None),
+        ),
+    ]
+
+    result = mock_tree_text_memory.add(mock_items)
+
+    assert result == dummy_ids
+    mock_tree_text_memory.memory_manager.add.assert_called_once_with(mock_items)

tests/memories/textual/test_tree_manager.py

@@ -145,3 +145,14 @@ def test_ensure_structure_path_reuses_existing(memory_manager, mock_graph_store)
     meta = TreeNodeTextualMemoryMetadata(key="hobby")
     node_id = memory_manager._ensure_structure_path("UserMemory", meta)
     assert node_id == "existing_node_id"
+
+
+def test_add_returns_written_node_ids(memory_manager):
+    memory = TextualMemoryItem(
+        memory="test memory",
+        metadata=TreeNodeTextualMemoryMetadata(embedding=[0.1] * 5, memory_type="UserMemory"),
+    )
+    ids = memory_manager.add([memory])
+    assert isinstance(ids, list)
+    assert all(isinstance(i, str) for i in ids)
+    assert len(ids) > 0