Add tests

Author: stellasia

examples/customize/build_graph/pipeline/pipeline_streaming.py

@@ -2,7 +2,7 @@
 
 from neo4j_graphrag.experimental.pipeline.component import Component, DataModel
 from neo4j_graphrag.experimental.pipeline.pipeline import Pipeline
-from neo4j_graphrag.experimental.pipeline.notification import EventType
+from neo4j_graphrag.experimental.pipeline.notification import EventType, Event
 from neo4j_graphrag.experimental.pipeline.types.context import RunContext
 
 
@@ -13,6 +13,7 @@ class OutputModel(DataModel):
 
 class SlowAdder(Component):
     """A component that slowly adds numbers and reports progress"""
+
     def __init__(self, number: int) -> None:
         self.number = number
 
@@ -21,14 +22,14 @@ async def run_with_context(self, context_: RunContext, value: int) -> OutputMode
         for i in range(value):
             await asyncio.sleep(0.5)  # Simulate work
             await context_.notify(
-                message=f"Added {i+1}/{value}", 
-                data={"current": i+1, "total": value}
+                message=f"Added {i+1}/{value}", data={"current": i + 1, "total": value}
             )
         return OutputModel(result=value + self.number)
 
 
 class SlowMultiplier(Component):
     """A component that slowly multiplies numbers and reports progress"""
+
     def __init__(self, multiplier: int) -> None:
         self.multiplier = multiplier
 
@@ -37,26 +38,25 @@ async def run_with_context(self, context_: RunContext, value: int) -> OutputMode
         for i in range(3):  # Always do 3 steps
             await asyncio.sleep(0.7)  # Simulate work
             await context_.notify(
-                message=f"Multiplication step {i+1}/3", 
-                data={"step": i+1, "total": 3}
+                message=f"Multiplication step {i+1}/3", data={"step": i + 1, "total": 3}
             )
         return OutputModel(result=value * self.multiplier)
 
 
-async def main():
+async def callback(event: Event) -> None:
+    await asyncio.sleep(0.1)
+
+
+async def main() -> None:
     # Create pipeline
-    pipeline = Pipeline()
-    
+    pipeline = Pipeline(callback=callback)
+
     # Add components
     pipeline.add_component(SlowAdder(number=3), "adder")
     pipeline.add_component(SlowMultiplier(multiplier=2), "multiplier")
-    
+
     # Connect components
-    pipeline.connect(
-        "adder", 
-        "multiplier",
-        {"value": "adder.result"}
-    )
+    pipeline.connect("adder", "multiplier", {"value": "adder.result"})
 
     print("\n=== Running pipeline with streaming ===")
     # Run pipeline with streaming - see events as they happen
@@ -66,12 +66,16 @@ async def main():
         elif event.event_type == EventType.PIPELINE_FINISHED:
             print(f"Stream: Pipeline finished! Final results: {event.payload}")
         elif event.event_type == EventType.TASK_STARTED:
-            print(f"Stream: Task {event.task_name} started with inputs: {event.payload}")
+            print(
+                f"Stream: Task {event.task_name} started with inputs: {event.payload}"  # type: ignore
+            )
         elif event.event_type == EventType.TASK_PROGRESS:
-            print(f"Stream: Task {event.task_name} progress - {event.message}")
+            print(f"Stream: Task {event.task_name} progress - {event.message}")  # type: ignore
         elif event.event_type == EventType.TASK_FINISHED:
-            print(f"Stream: Task {event.task_name} finished with result: {event.payload}")
+            print(
+                f"Stream: Task {event.task_name} finished with result: {event.payload}"  # type: ignore
+            )
 
 
 if __name__ == "__main__":
-    asyncio.run(main()) 
+    asyncio.run(main())

src/neo4j_graphrag/experimental/pipeline/notification.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import asyncio
 import datetime
 import enum
 from collections.abc import Awaitable
@@ -72,12 +73,14 @@ def __call__(self, event: Event) -> Awaitable[None]: ...
 
 
 class EventNotifier:
