Move pipeline started event notification so that we have both started and failure event - add unit test

Author: stellasia

src/neo4j_graphrag/experimental/pipeline/orchestrator.py

@@ -264,9 +264,5 @@ async def run(self, data: dict[str, Any]) -> None:
         (node without any parent). Then the callback on_task_complete
         will handle the task dependencies.
         """
-        await self.event_notifier.notify_pipeline_started(self.run_id, data)
         tasks = [self.run_task(root, data) for root in self.pipeline.roots()]
         await asyncio.gather(*tasks)
-        await self.event_notifier.notify_pipeline_finished(
-            self.run_id, await self.pipeline.get_final_results(self.run_id)
-        )

src/neo4j_graphrag/experimental/pipeline/pipeline.py

@@ -571,11 +571,16 @@ async def run(self, data: dict[str, Any]) -> PipelineResult:
         return res
 
     async def _run(self, run_id: str, data: dict[str, Any]) -> PipelineResult:
+        await self.event_notifier.notify_pipeline_started(run_id, data)
         self.invalidate()
         self.validate_input_data(data)
         orchestrator = Orchestrator(self, run_id)
         await orchestrator.run(data)
-        return PipelineResult(
+        result = PipelineResult(
             run_id=orchestrator.run_id,
             result=await self.get_final_results(orchestrator.run_id),
         )
+        await self.event_notifier.notify_pipeline_finished(
+            run_id, await self.get_final_results(run_id),
+        )
+        return result

tests/unit/experimental/pipeline/test_pipeline.py

@@ -483,6 +483,23 @@ async def test_pipeline_event_notification() -> None:
         previous_ts = actual_event.timestamp
 
 
+@pytest.mark.asyncio
+async def test_pipeline_event_notification_error_in_pipeline_run() -> None:
+    callback = AsyncMock(spec=EventCallbackProtocol)
+    pipe = Pipeline(callback=callback)
+    component_a = ComponentAdd()
+    component_b = ComponentAdd()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+
+    with pytest.raises(PipelineDefinitionError):
+        await pipe.run({"a": {"number1": 1, "number2": 2}})
+    assert len(callback.await_args_list) == 2
+    assert callback.await_args_list[0][0][0].event_type == EventType.PIPELINE_STARTED
+    assert callback.await_args_list[1][0][0].event_type == EventType.PIPELINE_FAILED
+
+
 def test_event_model_no_warning(recwarn: Sized) -> None:
     event = Event(
         event_type=EventType.PIPELINE_STARTED,
@@ -557,10 +574,9 @@ async def test_pipeline_streaming_error_in_pipeline_definition() -> None:
     with pytest.raises(PipelineDefinitionError):
         async for e in pipe.stream({"a": {"number1": 1, "number2": 2}}):
             events.append(e)
-    # validation happens before pipeline run actually starts
-    # but we have the PIPELINE_FAILED event
-    assert len(events) == 1
-    assert events[0].event_type == EventType.PIPELINE_FAILED
+    assert len(events) == 2
+    assert events[0].event_type == EventType.PIPELINE_STARTED
+    assert events[1].event_type == EventType.PIPELINE_FAILED
 
 
 @pytest.mark.asyncio