fix format

Author: gkorland

.github/workflows/test.yml

@@ -33,6 +33,10 @@ jobs:
         with:
           python-version: 3.x # Update with desired Python version
 
+      - uses: datadog/action-py-black-formatter@v2.1
+        with:
+          check_mode: "true"
+
       - uses: astral-sh/ruff-action@v3
         with:
           src: >-

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,6 @@
 from abc import ABC, abstractmethod
 
+
 class AgentResponseCode:
     """
     Represents the response codes for an agent.

graphrag_sdk/agents/kg_agent.py

@@ -1,6 +1,7 @@
 from graphrag_sdk.kg import KnowledgeGraph
 from .agent import Agent
 
+
 class KGAgent(Agent):
     """Represents an Agent for a FalkorDB Knowledge Graph.
 
@@ -134,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

@@ -41,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__(

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')
+            __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]