gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: Add fix for accessing dictionary elements that don't exits

Browse source 
Using get for the text key instead of direct access to handle situation if the text key doesn't exist
This commit is contained in:
Igor Ilic 2025-01-09 16:06:26 +01:00
parent 341f30fcdc
commit 6b6cc0f1d4
1 changed files with 3 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cognee/tasks/completion/graph_query_completion.py
View file

@ -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)