Add Lint to CI (Ruff & Black)

.github/workflows/test.yml

@@ -60,6 +60,11 @@ jobs:
       - name: Install dependencies
         run: poetry install --extras "all"
 
+      - name: Run Lint
+        run: |
+          poetry run ruff check ./graphrag_sdk/
+          poetry run black ./graphrag_sdk/ --check
+
       - name: Wait for Ollama to be ready
         run: |
           until curl -s http://localhost:11434; do

examples/movies/demo-movies.ipynb

@@ -34,7 +34,7 @@
     "\n",
     "# Load environment variables\n",
     "load_dotenv()\n",
-    "logging.disable(logging.CRITICAL)\n"
+    "logging.disable(logging.CRITICAL)"
    ]
   },
   {
@@ -57,7 +57,7 @@
     "falkor_host = \"\"\n",
     "falkor_port = None\n",
     "falkor_username = \"\"\n",
-    "falkor_password = \"\"\n"
+    "falkor_password = \"\""
    ]
   },
   {
@@ -75,12 +75,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "urls = [\"https://www.rottentomatoes.com/m/side_by_side_2012\",\n",
-    "\"https://www.rottentomatoes.com/m/matrix\",\n",
-    "\"https://www.rottentomatoes.com/m/matrix_revolutions\",\n",
-    "\"https://www.rottentomatoes.com/m/matrix_reloaded\",\n",
-    "\"https://www.rottentomatoes.com/m/speed_1994\",\n",
-    "\"https://www.rottentomatoes.com/m/john_wick_chapter_4\"]\n",
+    "urls = [\n",
+    "    \"https://www.rottentomatoes.com/m/side_by_side_2012\",\n",
+    "    \"https://www.rottentomatoes.com/m/matrix\",\n",
+    "    \"https://www.rottentomatoes.com/m/matrix_revolutions\",\n",
+    "    \"https://www.rottentomatoes.com/m/matrix_reloaded\",\n",
+    "    \"https://www.rottentomatoes.com/m/speed_1994\",\n",
+    "    \"https://www.rottentomatoes.com/m/john_wick_chapter_4\",\n",
+    "]\n",
     "\n",
     "sources = [URL(url) for url in urls]"
    ]
@@ -136,7 +138,7 @@
     "    host=falkor_host,\n",
     "    port=falkor_port,\n",
     "    username=falkor_username,\n",
-    "    password=falkor_password\n",
+    "    password=falkor_password,\n",
     ")\n",
     "kg.process_sources(sources)"
    ]
@@ -190,7 +192,9 @@
     "answer = chat.send_message(\"Who is the director of the movie Side by Side?\")\n",
     "print(f\"Q: {answer['question']} \\nA: {answer['response']}\\n\")\n",
     "\n",
-    "answer = chat.send_message(\"Order the directors that you mentioned in all of our conversation by lexical order.\")\n",
+    "answer = chat.send_message(\n",
+    "    \"Order the directors that you mentioned in all of our conversation by lexical order.\"\n",
+    ")\n",
     "print(f\"Q: {answer['question']} \\nA: {answer['response']}\\n\")"
    ]
   }

examples/trip/demo_orchestrator_trip.ipynb

@@ -41,7 +41,13 @@
     "from graphrag_sdk.agents.kg_agent import KGAgent\n",
     "from graphrag_sdk.models.openai import OpenAiGenerativeModel\n",
     "from graphrag_sdk import (\n",
-    "    Ontology, Entity, Relation, Attribute, AttributeType, KnowledgeGraph, KnowledgeGraphModelConfig\n",
+    "    Ontology,\n",
+    "    Entity,\n",
+    "    Relation,\n",
+    "    Attribute,\n",
+    "    AttributeType,\n",
+    "    KnowledgeGraph,\n",
+    "    KnowledgeGraphModelConfig,\n",
     ")\n",
     "\n",
     "# Load environment variables\n",
@@ -402,7 +408,9 @@
     "orchestrator.register_agent(attractions_agent)\n",
     "\n",
     "# Query the orchestrator\n",
-    "runner = orchestrator.ask(\"Create a two-day itinerary for a trip to Rome. Please don't ask me any questions. Just provide the best itinerary you can.\")"
+    "runner = orchestrator.ask(\n",
+    "    \"Create a two-day itinerary for a trip to Rome. Please don't ask me any questions. Just provide the best itinerary you can.\"\n",
+    ")"
    ]
   },
   {
@@ -464,7 +472,9 @@
     }
    ],
    "source": [
-    "runner = orchestrator.ask(\"Please tell me only the name of the restaurant for dinner at the first day that you mention in the itinerary of the trip\")\n",
+    "runner = orchestrator.ask(\n",
+    "    \"Please tell me only the name of the restaurant for dinner at the first day that you mention in the itinerary of the trip\"\n",
+    ")\n",
     "print(runner.output)"
    ]
   },
