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

feat: genarlizing getting entities from triplets

Browse source
This commit is contained in:
chinu0609 2025-11-03 00:59:04 +05:30
parent f1afd1f0a2
commit 5080e8f8a5
3 changed files with 17 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
cognee/modules/graph/utils/__init__.py
View file

@ -5,3 +5,4 @@ from .retrieve_existing_edges import retrieve_existing_edges
from .convert_node_to_data_point import convert_node_to_data_point from .convert_node_to_data_point import convert_node_to_data_point
from .deduplicate_nodes_and_edges import deduplicate_nodes_and_edges from .deduplicate_nodes_and_edges import deduplicate_nodes_and_edges
from .resolve_edges_to_text import resolve_edges_to_text from .resolve_edges_to_text import resolve_edges_to_text
from .get_entity_nodes_from_triplets import get_entity_nodes_from_triplets

13
cognee/modules/graph/utils/get_entity_nodes_from_triplets.py Normal file
View file

@ -0,0 +1,13 @@
def get_entity_nodes_from_triplets(triplets):
entity_nodes = []
seen_ids = set()
for triplet in triplets:
if hasattr(triplet, 'node1') and triplet.node1 and triplet.node1.id not in seen_ids:
entity_nodes.append({"id": str(triplet.node1.id)})
seen_ids.add(triplet.node1.id)
if hasattr(triplet, 'node2') and triplet.node2 and triplet.node2.id not in seen_ids:
entity_nodes.append({"id": str(triplet.node2.id)})
seen_ids.add(triplet.node2.id)
return entity_nodes

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

@ -22,6 +22,7 @@ from cognee.modules.engine.models.node_set import NodeSet
from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.context_global_variables import session_user from cognee.context_global_variables import session_user
from cognee.infrastructure.databases.cache.config import CacheConfig from cognee.infrastructure.databases.cache.config import CacheConfig
from cognee.modules.graph.utils import get_entity_nodes_from_triplets
logger = get_logger("GraphCompletionRetriever") logger = get_logger("GraphCompletionRetriever")
@ -139,15 +140,8 @@ class GraphCompletionRetriever(BaseGraphRetriever):
return [] return []
# context = await self.resolve_edges_to_text(triplets) # context = await self.resolve_edges_to_text(triplets)
entity_nodes = []
seen_ids = set() entity_nodes = get_entity_nodes_from_triplets(triplets)
for triplet in triplets:
if hasattr(triplet, 'node1') and triplet.node1 and triplet.node1.id not in seen_ids:
entity_nodes.append({"id": str(triplet.node1.id)})
seen_ids.add(triplet.node1.id)
if hasattr(triplet, 'node2') and triplet.node2 and triplet.node2.id not in seen_ids:
entity_nodes.append({"id": str(triplet.node2.id)})
seen_ids.add(triplet.node2.id)
await update_node_access_timestamps(entity_nodes) await update_node_access_timestamps(entity_nodes)
return triplets return triplets