Make close function async

stellasia · Dec 6, 2024 · 82b4598 · 82b4598
1 parent 8834351
commit 82b4598
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/neo4j_graphrag/experimental/pipeline/config/pipeline_config.py b/src/neo4j_graphrag/experimental/pipeline/config/pipeline_config.py
@@ -152,7 +152,7 @@ def parse(
     def get_run_params(self, user_input: dict[str, Any]) -> dict[str, Any]:
         return user_input
 
-    def close(self) -> None:
+    async def close(self) -> None:
         drivers = self._global_data.get("neo4j_config", {})
         for driver_name in drivers:
             driver = drivers[driver_name]

diff --git a/src/neo4j_graphrag/experimental/pipeline/config/runner.py b/src/neo4j_graphrag/experimental/pipeline/config/runner.py
@@ -123,10 +123,10 @@ async def run(self, user_input: dict[str, Any]) -> PipelineResult:
         )
         result = await self.pipeline.run(data=run_param)
         if self.do_cleaning:
-            self.close()
+            await self.close()
         return result
 
-    def close(self) -> None:
+    async def close(self) -> None:
         logger.debug("PIPELINE_RUNNER: cleaning up (closing instantiated drivers...)")
         if self.config:
-            self.config.close()
+            await self.config.close()