@@ -510,7 +520,9 @@
     }
    ],
    "source": [
-    "runner = orchestrator.ask(\"Can you change this restaurant to another one and give me the updated itinerary?\")\n",
+    "runner = orchestrator.ask(\n",
+    "    \"Can you change this restaurant to another one and give me the updated itinerary?\"\n",
+    ")\n",
     "print(runner.output)"
    ]
   },
@@ -528,7 +540,9 @@
     }
    ],
    "source": [
-    "runner = orchestrator.ask(\"Please tell me attraction on the morning at the first day that you mention in the itinerary of the trip\")\n",
+    "runner = orchestrator.ask(\n",
+    "    \"Please tell me attraction on the morning at the first day that you mention in the itinerary of the trip\"\n",
+    ")\n",
     "print(runner.output)"
    ]
   }

graphrag_sdk/__init__.py

@@ -37,4 +37,4 @@
     "Relation",
     "Attribute",
     "AttributeType",
-]
+]

graphrag_sdk/agents/__init__.py

@@ -1,3 +1,3 @@
 from .agent import Agent
 
-__all__ = ['Agent']
+__all__ = ["Agent"]

graphrag_sdk/agents/kg_agent.py

@@ -126,7 +126,7 @@ def run(self, params: dict) -> str:
 
         """
         output = self.chat_session.send_message(params["prompt"])
-        return output['response']
+        return output["response"]
 
     def __repr__(self) -> str:
         """

graphrag_sdk/attribute.py

@@ -58,18 +58,18 @@ def from_string(txt: str) -> "AttributeType":
         raise ValueError(f"Invalid attribute type: {txt}")
 
 class Attribute:
-    """ Represents an attribute of an entity or relation in the ontology.
-
-        Args:
-            name (str): The name of the attribute.
-            attr_type (AttributeType): The type of the attribute.
-            unique (bool): Whether the attribute is unique.
-            required (bool): Whether the attribute is required.
-
-        Examples:
-            >>> attr = Attribute("name", AttributeType.STRING, True, True)
-            >>> print(attr)
-            name: "string!*"
+    """Represents an attribute of an entity or relation in the ontology.
+
+    Args:
+        name (str): The name of the attribute.
+        attr_type (AttributeType): The type of the attribute.
+        unique (bool): Whether the attribute is unique.
+        required (bool): Whether the attribute is required.
+
+    Examples:
+        >>> attr = Attribute("name", AttributeType.STRING, True, True)
+        >>> print(attr)
+        name: "string!*"
     """
 
     def __init__(
@@ -168,4 +168,4 @@ def __str__(self) -> str:
         Returns:
             str: A string representation of the Attribute object.
         """
-        return f"{self.name}: \"{self.type}{'!' if self.unique else ''}{'*' if self.required else ''}\""
+        return f'{self.name}: "{self.type}{"!" if self.unique else ""}{"*" if self.required else ""}"'

graphrag_sdk/chat_session.py

@@ -25,9 +25,17 @@ class ChatSession:
         >>> chat_session.send_message("What is the capital of France?")
     """
 
-    def __init__(self, model_config: KnowledgeGraphModelConfig, ontology: Ontology, graph: Graph,
-                cypher_system_instruction: str, qa_system_instruction: str,
-                cypher_gen_prompt: str, qa_prompt: str, cypher_gen_prompt_history: str):
+    def __init__(
+        self,
+        model_config: KnowledgeGraphModelConfig,
+        ontology: Ontology,
+        graph: Graph,
+        cypher_system_instruction: str,
+        qa_system_instruction: str,
+        cypher_gen_prompt: str,
+        qa_prompt: str,
+        cypher_gen_prompt_history: str,
+    ):
         """
         Initializes a new ChatSession object.
 
@@ -73,9 +81,9 @@ def send_message(self, message: str) -> dict:
 
         Returns:
             dict: The response to the message in the following format:
