Merge 19fd31c into b0b33e3

Author: masci

src/workflows/handler.py

@@ -7,7 +7,7 @@
 from typing import Any, AsyncGenerator
 
 from .context import Context
-from .errors import WorkflowDone
+from .errors import WorkflowDone, WorkflowRuntimeError
 from .events import Event, StopEvent
 from .types import RunResultT
 from .utils import BUSY_WAIT_DELAY
@@ -24,6 +24,7 @@ def __init__(
         super().__init__(*args, **kwargs)
         self.run_id = run_id
         self._ctx = ctx
+        self._all_events_consumed = False
 
     @property
     def ctx(self) -> Context | None:
@@ -36,15 +37,60 @@ def is_done(self) -> bool:
         return self.done()
 
     async def stream_events(self) -> AsyncGenerator[Event, None]:
+        """
+        Stream events from the workflow execution as they occur.
+
+        This method provides real-time access to events generated during workflow
+        execution, allowing for monitoring and processing of intermediate results.
+        Events are yielded in the order they are generated by the workflow.
+
+        The stream includes all events written to the context's streaming queue,
+        and terminates when a StopEvent is encountered, indicating the workflow
+        has completed.
+
+        Returns:
+            AsyncGenerator[Event, None]: An async generator that yields Event objects
+                as they are produced by the workflow.
+
+        Raises:
+            ValueError: If the context is not set on the handler.
+            WorkflowRuntimeError: If all events have already been consumed by a
+                previous call to stream_events() on the same handler instance.
+
+        Examples:
+            ```python
+            handler = workflow.run()
+
+            # Stream and process events in real-time
+            async for event in handler.stream_events():
+                if isinstance(event, StopEvent):
+                    print(f"Workflow completed with result: {event.result}")
+                else:
+                    print(f"Received event: {event}")
+
+            # Get final result
+            result = await handler
+            ```
+
+        Note:
+            Events can only be streamed once per handler instance. Subsequent
+            calls to stream_events() will raise a WorkflowRuntimeError.
+        """
         if self.ctx is None:
             raise ValueError("Context is not set!")
 
+        # Check if we already consumed all the streamed events
+        if self._all_events_consumed:
+            msg = "All the streamed events have already been consumed."
+            raise WorkflowRuntimeError(msg)
+
         while True:
             ev = await self.ctx.streaming_queue.get()
 
             yield ev
 
             if isinstance(ev, StopEvent):
+                self._all_events_consumed = True
                 break
 
     async def run_step(self) -> list[Event] | None:

tests/test_handler.py

@@ -6,6 +6,7 @@
 import pytest
 
 from workflows.context import Context
+from workflows.errors import WorkflowRuntimeError
 from workflows.handler import WorkflowHandler
 
 
@@ -42,3 +43,18 @@ async def test_run_step_no_stepwise() -> None:
         match="Workflow must be created passing stepwise=True to call this method.",
     ):
         await h.run_step()
+
+
+@pytest.mark.asyncio
+async def test_stream_events_consume_only_once() -> None:
+    ctx = mock.MagicMock(spec=Context)
+
+    h = WorkflowHandler(ctx=ctx)
+    h._all_events_consumed = True
+
+    with pytest.raises(
+        WorkflowRuntimeError,
+        match="All the streamed events have already been consumed.",
+    ):
+        async for _ in h.stream_events():
+            pass

tests/test_streaming.py

@@ -116,6 +116,24 @@ async def test_multiple_sequential_streams() -> None:
     await r
 
 
+@pytest.mark.asyncio
+async def test_consume_only_once() -> None:
+    wf = StreamingWorkflow()
+    handler = wf.run()
+
+    async for _ in handler.stream_events():
+        pass
+
+    with pytest.raises(
+        WorkflowRuntimeError,
+        match="All the streamed events have already been consumed.",
+    ):
+        async for _ in handler.stream_events():
+            pass
+
+    await handler
+
+
 @pytest.mark.asyncio
 async def test_multiple_ongoing_streams() -> None:
     wf = StreamingWorkflow()