refactor: Change payload text use to use edge id
This commit is contained in:
parent
334574444c
commit
01a6382552
2 changed files with 7 additions and 8 deletions
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue