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

fix: fixes user-cognee interaction edges for kuzu

Browse source
This commit is contained in:
hajdul88 2025-07-30 12:13:30 +02:00
parent e0221feb74
commit 9d46453a7a
1 changed files with 47 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

79
cognee/modules/retrieval/graph_completion_retriever.py
View file

@ -88,6 +88,21 @@ class GraphCompletionRetriever(BaseRetriever):
) )
return f"Nodes:\n{node_section}\n\nConnections:\n{connection_section}" return f"Nodes:\n{node_section}\n\nConnections:\n{connection_section}"
def extract_uuid_from_node(self, node: Any) -> Optional[UUID]:
"""
Try to pull a UUID string out of node.id or node.properties['id'],
then return a UUID instance (or None if neither exists).
"""
id_str = None
if not id_str:
id_str = getattr(node, "id", None)
if hasattr(node, "attributes") and not id_str:
id_str = node.attributes.get("id", None)
id = UUID(id_str) if isinstance(id_str, str) else None
return id
async def save_qa(self, question: str, answer: str, context: str, triplets: List) -> None: async def save_qa(self, question: str, answer: str, context: str, triplets: List) -> None:
""" """
Saves a question and answer pair for later analysis or storage. Saves a question and answer pair for later analysis or storage.
@ -118,42 +133,42 @@ class GraphCompletionRetriever(BaseRetriever):
relationships = [] relationships = []
relationship_name = "used_graph_element_to_answer" relationship_name = "used_graph_element_to_answer"
for triplet in triplets: for triplet in triplets:
target_id_1 = UUID(triplet.node1.id) target_id_1 = self.extract_uuid_from_node(triplet.node1)
target_id_2 = UUID(triplet.node2.id) target_id_2 = self.extract_uuid_from_node(triplet.node2)
if target_id_1 and target_id_2:
# Defined qa node to triplet node 1 # Defined qa node to triplet node 1
relationships.append( relationships.append(
( (
source_id, source_id,
target_id_1, target_id_1,
relationship_name, relationship_name,
{ {
"relationship_name": relationship_name, "relationship_name": relationship_name,
"source_node_id": source_id, "source_node_id": source_id,
"target_node_id": target_id_1, "target_node_id": target_id_1,
"ontology_valid": False, "ontology_valid": False,
}, },
)
) )
)
# Defined qa node to triplet node 2 # Defined qa node to triplet node 2
relationships.append( relationships.append(
( (
source_id, source_id,
target_id_2, target_id_2,
relationship_name, relationship_name,
{ {
"relationship_name": relationship_name, "relationship_name": relationship_name,
"source_node_id": source_id, "source_node_id": source_id,
"target_node_id": target_id_2, "target_node_id": target_id_2,
"ontology_valid": False, "ontology_valid": False,
}, },
)
) )
)
if len(relationships) > 0: if len(relationships) > 0:
graph_engine = await get_graph_engine() graph_engine = await get_graph_engine()
await graph_engine.add_edges(relationships) await graph_engine.add_edges(relationships)
async def get_triplets(self, query: str) -> list: async def get_triplets(self, query: str) -> list:
""" """