mypy

stellasia · Jan 21, 2025 · ee16e54 · ee16e54
1 parent bca341a
commit ee16e54
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/neo4j_graphrag/experimental/pipeline/notification.py b/src/neo4j_graphrag/experimental/pipeline/notification.py
@@ -29,11 +29,11 @@ class EventType(enum.Enum):
 
     @property
     def is_pipeline_event(self) -> bool:
-        return self in [self.PIPELINE_STARTED, self.PIPELINE_FINISHED]
+        return self in [EventType.PIPELINE_STARTED, EventType.PIPELINE_FINISHED]
 
     @property
     def is_task_event(self) -> bool:
-        return self in [self.TASK_STARTED, self.TASK_FINISHED]
+        return self in [EventType.TASK_STARTED, EventType.TASK_FINISHED]
 
 
 class Event(BaseModel):

diff --git a/src/neo4j_graphrag/experimental/pipeline/pipeline.py b/src/neo4j_graphrag/experimental/pipeline/pipeline.py
@@ -402,7 +402,7 @@ def validate_parameter_mapping_for_task(self, task: TaskPipelineNode) -> bool:
         return True
 
     async def get_final_results(self, run_id: str) -> dict[str, Any]:
-        return await self.final_results.get(run_id)
+        return await self.final_results.get(run_id)  # type: ignore[no-any-return]
 
     async def run(self, data: dict[str, Any]) -> PipelineResult:
         logger.debug("PIPELINE START")