refactor: Change payload text use to use edge id

This commit is contained in:
Igor Ilic 2026-01-19 13:57:30 +01:00
parent 334574444c
commit 01a6382552
2 changed files with 7 additions and 8 deletions

View file

@ -1,5 +1,6 @@
import time import time
from cognee.shared.logging_utils import get_logger from cognee.shared.logging_utils import get_logger
from cognee.modules.engine.utils.generate_edge_id import generate_edge_id
from typing import List, Dict, Union, Optional, Type, Iterable, Tuple, Callable, Any from typing import List, Dict, Union, Optional, Type, Iterable, Tuple, Callable, Any
from cognee.modules.graph.exceptions import ( from cognee.modules.graph.exceptions import (
@ -205,6 +206,10 @@ class CogneeGraph(CogneeAbstractGraph):
key: properties.get(key) for key in edge_properties_to_project key: properties.get(key) for key in edge_properties_to_project
} }
edge_attributes["relationship_type"] = relationship_type edge_attributes["relationship_type"] = relationship_type
edge_text = properties.get("edge_text") or properties.get("relationship_name")
edge_attributes["edge_type_id"] = (
generate_edge_id(edge_id=edge_text) if edge_text else None
)
edge = Edge( edge = Edge(
source_node, source_node,
@ -284,13 +289,7 @@ class CogneeGraph(CogneeAbstractGraph):
for query_index, scored_results in enumerate(per_query_scored_results): for query_index, scored_results in enumerate(per_query_scored_results):
for result in scored_results: for result in scored_results:
payload = getattr(result, "payload", None) matching_edges = self.edges_by_distance_key.get(str(result.id))
if not isinstance(payload, dict):
continue
text = payload.get("text")
if not text:
continue
matching_edges = self.edges_by_distance_key.get(str(text))
if not matching_edges: if not matching_edges:
continue continue
for edge in matching_edges: for edge in matching_edges:

View file

@ -141,7 +141,7 @@ class Edge:
self.status = np.ones(dimension, dtype=int) self.status = np.ones(dimension, dtype=int)
def get_distance_key(self) -> Optional[str]: def get_distance_key(self) -> Optional[str]:
key = self.attributes.get("edge_text") or self.attributes.get("relationship_type") key = self.attributes.get("edge_type_id")
if key is None: if key is None:
return None return None
return str(key) return str(key)