diff --git a/src/memos/mem_feedback/feedback.py b/src/memos/mem_feedback/feedback.py index 28f495afb..18ebb0b49 100644 --- a/src/memos/mem_feedback/feedback.py +++ b/src/memos/mem_feedback/feedback.py @@ -657,11 +657,6 @@ def _info_comparison(self, memory: TextualMemoryItem, _info: dict, include_keys: def _retrieve(self, query: str, info=None, top_k=20, user_name=None): """Retrieve memory items""" - def check_has_edges(mem_item: TextualMemoryItem) -> tuple[TextualMemoryItem, bool]: - """Check if a memory item has edges.""" - edges = self.searcher.graph_store.get_edges(mem_item.id, user_name=user_name) - return (mem_item, len(edges) == 0) - logger.info(f"[feedback _retrieve] query: {query}, user_name: {user_name}") text_mems = self.searcher.search( query=query, @@ -686,22 +681,16 @@ def check_has_edges(mem_item: TextualMemoryItem) -> tuple[TextualMemoryItem, boo pref_mems = [item[0] for item in pref_mems if float(item[1]) > 0.01] text_mems.extend(pref_mems) - # Memory with edges is not modified by feedback - retrieved_mems = [] - with ContextThreadPoolExecutor(max_workers=10) as executor: - futures = {executor.submit(check_has_edges, item): item for item in text_mems} - for future in concurrent.futures.as_completed(futures): - try: - mem_item, has_no_edges = future.result() - if has_no_edges: - retrieved_mems.append(mem_item) - except Exception as e: - logger.error(f"[0107 Feedback Core: _retrieve] Error checking edges: {e}") + retrieved_mems = [ + item + for item in text_mems + if not (getattr(item.metadata, "knowledge_material_ids", None) or []) + ] if len(retrieved_mems) < len(text_mems): logger.info( f"[0107 Feedback Core: _retrieve] {len(text_mems) - len(retrieved_mems)} " - f"text memories are not modified by feedback due to edges." + f"knowledge memories are not modified by feedback." ) return retrieved_mems diff --git a/src/memos/mem_reader/multi_modal_struct.py b/src/memos/mem_reader/multi_modal_struct.py index c46331f1d..eb226c27d 100644 --- a/src/memos/mem_reader/multi_modal_struct.py +++ b/src/memos/mem_reader/multi_modal_struct.py @@ -769,30 +769,52 @@ def _process_one_item( except Exception as e: logger.error(f"[MultiModalFine] parse error: {e}") - # save rawfile node - if self.save_rawfile and prompt_type == "doc" and len(fine_items) > 0: - rawfile_chunk = mem_str - file_info = fine_items[0].metadata.sources[0].file_info - source = self.multi_modal_parser.file_content_parser.create_source( - message={"file": file_info}, - info=info_per_item, - chunk_index=chunk_idx, - chunk_total=total_chunks, - chunk_content="", - ) - rawfile_node = self._make_memory_item( - value=rawfile_chunk, - info=info_per_item, - memory_type="RawFileMemory", - tags=[ - "mode:fine", - "multimodal:file", - f"chunk:{chunk_idx + 1}/{total_chunks}", - ], - sources=[source], - ) - rawfile_node.metadata.summary_ids = [mem_node.id for mem_node in fine_items] - fine_items.append(rawfile_node) + # Save the original file chunk independently from prompt selection. + file_source = next( + ( + source + for source in sources + if ( + source.get("type") + if isinstance(source, dict) + else getattr(source, "type", None) + ) + == "file" + ), + None, + ) + if self.save_rawfile and file_source is not None and fine_items: + if isinstance(file_source, dict): + file_info = file_source.get("file_info") or {} + file_content = file_source.get("content") + else: + file_info = file_source.file_info or {} + file_content = file_source.content + + rawfile_chunk = file_content or file_info.get("file_data", "") + if rawfile_chunk: + source = self.multi_modal_parser.file_content_parser.create_source( + message={"file": file_info}, + info=info_per_item, + chunk_index=chunk_idx, + chunk_total=total_chunks, + chunk_content=rawfile_chunk, + ) + rawfile_node = self._make_memory_item( + value=rawfile_chunk, + info=info_per_item, + memory_type="RawFileMemory", + tags=[ + "mode:fine", + "multimodal:file", + f"chunk:{chunk_idx + 1}/{total_chunks}", + ], + sources=[source], + ) + rawfile_node.metadata.summary_ids = [mem_node.id for mem_node in fine_items] + fine_items.append(rawfile_node) + else: + logger.warning("[RawFile] File source has no content; skipping raw-file node") enriched_items = trigger_hook( H.MEMORY_ITEMS_AFTER_FINE_EXTRACT, items=fine_items, diff --git a/src/memos/mem_scheduler/task_schedule_modules/handlers/mem_read_handler.py b/src/memos/mem_scheduler/task_schedule_modules/handlers/mem_read_handler.py index a954794b7..50031b7aa 100644 --- a/src/memos/mem_scheduler/task_schedule_modules/handlers/mem_read_handler.py +++ b/src/memos/mem_scheduler/task_schedule_modules/handlers/mem_read_handler.py @@ -209,7 +209,7 @@ def _process_memories_with_reader( for memory in flattened_memories if memory.metadata.memory_type == "RawFileMemory" ] - text_mem.add_rawfile_nodes_n_edges( + text_mem.add_rawfile_nodes( raw_file_mem_group, enhanced_mem_ids, user_id=user_id, diff --git a/src/memos/memories/textual/item.py b/src/memos/memories/textual/item.py index 23474cc20..1ba148d8a 100644 --- a/src/memos/memories/textual/item.py +++ b/src/memos/memories/textual/item.py @@ -212,6 +212,14 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata): default_factory=list, description="The ids of the files associated with the memory.", ) + knowledge_material_ids: list[str] = Field( + default_factory=list, + description="RawFileMemory ids that support this knowledge summary.", + ) + knowledge_summary_ids: list[str] = Field( + default_factory=list, + description="Knowledge summary ids extracted from this raw file memory.", + ) @field_validator("sources", mode="before") @classmethod diff --git a/src/memos/memories/textual/tree.py b/src/memos/memories/textual/tree.py index 426a2359b..c9d74d43f 100644 --- a/src/memos/memories/textual/tree.py +++ b/src/memos/memories/textual/tree.py @@ -399,10 +399,13 @@ def get_all( ) return graph_output - def delete(self, memory_ids: list[str], user_name: str | None = None) -> None: + def delete(self, memory_ids: str | list[str], user_name: str | None = None) -> None: """Hard delete: permanently remove nodes and their edges from the graph.""" if not memory_ids: return + if isinstance(memory_ids, str): + memory_ids = [memory_ids] + self._unlink_knowledge_relations(memory_ids, user_name=user_name) for mid in memory_ids: try: self.graph_store.delete_node(mid, user_name=user_name) @@ -412,6 +415,7 @@ def delete(self, memory_ids: list[str], user_name: str | None = None) -> None: def delete_by_memory_ids(self, memory_ids: list[str]) -> None: """Delete memories by memory_ids.""" try: + self._unlink_knowledge_relations(memory_ids) self.graph_store.delete_node_by_prams(memory_ids=memory_ids) except Exception as e: logger.error(f"An error occurred while deleting memories by memory_ids: {e}") @@ -432,10 +436,71 @@ def delete_by_filter( filter: dict | None = None, ) -> None: """Delete memories by filter.""" + if file_ids: + user_names = writable_cube_ids or [None] + for user_name in user_names: + memory_ids: set[str] = set() + for file_id in file_ids: + memory_ids.update( + self.graph_store.get_by_metadata( + [{"field": "file_ids", "op": "contains", "value": file_id}], + user_name=user_name, + ) + ) + self._unlink_knowledge_relations(list(memory_ids), user_name=user_name) self.graph_store.delete_node_by_prams( writable_cube_ids=writable_cube_ids, file_ids=file_ids, filter=filter ) + @staticmethod + def _node_metadata(node: dict[str, Any] | TextualMemoryItem) -> dict[str, Any]: + if isinstance(node, TextualMemoryItem): + return node.metadata.model_dump() + return node.get("metadata") or {} + + def _get_node_map( + self, memory_ids: list[str] | set[str], user_name: str | None = None + ) -> dict[str, dict[str, Any] | TextualMemoryItem]: + if not memory_ids: + return {} + nodes = self.graph_store.get_nodes(ids=list(memory_ids), user_name=user_name) or [] + return { + node.id if isinstance(node, TextualMemoryItem) else node["id"]: node for node in nodes + } + + def _unlink_knowledge_relations( + self, memory_ids: list[str], user_name: str | None = None + ) -> None: + """Remove reverse knowledge ids before deleting or deactivating nodes.""" + deleting_ids = set(memory_ids) + deleting_nodes = self._get_node_map(deleting_ids, user_name=user_name) + removals: dict[str, dict[str, set[str]]] = {} + + for memory_id, node in deleting_nodes.items(): + metadata = self._node_metadata(node) + for material_id in metadata.get("knowledge_material_ids") or []: + if material_id not in deleting_ids: + removals.setdefault(material_id, {}).setdefault( + "knowledge_summary_ids", set() + ).add(memory_id) + for summary_id in metadata.get("knowledge_summary_ids") or []: + if summary_id not in deleting_ids: + removals.setdefault(summary_id, {}).setdefault( + "knowledge_material_ids", set() + ).add(memory_id) + + related_nodes = self._get_node_map(set(removals), user_name=user_name) + for related_id, fields in removals.items(): + node = related_nodes.get(related_id) + if node is None: + continue + metadata = self._node_metadata(node) + updates = { + field: [value for value in metadata.get(field) or [] if value not in removed_ids] + for field, removed_ids in fields.items() + } + self.graph_store.update_node(related_id, updates, user_name=user_name) + def load(self, dir: str, user_name: str | None = None) -> None: try: memory_file = os.path.join(dir, self.config.memory_filename) @@ -593,6 +658,109 @@ def add_rawfile_nodes_n_edges( f"[RawFile] Time taken to add edges: {end_time - start_time} seconds for {len(from_ids)} edges" ) + def add_rawfile_nodes( + self, + raw_file_mem_group: list[TextualMemoryItem], + mem_ids: list[str], + user_id: str | None = None, + user_name: str | None = None, + ) -> None: + """Persist raw-file nodes and their knowledge relationships without graph edges.""" + valid_summary_ids = set(mem_ids) + summary_to_material_ids: dict[str, list[str]] = {} + + for raw_file_mem in raw_file_mem_group: + summary_ids = list( + dict.fromkeys( + summary_id + for summary_id in getattr(raw_file_mem.metadata, "summary_ids", []) or [] + if summary_id in valid_summary_ids + ) + ) + raw_file_mem.metadata.knowledge_summary_ids = summary_ids + for summary_id in summary_ids: + summary_to_material_ids.setdefault(summary_id, []).append(raw_file_mem.id) + + self.add(raw_file_mem_group, user_name=user_name) + summary_nodes = self._get_node_map(set(summary_to_material_ids), user_name=user_name) + + merged_from_by_summary: dict[str, list[str]] = {} + old_summary_ids: set[str] = set() + for summary_id, node in summary_nodes.items(): + info = self._node_metadata(node).get("info") or {} + merged_from = info.get("merged_from") or [] + if isinstance(merged_from, str): + merged_from = [merged_from] + merged_from_by_summary[summary_id] = list(dict.fromkeys(merged_from)) + old_summary_ids.update(merged_from) + + old_summary_nodes = self._get_node_map(old_summary_ids, user_name=user_name) + material_relinks: dict[str, dict[str, Any]] = {} + + for summary_id, direct_material_ids in summary_to_material_ids.items(): + summary_node = summary_nodes.get(summary_id) + if summary_node is None: + logger.warning( + "Knowledge summary %s was not found after raw-file write", summary_id + ) + continue + + metadata = self._node_metadata(summary_node) + merged_from = merged_from_by_summary.get(summary_id, []) + inherited_material_ids: list[str] = [] + for old_summary_id in merged_from: + old_summary = old_summary_nodes.get(old_summary_id) + if old_summary is not None: + inherited_material_ids.extend( + self._node_metadata(old_summary).get("knowledge_material_ids") or [] + ) + + material_ids = list( + dict.fromkeys( + [ + *(metadata.get("knowledge_material_ids") or []), + *direct_material_ids, + *inherited_material_ids, + ] + ) + ) + self.graph_store.update_node( + summary_id, + {"knowledge_material_ids": material_ids}, + user_name=user_name, + ) + + for material_id in inherited_material_ids: + relink = material_relinks.setdefault(material_id, {"remove": set(), "add": []}) + relink["remove"].update(merged_from) + relink["add"].append(summary_id) + + material_nodes = self._get_node_map(set(material_relinks), user_name=user_name) + for material_id, relink in material_relinks.items(): + material_node = material_nodes.get(material_id) + if material_node is None: + continue + current_summary_ids = ( + self._node_metadata(material_node).get("knowledge_summary_ids") or [] + ) + updated_summary_ids = [ + summary_id + for summary_id in current_summary_ids + if summary_id not in relink["remove"] + ] + updated_summary_ids.extend(relink["add"]) + self.graph_store.update_node( + material_id, + {"knowledge_summary_ids": list(dict.fromkeys(updated_summary_ids))}, + user_name=user_name, + ) + + logger.info( + "[RawFile] Added %s chunks with ID mappings for user %s", + len(raw_file_mem_group), + user_id, + ) + def add_graph_edges( self, from_ids: list[str], to_ids: list[str], types: list[str], user_name: str | None = None ) -> None: @@ -624,6 +792,7 @@ def soft_delete( user_name: str, evolve_to_ids: list[str] | None = None, ) -> None: + self._unlink_knowledge_relations(memory_ids, user_name=user_name) # for ruff check... if not evolve_to_ids: update_fields = {"status": "deleted"} diff --git a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py index 5c9cce78e..b44d28218 100644 --- a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py +++ b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py @@ -3,8 +3,6 @@ import re import traceback -from concurrent.futures import as_completed - from memos.context.context import ContextThreadPoolExecutor from memos.embedders.factory import OllamaEmbedder from memos.graph_dbs.factory import Neo4jGraphDB @@ -1282,7 +1280,7 @@ def _sort_and_trim( @timed def _deduplicate_rawfile_results(self, results, user_name: str | None = None): """ - Deduplicate rawfile related memories by edge + Deduplicate raw-file results using their persisted summary ids. """ if not results: return results @@ -1292,29 +1290,10 @@ def _deduplicate_rawfile_results(self, results, user_name: str | None = None): if not rawfile_items: return results - with ContextThreadPoolExecutor(max_workers=min(len(rawfile_items), 10)) as executor: - futures = [ - executor.submit( - self.graph_store.get_edges, - rawfile_item.id, - type="SUMMARY", - direction="OUTGOING", - user_name=user_name, - ) - for rawfile_item in rawfile_items - ] - for future in as_completed(futures): - try: - edges = future.result() - for edge in edges: - summary_target_id = edge.get("to") - if summary_target_id: - summary_ids_to_remove.add(summary_target_id) - logger.debug( - f"[DEDUP] Marking summary node {summary_target_id} for removal (pointed by RawFileMemory)" - ) - except Exception as e: - logger.warning(f"[DEDUP] Failed to get summary target ids: {e}") + for rawfile_item in rawfile_items: + summary_ids_to_remove.update( + getattr(rawfile_item.metadata, "knowledge_summary_ids", None) or [] + ) filtered_results = [] for item in results: diff --git a/src/memos/multi_mem_cube/single_cube.py b/src/memos/multi_mem_cube/single_cube.py index 27399259a..0b9829b7f 100644 --- a/src/memos/multi_mem_cube/single_cube.py +++ b/src/memos/multi_mem_cube/single_cube.py @@ -441,22 +441,35 @@ def extract_edge_info(edges_info: list[dict], neighbor_relativity: float): ) return edge_mems + def get_neighbor(chunk_target_id: str | None, neighbor_relativity: float): + if not chunk_target_id: + return None + item_neighbor = self.searcher.graph_store.get_node(chunk_target_id, user_name=user_name) + if not item_neighbor: + return None + item_neighbor_mem = TextualMemoryItem(**item_neighbor) + item_neighbor_mem.metadata.relativity = neighbor_relativity + self.logger.info("Add neighbor chunk by id: %s", chunk_target_id) + return item_neighbor_mem + final_items = [] if neighbor_discovery: for item in search_results: if item.metadata.memory_type == "RawFileMemory": neighbor_relativity = item.metadata.relativity * 0.8 - preceding_info = self.searcher.graph_store.get_edges( - item.id, type="PRECEDING", direction="OUTGOING", user_name=user_name + preceding = get_neighbor( + getattr(item.metadata, "preceding_id", None), neighbor_relativity ) - final_items.extend(extract_edge_info(preceding_info, neighbor_relativity)) + if preceding is not None: + final_items.append(preceding) final_items.append(item) - following_info = self.searcher.graph_store.get_edges( - item.id, type="FOLLOWING", direction="OUTGOING", user_name=user_name + following = get_neighbor( + getattr(item.metadata, "following_id", None), neighbor_relativity ) - final_items.extend(extract_edge_info(following_info, neighbor_relativity)) + if following is not None: + final_items.append(following) else: final_items.append(item) @@ -748,7 +761,7 @@ def _process_text_mem( for memory in flattened_local if memory.metadata.memory_type == "RawFileMemory" ] - self.naive_mem_cube.text_mem.add_rawfile_nodes_n_edges( + self.naive_mem_cube.text_mem.add_rawfile_nodes( raw_file_mem_group, mem_ids_local, user_id=add_req.user_id, diff --git a/tests/mem_cube/test_single_cube_knowledge_neighbors.py b/tests/mem_cube/test_single_cube_knowledge_neighbors.py new file mode 100644 index 000000000..5cb433d77 --- /dev/null +++ b/tests/mem_cube/test_single_cube_knowledge_neighbors.py @@ -0,0 +1,50 @@ +import uuid + +from unittest.mock import Mock + +from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata + + +def make_raw(item_id: str, relativity: float, **extra) -> TextualMemoryItem: + return TextualMemoryItem( + id=item_id, + memory=item_id, + metadata=TreeNodeTextualMemoryMetadata( + memory_type="RawFileMemory", + relativity=relativity, + **extra, + ), + ) + + +def test_postformat_discovers_rawfile_neighbors_by_id_without_edges(): + import memos.api.handlers # noqa: F401 - initialize package before known circular import + + from memos.multi_mem_cube.single_cube import SingleCubeView + + previous = make_raw(str(uuid.uuid4()), 0.1) + following = make_raw(str(uuid.uuid4()), 0.1) + current = make_raw(str(uuid.uuid4()), 1.0, preceding_id=previous.id, following_id=following.id) + searcher = Mock() + neighbors = { + previous.id: previous.model_dump(), + following.id: following.model_dump(), + } + searcher.graph_store.get_node.side_effect = lambda node_id, **_: neighbors.get(node_id) + view = SingleCubeView( + cube_id="cube-1", + naive_mem_cube=Mock(), + mem_reader=Mock(), + mem_scheduler=Mock(), + logger=Mock(), + searcher=searcher, + ) + + result = view._postformat_memories( + [current], "cube-1", include_embedding=False, neighbor_discovery=True + ) + + assert [item["id"] for item in result] == [previous.id, current.id, following.id] + assert result[0]["metadata"]["relativity"] == 0.8 + assert result[2]["metadata"]["relativity"] == 0.8 + searcher.graph_store.get_edges.assert_not_called() diff --git a/tests/mem_feedback/test_feedback.py b/tests/mem_feedback/test_feedback.py index bbf30ae94..3f86ac482 100644 --- a/tests/mem_feedback/test_feedback.py +++ b/tests/mem_feedback/test_feedback.py @@ -1,3 +1,5 @@ +import uuid + from unittest.mock import Mock from memos.mem_feedback.feedback import MemFeedback @@ -139,3 +141,21 @@ def test_update_judgement_prompt_omits_memory_text_from_response_schema(): assert '"reason"' in output_section assert '"text"' not in output_section assert '"old_memory"' not in output_section + + +def test_retrieve_locks_only_knowledge_summaries_without_edge_queries(): + feedback = MemFeedback.__new__(MemFeedback) + feedback.pref_feedback = False + feedback.searcher = Mock() + normal = TextualMemoryItem(id=str(uuid.uuid4()), memory="normal") + knowledge = TextualMemoryItem( + id=str(uuid.uuid4()), + memory="knowledge", + metadata={"knowledge_material_ids": [str(uuid.uuid4())]}, + ) + feedback.searcher.search.return_value = [(normal, 0.9), (knowledge, 0.8)] + + result = feedback._retrieve("query", user_name="cube-1") + + assert result == [normal] + feedback.searcher.graph_store.get_edges.assert_not_called() diff --git a/tests/mem_reader/test_multi_modal_rawfile.py b/tests/mem_reader/test_multi_modal_rawfile.py new file mode 100644 index 000000000..f583b7173 --- /dev/null +++ b/tests/mem_reader/test_multi_modal_rawfile.py @@ -0,0 +1,89 @@ +from unittest.mock import MagicMock, patch + +from memos.mem_reader.multi_modal_struct import MultiModalStructMemReader +from memos.memories.textual.item import ( + SourceMessage, + TextualMemoryItem, + TreeNodeTextualMemoryMetadata, +) + + +def test_mixed_file_source_creates_rawfile_without_changing_prompt_type(): + reader = MultiModalStructMemReader.__new__(MultiModalStructMemReader) + reader.save_rawfile = True + reader.memory_version_switch = "off" + reader.embedder = MagicMock() + reader.embedder.embed.return_value = [[0.1]] + reader.multi_modal_parser = MagicMock() + + file_info = { + "file_id": "file-1", + "filename": "knowledge.txt", + "file_data": "Complete file content", + } + file_source = SourceMessage( + type="file", + role=None, + content="Current file chunk", + file_info=file_info, + ) + fast_item = TextualMemoryItem( + memory="Please parse this file\nCurrent file chunk\nRecord it completely", + metadata=TreeNodeTextualMemoryMetadata( + user_id="user-1", + session_id="session-1", + memory_type="LongTermMemory", + sources=[ + SourceMessage(type="chat", role="user", content="Please parse this file"), + file_source, + SourceMessage(type="chat", role="user", content="Record it completely"), + ], + file_ids=[file_info["file_id"]], + ), + ) + reader.multi_modal_parser.file_content_parser.create_source.return_value = SourceMessage( + type="file", + content=file_source.content, + doc_path=file_info["filename"], + ) + + llm_response = { + "memory list": [ + { + "key": "Knowledge summary", + "memory_type": "LongTermMemory", + "value": "Summarized knowledge", + "tags": ["knowledge"], + } + ], + "summary": "Summary background", + } + + with ( + patch.object(reader, "_get_llm_response", return_value=llm_response) as get_llm_response, + patch.object( + reader, + "_get_maybe_merged_memory", + side_effect=lambda extracted_memory_dict, **kwargs: extracted_memory_dict, + ), + patch( + "memos.mem_reader.multi_modal_struct.trigger_hook", + side_effect=lambda _hook, **kwargs: kwargs["items"], + ), + ): + result = reader._process_string_fine( + [fast_item], + {"user_id": "user-1", "session_id": "session-1"}, + ) + + summaries = [item for item in result if item.metadata.memory_type != "RawFileMemory"] + rawfiles = [item for item in result if item.metadata.memory_type == "RawFileMemory"] + + assert get_llm_response.call_args.args[3] == "chat" + assert len(summaries) == 1 + assert len(rawfiles) == 1 + assert rawfiles[0].memory == file_source.content + assert rawfiles[0].metadata.summary_ids == [summaries[0].id] + assert reader.multi_modal_parser.file_content_parser.create_source.call_args.kwargs[ + "message" + ] == {"file": file_info} diff --git a/tests/memories/textual/test_tree_knowledge_relations.py b/tests/memories/textual/test_tree_knowledge_relations.py new file mode 100644 index 000000000..79bd80743 --- /dev/null +++ b/tests/memories/textual/test_tree_knowledge_relations.py @@ -0,0 +1,182 @@ +import uuid + +from unittest.mock import Mock + +from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata +from memos.memories.textual.tree import TreeTextMemory + + +class FakeGraphStore: + def __init__(self, items: list[TextualMemoryItem]) -> None: + self.nodes = {item.id: item.model_dump() for item in items} + self.edge_calls: list[tuple] = [] + self.deleted_by_params: list[dict] = [] + + def get_nodes(self, ids, user_name=None): + return [self.nodes[node_id] for node_id in ids if node_id in self.nodes] + + def update_node(self, id, fields, user_name=None): + self.nodes[id]["metadata"].update(fields) + + def delete_node(self, id, user_name=None): + self.nodes.pop(id, None) + + def delete_node_by_prams(self, **kwargs): + self.deleted_by_params.append(kwargs) + for node_id in kwargs.get("memory_ids") or []: + self.nodes.pop(node_id, None) + + def get_by_metadata(self, filters, user_name=None, status=None): + matched = [] + for node_id, node in self.nodes.items(): + metadata = node["metadata"] + if all( + ( + item["value"] in (metadata.get(item["field"]) or []) + if item.get("op") == "contains" + else metadata.get(item["field"]) == item["value"] + ) + for item in filters + ): + matched.append(node_id) + return matched + + def add_edge(self, *args, **kwargs): + self.edge_calls.append((args, kwargs)) + + def get_edges(self, *args, **kwargs): + self.edge_calls.append((args, kwargs)) + return [] + + +def _test_id(value: str) -> str: + try: + return str(uuid.UUID(value)) + except ValueError: + return str(uuid.uuid5(uuid.NAMESPACE_DNS, value)) + + +def make_item( + item_id: str, + memory_type: str, + *, + knowledge_material_ids: list[str] | None = None, + knowledge_summary_ids: list[str] | None = None, + info: dict | None = None, + file_ids: list[str] | None = None, + summary_ids: list[str] | None = None, +) -> TextualMemoryItem: + metadata = TreeNodeTextualMemoryMetadata( + memory_type=memory_type, + knowledge_material_ids=[_test_id(value) for value in knowledge_material_ids or []], + knowledge_summary_ids=[_test_id(value) for value in knowledge_summary_ids or []], + info=( + { + **info, + "merged_from": [_test_id(value) for value in info.get("merged_from", [])], + } + if info and "merged_from" in info + else info + ), + file_ids=file_ids or [], + ) + if summary_ids is not None: + metadata.summary_ids = [_test_id(value) for value in summary_ids] + return TextualMemoryItem(id=_test_id(item_id), memory=item_id, metadata=metadata) + + +def make_tree(items: list[TextualMemoryItem]) -> tuple[TreeTextMemory, FakeGraphStore]: + tree = TreeTextMemory.__new__(TreeTextMemory) + graph_store = FakeGraphStore(items) + tree.graph_store = graph_store + + def add(memories, user_name=None, **kwargs): + for memory in memories: + graph_store.nodes[memory.id] = memory.model_dump() + return [memory.id for memory in memories] + + tree.add = Mock(side_effect=add) + return tree, graph_store + + +def test_add_rawfile_nodes_persists_bidirectional_ids_without_edges(): + summary = make_item("summary-1", "LongTermMemory") + material = make_item("material-1", "RawFileMemory", summary_ids=[summary.id]) + tree, graph_store = make_tree([summary]) + + tree.add_rawfile_nodes([material], [summary.id], user_name="cube-1") + + assert graph_store.nodes[material.id]["metadata"]["knowledge_summary_ids"] == [summary.id] + assert graph_store.nodes[summary.id]["metadata"]["knowledge_material_ids"] == [material.id] + assert graph_store.edge_calls == [] + + +def test_add_rawfile_nodes_reconnects_materials_when_summary_is_merged(): + old_material = make_item("material-old", "RawFileMemory", knowledge_summary_ids=["summary-old"]) + old_summary = make_item( + "summary-old", "LongTermMemory", knowledge_material_ids=[old_material.id] + ) + new_summary = make_item("summary-new", "LongTermMemory", info={"merged_from": [old_summary.id]}) + new_material = make_item("material-new", "RawFileMemory", summary_ids=[new_summary.id]) + tree, graph_store = make_tree([old_material, old_summary, new_summary]) + + tree.add_rawfile_nodes([new_material], [new_summary.id], user_name="cube-1") + + assert graph_store.nodes[new_summary.id]["metadata"]["knowledge_material_ids"] == [ + new_material.id, + old_material.id, + ] + assert graph_store.nodes[old_material.id]["metadata"]["knowledge_summary_ids"] == [ + new_summary.id + ] + + +def test_delete_summary_unlinks_material_without_deleting_it(): + material = make_item("material-1", "RawFileMemory", knowledge_summary_ids=["summary-1"]) + summary = make_item("summary-1", "LongTermMemory", knowledge_material_ids=[material.id]) + tree, graph_store = make_tree([material, summary]) + + tree.delete_by_memory_ids([summary.id]) + + assert summary.id not in graph_store.nodes + assert graph_store.nodes[material.id]["metadata"]["knowledge_summary_ids"] == [] + + +def test_delete_material_unlinks_summary_without_deleting_it(): + material = make_item("material-1", "RawFileMemory", knowledge_summary_ids=["summary-1"]) + summary = make_item("summary-1", "LongTermMemory", knowledge_material_ids=[material.id]) + tree, graph_store = make_tree([material, summary]) + + tree.delete_by_memory_ids([material.id]) + + assert material.id not in graph_store.nodes + assert graph_store.nodes[summary.id]["metadata"]["knowledge_material_ids"] == [] + + +def test_file_delete_unlinks_knowledge_relations_before_bulk_delete(): + material = make_item( + "material-1", + "RawFileMemory", + knowledge_summary_ids=["summary-1"], + file_ids=["file-1"], + ) + summary = make_item("summary-1", "LongTermMemory", knowledge_material_ids=[material.id]) + tree, graph_store = make_tree([material, summary]) + + tree.delete_by_filter(writable_cube_ids=["cube-1"], file_ids=["file-1"]) + + assert graph_store.nodes[summary.id]["metadata"]["knowledge_material_ids"] == [] + assert graph_store.deleted_by_params == [ + {"writable_cube_ids": ["cube-1"], "file_ids": ["file-1"], "filter": None} + ] + + +def test_soft_delete_unlinks_summary_before_marking_it_deleted(): + material = make_item("material-1", "RawFileMemory", knowledge_summary_ids=["summary-1"]) + summary = make_item("summary-1", "LongTermMemory", knowledge_material_ids=[material.id]) + tree, graph_store = make_tree([material, summary]) + + tree.soft_delete([summary.id], user_name="cube-1") + + assert graph_store.nodes[material.id]["metadata"]["knowledge_summary_ids"] == [] + assert graph_store.nodes[summary.id]["metadata"]["status"] == "deleted" diff --git a/tests/memories/textual/test_tree_searcher.py b/tests/memories/textual/test_tree_searcher.py index b79958ca1..fe40cbb7d 100644 --- a/tests/memories/textual/test_tree_searcher.py +++ b/tests/memories/textual/test_tree_searcher.py @@ -1,3 +1,5 @@ +import uuid + from unittest.mock import MagicMock import pytest @@ -143,3 +145,26 @@ def test_searcher_respects_memory_type(mock_searcher): ) # WorkingMemory triggers only once path A assert mock_searcher.graph_retriever.retrieve.call_args[1]["memory_scope"] == "WorkingMemory" + + +def test_deduplicate_rawfile_results_uses_knowledge_summary_ids(mock_searcher): + summary_id = str(uuid.uuid4()) + material_id = str(uuid.uuid4()) + summary = TextualMemoryItem( + id=summary_id, + memory="summary", + metadata=TreeNodeTextualMemoryMetadata(memory_type="LongTermMemory"), + ) + raw_file = TextualMemoryItem( + id=material_id, + memory="material", + metadata=TreeNodeTextualMemoryMetadata( + memory_type="RawFileMemory", + knowledge_summary_ids=[summary.id], + ), + ) + + result = mock_searcher._deduplicate_rawfile_results([summary, raw_file], user_name="cube-1") + + assert result == [raw_file] + mock_searcher.graph_store.get_edges.assert_not_called()