Ontology creation from Knowledge Graph

graphrag_sdk/kg.py

@@ -68,7 +68,7 @@ def __init__(
             # Load ontology from DB
             ontology = Ontology.from_schema_graph(ontology_graph)
 
-            if len(ontology.entities) == 0 and len(ontology.relations) == 0:
+            if len(ontology.entities) == 0:
                 raise Exception("The ontology is empty. Load a valid ontology or create one using the ontology module.")
         else:
             # Save ontology to DB

graphrag_sdk/ontology.py

@@ -136,17 +136,18 @@ def from_kg_graph(graph: Graph, sample_size: int = 100,):
 
         # Extract attributes for each edge type, limited by the specified sample size.
         for e_type in e_types:
+            e_t = e_type[0]
+            attributes = graph.query(
+                            f"""MATCH ()-[a:{e_t}]->() call {{ with a return [k in keys(a) | [k, typeof(a[k])]] as types }} 
+                            WITH types limit {sample_size} unwind types as kt RETURN kt, count(1)""").result_set
+            attributes = [Attribute(attr[0][0], attr[0][1]) for attr in attributes]
             for s_lbls in n_labels:
                 for t_lbls in n_labels:
-                    e_t = e_type[0]
                     s_l = s_lbls[0]
                     t_l = t_lbls[0]
                     # Check if a relationship exists between the source and target entity labels
                     if graph.query(f"MATCH (s:{s_l})-[a:{e_t}]->(t:{t_l}) return a limit 1").result_set:
-                        attributes = graph.query(
-                            f"""MATCH ()-[a:{e_t}]->() call {{ with a return [k in keys(a) | [k, typeof(a[k])]] as types }} 
-                            WITH types limit {sample_size} unwind types as kt RETURN kt, count(1)""").result_set
-                        ontology.add_relation(Relation(e_t, s_l, t_l, [Attribute(attr[0][0], attr[0][1]) for attr in attributes]))
+                        ontology.add_relation(Relation(e_t, s_l, t_l, attributes))
 
         return ontology