Pipeline streaming

Author: stellasia

examples/customize/build_graph/pipeline/pipeline_streaming.py

@@ -0,0 +1,77 @@
+import asyncio
+
+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.types.context import RunContext
+
+
+# Define some example components with progress notifications
+class OutputModel(DataModel):
+    result: int
+
+
+class SlowAdder(Component):
+    """A component that slowly adds numbers and reports progress"""
+    def __init__(self, number: int) -> None:
+        self.number = number
+
+    async def run_with_context(self, context_: RunContext, value: int) -> OutputModel:
+        # Simulate work with progress updates
+        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}
+            )
+        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
+
+    async def run_with_context(self, context_: RunContext, value: int) -> OutputModel:
+        # Simulate work with progress updates
+        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}
+            )
+        return OutputModel(result=value * self.multiplier)
+
+
+async def main():
+    # Create pipeline
+    pipeline = Pipeline()
+    
+    # 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"}
+    )
+
+    print("\n=== Running pipeline with streaming ===")
+    # Run pipeline with streaming - see events as they happen
+    async for event in pipeline.stream({"adder": {"value": 2}}):
+        if event.event_type == EventType.PIPELINE_STARTED:
+            print("Stream: Pipeline started!")
+        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}")
+        elif event.event_type == EventType.TASK_PROGRESS:
+            print(f"Stream: Task {event.task_name} progress - {event.message}")
+        elif event.event_type == EventType.TASK_FINISHED:
+            print(f"Stream: Task {event.task_name} finished with result: {event.payload}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main()) 

src/neo4j_graphrag/experimental/pipeline/pipeline.py

@@ -18,7 +18,8 @@
 import warnings
 from collections import defaultdict
 from timeit import default_timer
-from typing import Any, Optional
+from typing import Any, Optional, AsyncGenerator, Callable, List
+import asyncio
 
 from neo4j_graphrag.utils.logging import prettify
 
@@ -47,7 +48,7 @@
 )
 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
+from neo4j_graphrag.experimental.pipeline.notification import EventCallbackProtocol, Event
 
 
 logger = logging.getLogger(__name__)
@@ -412,6 +413,64 @@ def validate_parameter_mapping_for_task(self, task: TaskPipelineNode) -> bool:
     async def get_final_results(self, run_id: str) -> dict[str, Any]:
         return await self.final_results.get(run_id)  # type: ignore[no-any-return]
 
+    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:
+            # 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
+        
+        try:
+            # Start pipeline execution in background task
+            run_task = asyncio.create_task(self.run(data))
+            
+            while True:
+                # Wait for next event or pipeline completion
+                done, pending = await asyncio.wait(
+                    [run_task, event_queue.get()],
+                    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
+        
+        finally:
+            # Restore original callback
+            self.callback = original_callback
+
     async def run(self, data: dict[str, Any]) -> PipelineResult:
         logger.debug("PIPELINE START")
         start_time = default_timer()