Add Lint to CI (Ruff & Black)

.github/workflows/test.yml

@@ -33,6 +33,11 @@ jobs:
         with:
           python-version: 3.x # Update with desired Python version
 
+      - uses: astral-sh/ruff-action@v3
+        with:
+          src: >-
+            "./graphrag_sdk"
+
       - name: Cache Poetry virtualenv
         id: cache
         uses: actions/cache@v4

graphrag_sdk/agents/agent.py

@@ -1,6 +1,4 @@
 from abc import ABC, abstractmethod
-from graphrag_sdk.models.model import GenerativeModelChatSession
-
 
 class AgentResponseCode:
     """

graphrag_sdk/agents/kg_agent.py

@@ -1,7 +1,5 @@
 from graphrag_sdk.kg import KnowledgeGraph
 from .agent import Agent
-from graphrag_sdk.models import GenerativeModelChatSession
-
 
 class KGAgent(Agent):
     """Represents an Agent for a FalkorDB Knowledge Graph.

graphrag_sdk/attribute.py

@@ -1,5 +1,4 @@
 import json
-from graphrag_sdk.fixtures.regex import *
 import logging
 import re
 

graphrag_sdk/document_loaders/pdf.py

@@ -15,9 +15,9 @@ def __init__(self, path: str) -> None:
         """
 
         try:
-            import pypdf
-        except ImportError:
-            raise ImportError(
+            __import__('pypdf')
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
                 "pypdf package not found, please install it with " "`pip install pypdf`"
         )
 

graphrag_sdk/helpers.py

@@ -55,12 +55,12 @@ def stringify_falkordb_response(response):
     elif not isinstance(response[0], list):
         data = str(response).strip()
     else:
-        for l, _ in enumerate(response):
-            if not isinstance(response[l], list):
-                response[l] = str(response[l])
+        for line, _ in enumerate(response):
+            if not isinstance(response[line], list):
+                response[line] = str(response[line])
             else:
-                for i, __ in enumerate(response[l]):
-                    response[l][i] = str(response[l][i])
+                for i, __ in enumerate(response[line]):
+                    response[line][i] = str(response[line][i])
         data = str(response).strip()
 
     return data

graphrag_sdk/models/litellm.py

@@ -74,7 +74,7 @@ def check_valid_key(self, model: str):
                 model=model, messages=messages, max_tokens=10
             )
             return True
-        except:
+        except:  # noqa: E722
             return False
 
     def check_and_pull_model(self) -> None:

graphrag_sdk/orchestrator/__init__.py

@@ -1,7 +1,6 @@
 from .orchestrator import Orchestrator
 from .orchestrator_runner import OrchestratorRunner
 from .execution_plan import ExecutionPlan
-from .step import StepResult, PlanStep, StepBlockType
 
 __all__ = [
     'Orchestrator',

graphrag_sdk/orchestrator/execution_plan.py

@@ -1,5 +1,5 @@
 from json import loads
-from graphrag_sdk.orchestrator.step import PlanStep, StepBlockType
+from graphrag_sdk.orchestrator.step import PlanStep
 
 
 class ExecutionPlan:

graphrag_sdk/orchestrator/orchestrator.py

@@ -1,6 +1,5 @@
 from graphrag_sdk.agents import Agent
 from graphrag_sdk.models import GenerativeModel
-from graphrag_sdk.models.model import OutputMethod
 from .orchestrator_runner import OrchestratorRunner
 from graphrag_sdk.fixtures.prompts import (
     ORCHESTRATOR_SYSTEM,

graphrag_sdk/orchestrator/steps/parallel.py

@@ -2,6 +2,7 @@
 from concurrent.futures import ThreadPoolExecutor, wait
 from graphrag_sdk.orchestrator.step_result import StepResult
 from graphrag_sdk.orchestrator.orchestrator_runner import OrchestratorRunner
+from graphrag_sdk.orchestrator.step import PlanStep
 
 
 class ParallelStepResult(StepResult):
@@ -34,16 +35,16 @@ def output(self) -> str:
 
 
 class ParallelProperties:
-    steps: list["PlanStep"]
+    steps: list[PlanStep]
 
-    def __init__(self, steps: list["PlanStep"]):
+    def __init__(self, steps: list[PlanStep]):
         self.steps = steps
 
     @staticmethod
     def from_json(json: dict) -> "ParallelProperties":
         return ParallelProperties(
             [
-                graphrag_sdk.orchestrator.step.PlanStep.from_json(step)
+                PlanStep.from_json(step)
                 for step in (json if isinstance(json, list) else json["steps"])
             ]
         )
@@ -58,7 +59,7 @@ def __repr__(self) -> str:
         return str(self)
 
 
-class ParallelStep(graphrag_sdk.orchestrator.step.PlanStep):
+class ParallelStep(PlanStep):
 
     def __init__(self, id: str, properties: ParallelProperties):
         self._id = id

graphrag_sdk/orchestrator/steps/summary.py

@@ -1,6 +1,5 @@
 import graphrag_sdk.orchestrator.step
 from graphrag_sdk.orchestrator.step_result import StepResult
-from concurrent.futures import ThreadPoolExecutor, wait
 from graphrag_sdk.orchestrator.orchestrator_runner import OrchestratorRunner
 from graphrag_sdk.fixtures.prompts import ORCHESTRATOR_SUMMARY_PROMPT
 

graphrag_sdk/source.py

@@ -1,5 +1,5 @@
 from typing import Iterator
-from abc import ABC, abstractmethod
+from abc import ABC
 from graphrag_sdk.document import Document
 from graphrag_sdk.document_loaders import (
     PDFLoader,

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
@@ -138,7 +137,7 @@ def _process_source(
                 data = json.loads(extract_json(combined_text))
             except json.decoder.JSONDecodeError as e:
                 logger.debug(f"Error extracting JSON: {e}")
-                logger.debug(f"Prompting model to fix JSON")
+                logger.debug("Prompting model to fix JSON")
                 json_fix_response = self._call_model(
                     self._create_chat(),
                     FIX_JSON_PROMPT.format(json=combined_text, error=str(e)),
@@ -170,7 +169,7 @@ def _process_source(
             return o
 
     def _fix_ontology(self, chat_session: GenerativeModelChatSession, o: Ontology):
-        logger.debug(f"Fixing ontology...")
+        logger.debug("Fixing ontology...")
 
         user_message = FIX_ONTOLOGY_PROMPT.format(ontology=o)
 
@@ -196,7 +195,7 @@ def _fix_ontology(self, chat_session: GenerativeModelChatSession, o: Ontology):
             data = json.loads(extract_json(combined_text))
         except json.decoder.JSONDecodeError as e:
             logger.debug(f"Error extracting JSON: {e}")
-            logger.debug(f"Prompting model to fix JSON")
+            logger.debug("Prompting model to fix JSON")
             json_fix_response = self._call_model(
                 self._create_chat(),
                 FIX_JSON_PROMPT.format(json=combined_text, error=str(e)),

graphrag_sdk/steps/extract_data_step.py

@@ -7,7 +7,6 @@
 from falkordb import Graph
 from threading import Lock
 from graphrag_sdk.steps.Step import Step
-from graphrag_sdk.document import Document
 from ratelimit import limits, sleep_and_retry
 from graphrag_sdk.source import AbstractSource
 from graphrag_sdk.models.model import OutputMethod
@@ -174,7 +173,7 @@ def _process_source(
                 data = json.loads(extract_json(last_respond))
             except Exception as e:
                 _task_logger.debug(f"Error extracting JSON: {e}")
-                _task_logger.debug(f"Prompting model to fix JSON")
+                _task_logger.debug("Prompting model to fix JSON")
                 json_fix_response = self._call_model(
                     self._create_chat(),
                     FIX_JSON_PROMPT.format(json=last_respond, error=str(e)),
@@ -188,7 +187,7 @@ def _process_source(
                     f"Invalid data format. Missing entities or relations. {data}"
                 )
                 raise Exception(
-                    f"Invalid data format. Missing 'entities' or 'relations' in JSON."
+                    "Invalid data format. Missing 'entities' or 'relations' in JSON."
                 )
             for entity in data["entities"]:
                 try: