From 82b45988bee99427714914595fa7c8029bead858 Mon Sep 17 00:00:00 2001 From: estelle Date: Fri, 6 Dec 2024 10:58:20 +0100 Subject: [PATCH] Make close function async --- .../experimental/pipeline/config/pipeline_config.py | 2 +- src/neo4j_graphrag/experimental/pipeline/config/runner.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/neo4j_graphrag/experimental/pipeline/config/pipeline_config.py b/src/neo4j_graphrag/experimental/pipeline/config/pipeline_config.py index 3737b2c0..c3871179 100644 --- 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 index c8b7648d..a1a22585 100644 --- 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()