Add Lint to CI (Ruff & Black)

.github/workflows/test.yml

@@ -32,7 +32,7 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: 3.x # Update with desired Python version
-
+            
       - name: Cache Poetry virtualenv
         id: cache
         uses: actions/cache@v4
@@ -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/agent.py

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

graphrag_sdk/agents/kg_agent.py

@@ -1,6 +1,5 @@
 from graphrag_sdk.kg import KnowledgeGraph
 from .agent import Agent
-from graphrag_sdk.models import GenerativeModelChatSession
 
 
 class KGAgent(Agent):
@@ -136,7 +135,7 @@ def run(self, params: dict) -> str:
 
         """
         output = self.chat_session.send_message(params["prompt"])
-        return output['response']
+        return output["response"]
 
     def __repr__(self):
         """

graphrag_sdk/attribute.py

@@ -1,5 +1,4 @@
 import json
-from graphrag_sdk.fixtures.regex import *
 import logging
 import re
 
@@ -42,18 +41,18 @@ def from_string(txt: str):
 
 
 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__(
@@ -158,4 +157,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

@@ -24,9 +24,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.
 
@@ -45,21 +53,22 @@ def __init__(self, model_config: KnowledgeGraphModelConfig, ontology: Ontology,
         self.model_config = model_config
         self.graph = graph
         self.ontology = ontology
-        cypher_system_instruction = cypher_system_instruction.format(ontology=str(ontology.to_json()))
+        cypher_system_instruction = cypher_system_instruction.format(
+            ontology=str(ontology.to_json())
+        )
 
-
         self.cypher_prompt = cypher_gen_prompt
         self.qa_prompt = qa_prompt
         self.cypher_prompt_with_history = cypher_gen_prompt_history
-        
+
         self.cypher_chat_session = (
             model_config.cypher_generation.with_system_instruction(
                 cypher_system_instruction
             ).start_chat()
         )
         self.qa_chat_session = model_config.qa.with_system_instruction(
-                qa_system_instruction
-            ).start_chat()
+            qa_system_instruction
+        ).start_chat()
         self.last_answer = None
 
     def send_message(self, message: str):
@@ -71,9 +80,9 @@ def send_message(self, message: str):
 
         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(
@@ -82,7 +91,7 @@ def send_message(self, message: str):
             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)
@@ -92,8 +101,8 @@ def send_message(self, message: str):
                 "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,
@@ -102,10 +111,10 @@ def send_message(self, message: str):
 
         answer = qa_step.run(message, cypher, context)
         self.last_answer = answer
-        
+
         return {
-            "question": message, 
-            "response": answer, 
-            "context": context, 
-            "cypher": cypher
-            }
+            "question": message,
+            "response": answer,
+            "context": context,
+            "cypher": cypher,
+        }

graphrag_sdk/document.py

@@ -1,4 +1,4 @@
-class Document():
+class Document:
     """
     Common class containing text extracted from a source
     """

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)

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,11 +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 [
-            Document(page.extract_text())
-            for page in reader.pages
-        ]
+        yield from [Document(page.extract_text()) for page in reader.pages]

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
     """
@@ -24,7 +25,5 @@ def load(self) -> Iterator[Document]:
             Iterator[Document]: document iterator
         """
 
-        with open(self.path, 'r') as f:
-            yield Document(
-                f.read()
-            )
+        with open(self.path, "r") as f:
+            yield Document(f.read())