-    def __init__(self, callback: EventCallbackProtocol | None) -> None:
-        self.callback = callback
+    def __init__(self, callbacks: list[EventCallbackProtocol]) -> None:
+        self.callbacks = callbacks
 
     async def notify(self, event: Event) -> None:
-        if self.callback:
-            await self.callback(event)
+        await asyncio.gather(
+            *[c(event) for c in self.callbacks],
+            return_exceptions=True,
+        )
 
     async def notify_pipeline_started(
         self, run_id: str, input_data: Optional[dict[str, Any]] = None

src/neo4j_graphrag/experimental/pipeline/orchestrator.py

@@ -54,7 +54,7 @@ class Orchestrator:
 
     def __init__(self, pipeline: Pipeline):
         self.pipeline = pipeline
-        self.event_notifier = EventNotifier(pipeline.callback)
+        self.event_notifier = EventNotifier(pipeline.callbacks)
         self.run_id = str(uuid.uuid4())
 
     async def run_task(self, task: TaskPipelineNode, data: dict[str, Any]) -> None:

src/neo4j_graphrag/experimental/pipeline/pipeline.py

@@ -18,7 +18,7 @@
 import warnings
 from collections import defaultdict
 from timeit import default_timer
-from typing import Any, Optional, AsyncGenerator, Callable, List
+from typing import Any, Optional, AsyncGenerator
 import asyncio
 
 from neo4j_graphrag.utils.logging import prettify
@@ -48,7 +48,10 @@
 )
 from neo4j_graphrag.experimental.pipeline.types.orchestration import RunResult
 from neo4j_graphrag.experimental.pipeline.types.context import RunContext
-from neo4j_graphrag.experimental.pipeline.notification import EventCallbackProtocol, Event
+from neo4j_graphrag.experimental.pipeline.notification import (
+    EventCallbackProtocol,
+    Event,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -118,7 +121,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.store = store or InMemoryStore()
-        self.callback = callback
+        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)
@@ -415,61 +418,60 @@ async def get_final_results(self, run_id: str) -> dict[str, Any]:
 
     async def stream(self, data: dict[str, Any]) -> AsyncGenerator[Event, None]:
         """Run the pipeline and stream events for task progress.
-        
+
         Args:
             data: Input data for the pipeline components
-            
+
         Yields:
             Event: Pipeline and task events including start, progress, and completion
         """
         # Create queue for events
         event_queue: asyncio.Queue[Event] = asyncio.Queue()
-        
-        # Store original callback
-        original_callback = self.callback
-        
-        async def callback_and_event_stream(event: Event) -> None:
+
+        async def event_stream(event: Event) -> None:
             # Put event in queue for streaming
             await event_queue.put(event)
-            # Call original callback if it exists
-            if original_callback:
-                await original_callback(event)
-        
-        # Set up event callback
-        self.callback = callback_and_event_stream
-        
+
+        # Add event streaming callback
+        self.callbacks.append(event_stream)
+
         try:
             # Start pipeline execution in background task
             run_task = asyncio.create_task(self.run(data))
-            
-            while True:
+
+            # loop until the run task is done, and we do not have
+            # any more pending tasks in queue
+            is_run_task_running = True
+            is_queue_empty = False
+            while is_run_task_running or not is_queue_empty:
                 # Wait for next event or pipeline completion
+                event_queue_getter_task = asyncio.create_task(event_queue.get())
                 done, pending = await asyncio.wait(
-                    [run_task, event_queue.get()],
-                    return_when=asyncio.FIRST_COMPLETED
+                    [run_task, event_queue_getter_task],
+                    return_when=asyncio.FIRST_COMPLETED,
                 )
