ruff & mypy

Author: stellasia

examples/question_answering/graphrag_pipeline.py

@@ -1,24 +1,22 @@
 import asyncio
-import neo4j
 
+import neo4j
 from neo4j_graphrag.embeddings import OpenAIEmbeddings
-from neo4j_graphrag.experimental.pipeline import Pipeline
-from neo4j_graphrag.experimental.components.rag.retrievers import RetrieverWrapper
-from neo4j_graphrag.experimental.components.rag.prompt_builder import PromptBuilder
 from neo4j_graphrag.experimental.components.rag.generate import Generator
+from neo4j_graphrag.experimental.components.rag.prompt_builder import PromptBuilder
+from neo4j_graphrag.experimental.components.rag.retrievers import RetrieverWrapper
+from neo4j_graphrag.experimental.pipeline import Pipeline
 from neo4j_graphrag.generation import RagTemplate
 from neo4j_graphrag.llm import OpenAILLM
-
 from neo4j_graphrag.retrievers import VectorRetriever
 
-
 URI = "neo4j+s://demo.neo4jlabs.com"
 AUTH = ("recommendations", "recommendations")
 DATABASE = "recommendations"
 INDEX_NAME = "moviePlotsEmbedding"
 
 
-async def main():
+async def main() -> None:
     pipeline = Pipeline()
     driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
     llm = OpenAILLM(model_name="gpt-4o")
@@ -33,25 +31,33 @@ async def main():
     pipeline.add_component(PromptBuilder(RagTemplate()), "prompt")
     pipeline.add_component(Generator(llm), "generate")
 
-    pipeline.connect("retriever", "prompt", {
-        "context": "retriever.result",
-    })
-    pipeline.connect("prompt", "generate", {
-        "prompt": "prompt.prompt",
-    })
+    pipeline.connect(
+        "retriever",
+        "prompt",
+        {
+            "context": "retriever.result",
+        },
+    )
+    pipeline.connect(
+        "prompt",
+        "generate",
+        {
+            "prompt": "prompt.prompt",
+        },
+    )
 
     query = "show me a movie with cats"
-    res = await pipeline.run({
-        "retriever": {"query_text": query},
-        "prompt": {"query_text": query, "examples": ""},
-    })
+    res = await pipeline.run(
+        {
+            "retriever": {"query_text": query},
+            "prompt": {"query_text": query, "examples": ""},
+        }
+    )
 
     driver.close()
     await llm.async_client.close()
     return res
 
 
 if __name__ == "__main__":
-    print(
-        asyncio.run(main())
-    )
+    print(asyncio.run(main()))

examples/question_answering/graphrag_simple_pipeline.py

@@ -1,7 +1,6 @@
 from neo4j_graphrag.experimental.pipeline.config.runner import PipelineRunner
 
 if __name__ == "__main__":
-
     import asyncio
     import os
 
@@ -14,12 +13,14 @@
     )
     print(
         asyncio.run(
-            runner.run(dict(
-                query_text="show me a movie about cats",
-                retriever_config={
-                    "top_k": 2,
-                },
-                return_context=True,
-            ))
+            runner.run(
+                dict(
+                    query_text="show me a movie about cats",
+                    retriever_config={
+                        "top_k": 2,
+                    },
+                    return_context=True,
+                )
+            )
         )
     )

src/neo4j_graphrag/experimental/components/rag/prompt_builder.py

@@ -17,7 +17,6 @@
 from neo4j_graphrag.experimental.pipeline import Component, DataModel
 from neo4j_graphrag.generation import PromptTemplate
 
-
 # class PromptData(DataModel):
 #     inputs: dict[str, Any]
 

src/neo4j_graphrag/experimental/pipeline/config/runner.py

@@ -42,12 +42,12 @@
     AbstractPipelineConfig,
     PipelineConfig,
 )
-from neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_rag_pipeline import (
-    SimpleRAGPipelineConfig,
-)
 from neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder import (
     SimpleKGPipelineConfig,
 )
+from neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_rag_pipeline import (
+    SimpleRAGPipelineConfig,
+)
 from neo4j_graphrag.experimental.pipeline.config.types import PipelineType
 from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
 from neo4j_graphrag.experimental.pipeline.types import PipelineDefinition

src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_rag_pipeline.py

@@ -40,7 +40,7 @@ def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> RetrieverWrap
         return RetrieverWrapper(retriever)
 
 
-class RetrieverType(RootModel):
+class RetrieverType(RootModel):  # type: ignore
     root: Union[RetrieverWrapper, RetrieverConfig]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)