cleaning

Author: stellasia

examples/question_answering/graphrag_pipeline.py

@@ -6,6 +6,7 @@
 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.experimental.pipeline.pipeline import PipelineResult
 from neo4j_graphrag.generation import RagTemplate
 from neo4j_graphrag.llm import OpenAILLM
 from neo4j_graphrag.retrievers import VectorRetriever
@@ -16,7 +17,7 @@
 INDEX_NAME = "moviePlotsEmbedding"
 
 
-async def main() -> None:
+async def main() -> PipelineResult:
     pipeline = Pipeline()
     driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
     llm = OpenAILLM(model_name="gpt-4o")
@@ -56,6 +57,12 @@ async def main() -> None:
 
     driver.close()
     await llm.async_client.close()
+
+    # context_result = await pipeline.store.get_result_for_component(
+    #     res.run_id, "retriever"
+    # )
+    # context = context_result.get("result")
+    # res.result["context"] = context
     return res
 
 

examples/question_answering/graphrag_simple_pipeline.py

@@ -19,7 +19,7 @@
                     retriever_config={
                         "top_k": 2,
                     },
-                    return_context=True,
+                    # return_context=True,
                 )
             )
         )

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

@@ -156,27 +156,3 @@ async def search(self, **kwargs: Any) -> RagResultModel:
             answer=result.result["generator"]["content"],
             retriever_result=context,
         )
-
-
-if __name__ == "__main__":
-    import asyncio
-    import os
-
-    os.environ["NEO4J_URI"] = "neo4j+s://demo.neo4jlabs.com"
-    os.environ["NEO4J_USER"] = "recommendations"
-    os.environ["NEO4J_PASSWORD"] = "recommendations"
-
-    runner = RagPipelineRunner.from_config_file(
-        "src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_rag_pipeline_config.json"
-    )
-    print(
-        asyncio.run(
-            runner.search(
-                query_text="Recommend a Christmas movie",
-                retriever_config={
-                    "top_k": 2,
-                },
-                return_context=True,
-            )
-        )
-    )