-                
-                # Pipeline finished
-                if run_task in done:
-                    if run_task.exception():
-                        raise run_task.exception()
-                    # Drain any remaining events
-                    while not event_queue.empty():
-                        yield await event_queue.get()
-                    break
-                    
-                # Got an event from queue
-                event_future = next(f for f in done if f != run_task)
-                try:
-                    event = event_future.result()
-                    yield event
-                except Exception as e:
-                    logger.error(f"Error processing event: {e}")
-                    raise
-        
+
+                is_run_task_running = run_task not in done
+                is_queue_empty = event_queue.empty()
+
+                for event_future in done:
+                    if event_future == run_task:
+                        continue
+                    yield event_future.result()  # type: ignore
+
+                # cancel remaining task
+                event_queue_getter_task.cancel()
+
+            # # Drain any remaining events
+            # while not event_queue.empty():
+            #     yield await event_queue.get()
+            # Pipeline finished
+            if run_task.exception():
+                raise run_task.exception()  # type: ignore
+
         finally:
             # Restore original callback
-            self.callback = original_callback
+            self.callbacks.remove(event_stream)
 
     async def run(self, data: dict[str, Any]) -> PipelineResult:
         logger.debug("PIPELINE START")

tests/unit/experimental/pipeline/components.py

@@ -12,6 +12,8 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+import asyncio
+
 from neo4j_graphrag.experimental.pipeline import Component, DataModel
 from neo4j_graphrag.experimental.pipeline.types.context import RunContext
 
@@ -52,3 +54,12 @@ async def run_with_context(
             message="my message", data={"number1": number1, "number2": number2}
         )
         return IntResultModel(result=number1 * number2)
+
+
+class SlowComponentMultiply(Component):
+    def __init__(self, sleep: float = 1.0) -> None:
+        self.sleep = sleep
+
+    async def run(self, number1: int, number2: int = 2) -> IntResultModel:
+        await asyncio.sleep(self.sleep)
+        return IntResultModel(result=number1 * number2)

tests/unit/experimental/pipeline/test_pipeline.py

@@ -39,6 +39,7 @@
     ComponentNoParam,
     ComponentPassThrough,
     StringResultModel,
+    SlowComponentMultiply,
 )
 
 
@@ -491,3 +492,97 @@ def test_event_model_no_warning(recwarn: Sized) -> None:
     )
     assert event.timestamp is not None
     assert len(recwarn) == 0
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_no_user_callback_happy_path() -> None:
+    pipe = Pipeline()
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    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
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_with_user_callback_happy_path() -> None:
+    callback = AsyncMock()
+    pipe = Pipeline(callback=callback)
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert len(callback.call_args_list) == 2
+    assert len(pipe.callbacks) == 1
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_very_long_running_user_callback() -> None:
+    async def callback(event: Event) -> None:
+        await asyncio.sleep(2)
+
+    pipe = Pipeline(callback=callback)
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert len(pipe.callbacks) == 1
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_very_long_running_pipeline() -> None:
+    slow_component = SlowComponentMultiply()
+    pipe = Pipeline()
+    pipe.add_component(slow_component, "slow_component")
+    events = []
+    async for e in pipe.stream({"slow_component": {"number1": 1, "number2": 2}}):
+        events.append(e)
+    assert len(events) == 4
+    last_event = events[-1]
+    assert last_event.event_type == EventType.PIPELINE_FINISHED
+    assert last_event.payload == {"slow_component": {"result": 2}}
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_error_in_pipeline_definition() -> None:
+    pipe = Pipeline()
+    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"})
+    events = []
+    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
+    assert len(events) == 0
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_error_in_component() -> None:
+    component = ComponentMultiply()
+    pipe = Pipeline()
+    pipe.add_component(component, "component")
+    events = []
+    with pytest.raises(TypeError):
+        async for e in pipe.stream({"component": {"number1": None, "number2": 2}}):
+            events.append(e)
+    assert len(events) == 2
+    assert events[0].event_type == EventType.PIPELINE_STARTED
+    assert events[1].event_type == EventType.TASK_STARTED
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_error_in_user_callback() -> None:
+    async def callback(event: Event) -> None:
+        raise Exception("error in callback")
+
+    pipe = Pipeline(callback=callback)
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert len(pipe.callbacks) == 1