Send pipeline failed events to the event callback

Author: stellasia

src/neo4j_graphrag/experimental/pipeline/notification.py

@@ -86,6 +86,12 @@ class EventNotifier:
     def __init__(self, callbacks: list[EventCallbackProtocol]) -> None:
         self.callbacks = callbacks
 
+    def add_callback(self, callback: EventCallbackProtocol) -> None:
+        self.callbacks.append(callback)
+
+    def remove_callback(self, callback: EventCallbackProtocol) -> None:
+        self.callbacks.remove(callback)
+
     async def notify(self, event: Event) -> None:
         await asyncio.gather(
             *[c(event) for c in self.callbacks],
@@ -114,6 +120,17 @@ async def notify_pipeline_finished(
         )
         await self.notify(event)
 
+    async def notify_pipeline_failed(
+        self, run_id: str, message: Optional[str] = None
+    ) -> None:
+        event = PipelineEvent(
+            event_type=EventType.PIPELINE_FAILED,
+            run_id=run_id,
+            message=message,
+            payload=None,
+        )
+        await self.notify(event)
+
     async def notify_task_started(
         self,
         run_id: str,

src/neo4j_graphrag/experimental/pipeline/orchestrator.py

@@ -52,10 +52,10 @@ class Orchestrator:
     (checking that all dependencies are met), and run them.
     """
 
-    def __init__(self, pipeline: Pipeline):
+    def __init__(self, pipeline: Pipeline, run_id: Optional[str] = None):
         self.pipeline = pipeline
-        self.event_notifier = EventNotifier(pipeline.callbacks)
-        self.run_id = str(uuid.uuid4())
+        self.event_notifier = self.pipeline.event_notifier
+        self.run_id = run_id or str(uuid.uuid4())
 
     async def run_task(self, task: TaskPipelineNode, data: dict[str, Any]) -> None:
         """Get inputs and run a specific task. Once the task is done,

src/neo4j_graphrag/experimental/pipeline/pipeline.py

@@ -21,6 +21,8 @@
 from timeit import default_timer
 from typing import Any, AsyncGenerator, Optional
 
+import uuid
+
 from neo4j_graphrag.utils.logging import prettify
 
 try:
@@ -41,6 +43,7 @@
     EventCallbackProtocol,
     EventType,
     PipelineEvent,
+    EventNotifier,
 )
 from neo4j_graphrag.experimental.pipeline.orchestrator import Orchestrator
 from neo4j_graphrag.experimental.pipeline.pipeline_graph import (
@@ -124,7 +127,6 @@ def __init__(
     ) -> None:
         super().__init__()
         self.store = store or InMemoryStore()
-        self.callbacks = [callback] if callback else []
         self.final_results = InMemoryStore()
         self.is_validated = False
         self.param_mapping: dict[str, dict[str, dict[str, str]]] = defaultdict(dict)
@@ -139,6 +141,7 @@ def __init__(
         }
         """
         self.missing_inputs: dict[str, list[str]] = defaultdict()
+        self.event_notifier = EventNotifier([callback] if callback else [])
 
     @classmethod
     def from_template(
@@ -514,7 +517,7 @@ async def event_stream(event: Event) -> None:
             await event_queue.put(event)
 
         # Add event streaming callback
-        self.callbacks.append(event_stream)
+        self.event_notifier.add_callback(event_stream)
 
         event_queue_getter_task = None
         try:
@@ -546,34 +549,36 @@ async def event_stream(event: Event) -> None:
                     yield event  # type: ignore
 
             if exc := run_task.exception():
-                yield PipelineEvent(
-                    event_type=EventType.PIPELINE_FAILED,
-                    # run_id is null if pipeline fails before even starting
-                    # ie during pipeline validation
-                    run_id=run_id or "",
-                    message=str(exc),
-                )
                 if raise_exception:
                     raise exc
 
         finally:
             # Restore original callback
-            self.callbacks.remove(event_stream)
+            self.event_notifier.remove_callback(event_stream)
             if event_queue_getter_task and not event_queue_getter_task.done():
                 event_queue_getter_task.cancel()
 
     async def run(self, data: dict[str, Any]) -> PipelineResult:
-        logger.debug("PIPELINE START")
         start_time = default_timer()
+        run_id = str(uuid.uuid4())
+        logger.debug(f"PIPELINE START with {run_id=}")
+        try:
+            res = await self._run(run_id, data)
+        except Exception as e:
+            await self.event_notifier.notify_pipeline_failed(
+                run_id,
+                message=f"Pipeline failed with error {e}",
+            )
+            raise e
+        end_time = default_timer()
+        logger.debug(f"PIPELINE FINISHED {run_id} in {end_time - start_time}s")
+        return res
+
+    async def _run(self, run_id: str, data: dict[str, Any]) -> PipelineResult:
         self.invalidate()
         self.validate_input_data(data)
-        orchestrator = Orchestrator(self)
-        logger.debug(f"PIPELINE ORCHESTRATOR: {orchestrator.run_id}")
+        orchestrator = Orchestrator(self, run_id)
         await orchestrator.run(data)
-        end_time = default_timer()
-        logger.debug(
-            f"PIPELINE FINISHED {orchestrator.run_id} in {end_time - start_time}s"
-        )
         return PipelineResult(
             run_id=orchestrator.run_id,
             result=await self.get_final_results(orchestrator.run_id),

tests/unit/experimental/pipeline/test_pipeline.py

@@ -503,7 +503,7 @@ async def test_pipeline_streaming_no_user_callback_happy_path() -> None:
     assert len(events) == 2
     assert events[0].event_type == EventType.PIPELINE_STARTED
     assert events[1].event_type == EventType.PIPELINE_FINISHED
-    assert len(pipe.callbacks) == 0
+    assert len(pipe.event_notifier.callbacks) == 0
 
 
 @pytest.mark.asyncio
@@ -515,7 +515,7 @@ async def test_pipeline_streaming_with_user_callback_happy_path() -> None:
         events.append(e)
     assert len(events) == 2
     assert len(callback.call_args_list) == 2
-    assert len(pipe.callbacks) == 1
+    assert len(pipe.event_notifier.callbacks) == 1
 
 
 @pytest.mark.asyncio
@@ -528,7 +528,7 @@ async def callback(event: Event) -> None:
     async for e in pipe.stream({}):
         events.append(e)
     assert len(events) == 2
-    assert len(pipe.callbacks) == 1
+    assert len(pipe.event_notifier.callbacks) == 1
 
 
 @pytest.mark.asyncio
@@ -559,9 +559,9 @@ async def test_pipeline_streaming_error_in_pipeline_definition() -> None:
             events.append(e)
     # validation happens before pipeline run actually starts
     # but we have the PIPELINE_FAILED event
+    print(events)
     assert len(events) == 1
     assert events[0].event_type == EventType.PIPELINE_FAILED
-    assert events[0].run_id == ""
 
 
 @pytest.mark.asyncio
@@ -573,6 +573,7 @@ async def test_pipeline_streaming_error_in_component() -> None:
     with pytest.raises(TypeError):
         async for e in pipe.stream({"component": {"number1": None, "number2": 2}}):
             events.append(e)
+    print(events)
     assert len(events) == 3
     assert events[0].event_type == EventType.PIPELINE_STARTED
     assert events[1].event_type == EventType.TASK_STARTED
@@ -589,4 +590,4 @@ async def callback(event: Event) -> None:
     async for e in pipe.stream({}):
         events.append(e)
     assert len(events) == 2
-    assert len(pipe.callbacks) == 1
+    assert len(pipe.event_notifier.callbacks) == 1