From e25c7c93fe4cae0a7614bf03ce72ce87de7c3f46 Mon Sep 17 00:00:00 2001 From: lxobr <122801072+lxobr@users.noreply.github.com> Date: Thu, 20 Feb 2025 12:52:34 +0100 Subject: [PATCH] fix: correctly add nodes to chunks [COG-1370] (#568) ## Description - Fix expand_with_nodes_and_edges to correctly add nodes to chunks ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin ## Summary by CodeRabbit - **Refactor** - Enhanced the internal processing for data associations to ensure more reliable and consistent handling of connections. - Streamlined the logic to better manage edge cases, improving overall stability and error handling. --- .../graph/utils/expand_with_nodes_and_edges.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py index 23a6f4f63..8ebc5788a 100644 --- a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py +++ b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py @@ -19,6 +19,7 @@ def expand_with_nodes_and_edges( existing_edges_map = {} added_nodes_map = {} + relationships = [] for index, data_chunk in enumerate(data_chunks): @@ -53,12 +54,14 @@ def expand_with_nodes_and_edges( description=node.description, ) - if data_chunk.contains is None: - data_chunk.contains = [] - - data_chunk.contains.append(entity_node) - added_nodes_map[f"{str(node_id)}_entity"] = entity_node + else: + entity_node = added_nodes_map[f"{str(node_id)}_entity"] + + if data_chunk.contains is None: + data_chunk.contains = [] + + data_chunk.contains.append(entity_node) # Add relationship that came from graphs. for edge in graph.edges: