Merge pull request #179 from FalkorDB/fix-ontology-imports

Revert "add delay to avoid rate limits"  & Fix curricular import

Author: galshubeli

graphrag_sdk/kg.py

@@ -1,6 +1,5 @@
 import logging
 import warnings
-import time
 from falkordb import FalkorDB
 from typing import Optional, Union
 from graphrag_sdk.ontology import Ontology
@@ -203,16 +202,15 @@ def list_sources(self) -> list[AbstractSource]:
         return [s.source for s in self.sources]
 
     def process_sources(
-        self, sources: list[AbstractSource], instructions: Optional[str] = None, hide_progress: Optional[bool] = False,  delay: Optional[float] = 0
+        self, sources: list[AbstractSource], instructions: Optional[str] = None, hide_progress: Optional[bool] = False
     ) -> None:
         """
         Add entities and relations found in sources into the knowledge-graph
 
         Args:
-            sources (list[AbstractSource]): list of sources to extract knowledge from            
+            sources (list[AbstractSource]): list of sources to extract knowledge from
             instructions (Optional[str]): Instructions for processing.
             hide_progress (Optional[bool]): hide progress bar
-            delay (float): seconds to delay each iteration through sources to avoid rate limits
         """
 
         if self.ontology is None:
@@ -221,11 +219,6 @@ def process_sources(
         # Create graph with sources
         self._create_graph_with_sources(sources, instructions, hide_progress)
 
-        # Add processed sources
-        for src in sources:
-            self.sources.add(src)
-            time.sleep(delay)
-
 
     def _create_graph_with_sources(
         self, sources: Optional[list[AbstractSource]] = None, instructions: Optional[str] = None, hide_progress: Optional[bool] = False

graphrag_sdk/ontology.py

@@ -1,13 +1,11 @@
 import json
 import logging
-import graphrag_sdk
 from falkordb import Graph
 from .entity import Entity
 from .relation import Relation
 from typing import Optional, Union
 from graphrag_sdk.source import AbstractSource
 from graphrag_sdk.models import GenerativeModel
-from graphrag_sdk.steps.create_ontology_step import CreateOntologyStep
 from .attribute import Attribute, AttributeType
 
 logger = logging.getLogger(__name__)
@@ -72,7 +70,10 @@ def from_sources(
         Returns:
             The created Ontology object.
         """
-        step = graphrag_sdk.CreateOntologyStep(
+        # Import here to avoid circular import
+        from graphrag_sdk.steps.create_ontology_step import CreateOntologyStep
+        
+        step = CreateOntologyStep(
             sources=sources,
             ontology=Ontology(),
             model=model,

graphrag_sdk/steps/create_ontology_step.py

@@ -5,7 +5,6 @@
 from threading import Lock
 from typing import Optional
 from graphrag_sdk.steps.Step import Step
-from graphrag_sdk.document import Document
 from graphrag_sdk.ontology import Ontology
 from graphrag_sdk.helpers import extract_json
 from ratelimit import limits, sleep_and_retry