-                    {"question": message, 
-                    "response": answer, 
-                    "context": context, 
+                    {"question": message,
+                    "response": answer,
+                    "context": context,
                     "cypher": cypher}
         """
         cypher_step = GraphQueryGenerationStep(
@@ -84,7 +92,7 @@ def send_message(self, message: str) -> dict:
             ontology=self.ontology,
             last_answer=self.last_answer,
             cypher_prompt=self.cypher_prompt,
-            cypher_prompt_with_history=self.cypher_prompt_with_history
+            cypher_prompt_with_history=self.cypher_prompt_with_history,
         )
 
         (context, cypher) = cypher_step.run(message)
@@ -94,8 +102,8 @@ def send_message(self, message: str) -> dict:
                 "question": message,
                 "response": "I am sorry, I could not find the answer to your question",
                 "context": None,
-                "cypher": None
-                }
+                "cypher": None,
+            }
 
         qa_step = QAStep(
             chat_session=self.qa_chat_session,
@@ -104,7 +112,7 @@ def send_message(self, message: str) -> dict:
 
         answer = qa_step.run(message, cypher, context)
         self.last_answer = answer
-        
+
         return {
             "question": message, 
             "response": answer, 

graphrag_sdk/document_loaders/jsonl.py

@@ -18,10 +18,6 @@ def load(self) -> Iterator[Document]:
             num_documents = num_rows // self.rows_per_document
             for i in range(num_documents):
                 content = "\n".join(
-                    rows[
-                        i
-                        * self.rows_per_document : (i + 1)
-                        * self.rows_per_document
-                    ]
+                    rows[i * self.rows_per_document : (i + 1) * self.rows_per_document]
                 )
                 yield Document(content, f"{self.path}#{i}")

graphrag_sdk/document_loaders/pdf.py

@@ -1,7 +1,8 @@
 from typing import Iterator
 from graphrag_sdk.document import Document
 
-class PDFLoader():
+
+class PDFLoader:
     """
     Load PDF
     """
@@ -15,11 +16,11 @@ def __init__(self, path: str) -> None:
         """
 
         try:
-            import pypdf
-        except ImportError:
-            raise ImportError(
-                "pypdf package not found, please install it with " "`pip install pypdf`"
-        )
+            __import__("pypdf")
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "pypdf package not found, please install it with `pip install pypdf`"
+            )
 
         self.path = path
 
@@ -30,8 +31,8 @@ def load(self) -> Iterator[Document]:
         Returns:
             Iterator[Document]: document iterator
         """
-        
-        from pypdf import PdfReader # pylint: disable=import-outside-toplevel
+
+        from pypdf import PdfReader  # pylint: disable=import-outside-toplevel
 
         reader = PdfReader(self.path)
         yield from [

graphrag_sdk/document_loaders/text.py

@@ -1,7 +1,8 @@
 from typing import Iterator
 from graphrag_sdk.document import Document
 
-class TextLoader():
+
+class TextLoader:
     """
     Load Text
     """

graphrag_sdk/document_loaders/url.py

@@ -4,7 +4,8 @@
 from bs4 import BeautifulSoup
 from graphrag_sdk.document import Document
 
-class URLLoader():
+
+class URLLoader:
     """
     Load URL
     """
@@ -21,7 +22,7 @@ def __init__(self, url: str) -> None:
 
     def _download(self) -> str:
         try:
-            response = requests.get(self.url, headers={'User-Agent': 'Mozilla/5.0'})
+            response = requests.get(self.url, headers={"User-Agent": "Mozilla/5.0"})
             response.raise_for_status()  # Raise an HTTPError for bad responses (4xx and 5xx)
             return response.text
         except requests.exceptions.RequestException as e:
@@ -39,12 +40,12 @@ def load(self) -> Iterator[Document]:
         content = self._download()
 
         # extract text from HTML, populate content
-        soup = BeautifulSoup(content, 'html.parser')
+        soup = BeautifulSoup(content, "html.parser")
 
         # Extract text from the HTML
         content = soup.get_text()
 
         # Remove extra newlines
-        content = re.sub(r'\n{2,}', '\n', content)
+        content = re.sub(r"\n{2,}", "\n", content)
 
         yield Document(content, self.url)

graphrag_sdk/fixtures/regex.py

@@ -4,4 +4,4 @@
 
 NODE_LABEL_REGEX = r"\(.+:(.*?)\)"
 
-NODE_REGEX = r"\(.*?\)"
+NODE_REGEX = r"\(.*?\)"

graphrag_sdk/helpers.py

@@ -84,12 +84,12 @@ def stringify_falkordb_response(response: Union[list, str]) -> str:
     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
@@ -194,10 +194,10 @@ def validate_cypher_relations_exist(cypher: str, ontology: Ontology) -> list[str
     for relation in relation_labels:
         for label in relation.split("|"):
             max_idx = min(
-                    label.index("*") if "*" in label else len(label),
-                    label.index("{") if "{" in label else len(label),
-                    label.index("]") if "]" in label else len(label),
-                    )
+                label.index("*") if "*" in label else len(label),
+                label.index("{") if "{" in label else len(label),
+                label.index("]") if "]" in label else len(label),
+            )
             label = label[:max_idx]
             if label not in [relation.label for relation in ontology.relations]:
                 not_found_relation_labels.append(label)