From 6b6cc0f1d45863328a7b2b150f4234910f622df0 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 9 Jan 2025 16:06:26 +0100 Subject: [PATCH] fix: Add fix for accessing dictionary elements that don't exits Using get for the text key instead of direct access to handle situation if the text key doesn't exist --- cognee/tasks/completion/graph_query_completion.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cognee/tasks/completion/graph_query_completion.py b/cognee/tasks/completion/graph_query_completion.py index d579aea4d..b130d4f7b 100644 --- a/cognee/tasks/completion/graph_query_completion.py +++ b/cognee/tasks/completion/graph_query_completion.py @@ -5,11 +5,11 @@ from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search -def retrieved_edges_to_string(retrieved_edges): +def retrieved_edges_to_string(retrieved_edges: list) -> str: edge_strings = [] for edge in retrieved_edges: - node1_string = edge.node1.attributes["text"] or edge.node1.attributes.get("name") - node2_string = edge.node2.attributes["text"] or edge.node2.attributes.get("name") + node1_string = edge.node1.attributes.get("text") or edge.node1.attributes.get("name") + node2_string = edge.node2.attributes.get("text") or edge.node2.attributes.get("name") edge_string = edge.attributes["relationship_type"] edge_str = f"{node1_string} -- {edge_string} -- {node2_string}" edge_strings.append(edge_str)