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

Merge branch 'dev' into limit_major_dep_version_change

Browse source
This commit is contained in:
Vasilije 2025-07-13 20:40:15 +02:00 • committed by GitHub
commit 0db56c6e36
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 101 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
cognee/modules/graph/utils/expand_with_nodes_and_edges.py
View file

@ -25,7 +25,7 @@ def _process_ontology_nodes(
ontology_nodes: list, ontology_nodes: list,
data_chunk: DocumentChunk, data_chunk: DocumentChunk,
added_nodes_map: dict, added_nodes_map: dict,
added_ontology_nodes_map: dict added_ontology_nodes_map: dict,
) -> None: ) -> None:
"""Process and store ontology nodes""" """Process and store ontology nodes"""
for ontology_node in ontology_nodes: for ontology_node in ontology_nodes:
@ -55,9 +55,7 @@ def _process_ontology_nodes(
def _process_ontology_edges( def _process_ontology_edges(
ontology_edges: list, ontology_edges: list, existing_edges_map: dict, ontology_relationships: list
existing_edges_map: dict,
ontology_relationships: list
) -> None: ) -> None:
"""Process ontology edges and add them if new""" """Process ontology edges and add them if new"""
for source, relation, target in ontology_edges: for source, relation, target in ontology_edges:
@ -92,7 +90,7 @@ def _create_type_node(
key_mapping: dict, key_mapping: dict,
data_chunk: DocumentChunk, data_chunk: DocumentChunk,
existing_edges_map: dict, existing_edges_map: dict,
ontology_relationships: list ontology_relationships: list,
) -> EntityType: ) -> EntityType:
"""Create or retrieve a type node with ontology validation""" """Create or retrieve a type node with ontology validation"""
node_id = generate_node_id(node_type) node_id = generate_node_id(node_type)
@ -100,7 +98,9 @@ def _create_type_node(
type_node_key = _create_node_key(node_id, "type") type_node_key = _create_node_key(node_id, "type")
if type_node_key in added_nodes_map or type_node_key in key_mapping: if type_node_key in added_nodes_map or type_node_key in key_mapping:
return added_nodes_map.get(type_node_key) or added_nodes_map.get(key_mapping.get(type_node_key)) return added_nodes_map.get(type_node_key) or added_nodes_map.get(
key_mapping.get(type_node_key)
)
# Get ontology validation # Get ontology validation
ontology_nodes, ontology_edges, closest_class = ontology_resolver.get_subgraph( ontology_nodes, ontology_edges, closest_class = ontology_resolver.get_subgraph(
@ -148,7 +148,7 @@ def _create_entity_node(
key_mapping: dict, key_mapping: dict,
data_chunk: DocumentChunk, data_chunk: DocumentChunk,
existing_edges_map: dict, existing_edges_map: dict,
ontology_relationships: list ontology_relationships: list,
) -> Entity: ) -> Entity:
"""Create or retrieve an entity node with ontology validation""" """Create or retrieve an entity node with ontology validation"""
generated_node_id = generate_node_id(node_id) generated_node_id = generate_node_id(node_id)
@ -156,7 +156,9 @@ def _create_entity_node(
entity_node_key = _create_node_key(generated_node_id, "entity") entity_node_key = _create_node_key(generated_node_id, "entity")
if entity_node_key in added_nodes_map or entity_node_key in key_mapping: if entity_node_key in added_nodes_map or entity_node_key in key_mapping:
return added_nodes_map.get(entity_node_key) or added_nodes_map.get(key_mapping.get(entity_node_key)) return added_nodes_map.get(entity_node_key) or added_nodes_map.get(
key_mapping.get(entity_node_key)
)
# Get ontology validation # Get ontology validation
ontology_nodes, ontology_edges, start_ent_ont = ontology_resolver.get_subgraph( ontology_nodes, ontology_edges, start_ent_ont = ontology_resolver.get_subgraph(
@ -202,21 +204,37 @@ def _process_graph_nodes(
name_mapping: dict, name_mapping: dict,
key_mapping: dict, key_mapping: dict,
existing_edges_map: dict, existing_edges_map: dict,
ontology_relationships: list ontology_relationships: list,
) -> None: ) -> None:
"""Process nodes in a knowledge graph""" """Process nodes in a knowledge graph"""
for node in graph.nodes: for node in graph.nodes:
# Create type node # Create type node
type_node = _create_type_node( type_node = _create_type_node(
node.type, ontology_resolver, added_nodes_map, added_ontology_nodes_map, node.type,
name_mapping, key_mapping, data_chunk, existing_edges_map, ontology_relationships ontology_resolver,
added_nodes_map,
added_ontology_nodes_map,
name_mapping,
key_mapping,
data_chunk,
existing_edges_map,
ontology_relationships,
) )
# Create entity node # Create entity node
entity_node = _create_entity_node( entity_node = _create_entity_node(
node.id, node.name, node.description, type_node, ontology_resolver, node.id,
added_nodes_map, added_ontology_nodes_map, name_mapping, key_mapping, node.name,
data_chunk, existing_edges_map, ontology_relationships node.description,
type_node,
ontology_resolver,
added_nodes_map,
added_ontology_nodes_map,
name_mapping,
key_mapping,
data_chunk,
existing_edges_map,
ontology_relationships,
) )
# Add entity to data chunk # Add entity to data chunk
@ -226,10 +244,7 @@ def _process_graph_nodes(
def _process_graph_edges( def _process_graph_edges(
graph: KnowledgeGraph, graph: KnowledgeGraph, name_mapping: dict, existing_edges_map: dict, relationships: list
name_mapping: dict,
existing_edges_map: dict,
relationships: list
) -> None: ) -> None:
"""Process edges in a knowledge graph""" """Process edges in a knowledge graph"""
for edge in graph.edges: for edge in graph.edges:
@ -288,8 +303,15 @@ def expand_with_nodes_and_edges(
# Process nodes first # Process nodes first
_process_graph_nodes( _process_graph_nodes(
data_chunk, graph, ontology_resolver, added_nodes_map, added_ontology_nodes_map, data_chunk,
name_mapping, key_mapping, existing_edges_map, ontology_relationships graph,
ontology_resolver,
added_nodes_map,
added_ontology_nodes_map,
name_mapping,
key_mapping,
existing_edges_map,
ontology_relationships,
) )
# Then process edges # Then process edges