Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 2 additions & 9 deletions graphrag_sdk/kg.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
7 changes: 4 additions & 3 deletions graphrag_sdk/ontology.py
Original file line number Diff line number Diff line change
@@ -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__)
Expand Down Expand Up @@ -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,
Expand Down
1 change: 0 additions & 1 deletion graphrag_sdk/steps/create_ontology_step.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After removing Document, _process_source’s docstring still documents a document (Document) parameter that isn’t in the signature; consider updating the docstring to match (source: AbstractSource) to avoid misleading API docs.

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎

from graphrag_sdk.ontology import Ontology
from graphrag_sdk.helpers import extract_json
from ratelimit import limits, sleep_and_retry
Expand Down
Loading