Repository: run-llama/llama_index
Pull request: feat: Allow streaming events from steps
# Description

Allow users to consume a stream of events produced by steps during a workflow execution.

Fixes #15394 

Main takeaways:
- We expose the API to stream events from steps through the session instance: `ctx.session.write_stream_event(ev)`
- We expose the events to the user through a method of the workflow: `wf.stream_events()`
- If the workflow was `run` more than once, streaming will produce a meaningful error

Alternatives:
- We expose the API to stream events from steps through a built-in "service", depending on [this ongoing PR](https://github.com/run-llama/llama_index/pull/15325)
- We expose the API to stream events through `self` (this won't work for unbound steps implemented as free functions)

Further development
- We might introduce "topics" in the stream of events, so that consumers can filter what to get in case multiple steps are streaming data: `async for llm_ev in wf.stream_events(topic="llm")`
- If we let users manage sessions explicitly, we can support streaming from multiple `workflow.run`s at the same time.

Example:
```py
class StreamingWorkflow(Workflow):
    @step
    async def chat(self, ctx: Context, ev: StartEvent) -> StopEvent:
        llm = OpenAI()
        messages = [
            ChatMessage(
                role=MessageRole.SYSTEM,
                content="You are a pirate with a colorful personality",
            ),
            ChatMessage(role=MessageRole.USER, content="Who is Paul Graham?"),
        ]
        gen = await llm.astream_chat(messages)
        async for resp in gen:
            # Allow the workflow to stream this piece of response
            ctx.session.write_stream_event(Event(msg=resp.delta))

        return StopEvent(result=None)


#
# in a main loop
#
wf = StreamingWorkflow()

# run the workflow in the background instead of awaiting it
r = asyncio.create_task(wf.run())

# consume events possibly produced by steps in the workflow
# the loop will exit when the workflow is done in any case
async for ev in wf.stream_events():
    print(ev.msg)

# await the workflow result and access its results as usual
await r
```










Problem statement:
[Feature Request]: Allow streaming of progress during workflow execution
### Feature Description

During execution of a long-running workflow, it would be useful to expose the progress somehow

One idea for doing this is introducing a special `ProgressEvent` that the user can emit. Or alternatively, a new method like `self.send_progress("progress")` ?

Then, the usage might look like
```python
class LongWorkflow(Workflow):

  @step()
  async def my_step(self, ev: StartEvent) -> StopEvent:
    self.send_event(ProgressEvent(progress="about to run"))
    await asyncio.sleep(10)
    self.send_event(ProgressEvent(progress="done running"))
    return StopEvent(result="All done!")

gen = await workflow.run_with_progress(...)
async for progress in gen:
  print(progress) # would print 'about to run' -> 'done running' -> 'All Done!'
```

### Reason

Exposing inner actions is pretty useful for long-running workflows

### Value of Feature

_No response_

Hints from issue discussion:

Commit message: add proper testing

Patch to review:
diff --git a/llama-index-core/llama_index/core/workflow/session.py b/llama-index-core/llama_index/core/workflow/session.py
index 1d48734592..1c868c7664 100644
--- a/llama-index-core/llama_index/core/workflow/session.py
+++ b/llama-index-core/llama_index/core/workflow/session.py
@@ -24,6 +24,8 @@ def __init__(self, workflow: "Workflow") -> None:
         self._root_context = Context(self)
         # Context management
         self._step_to_context: Dict[str, Context] = {}
+        # Streaming machinery
+        self._streamed_events: asyncio.Queue = asyncio.Queue()
 
     def send_event(self, message: Event, step: Optional[str] = None) -> None:
         """Sends an event to a specific step in the workflow.
@@ -52,6 +54,9 @@ def send_event(self, message: Event, step: Optional[str] = None) -> None:
 
         self._broker_log.append(message)
 
+    def write_stream_event(self, ev: Optional[Event]) -> None:
+        self._streamed_events.put_nowait(ev)
+
     def get_context(self, step_name: str) -> Context:
         """Get the global context for this workflow.
 
@@ -66,3 +71,7 @@ def get_context(self, step_name: str) -> Context:
     def get_result(self) -> Any:
         """Returns the result of the workflow."""
         return self._retval
+
+    @property
+    def streaming_queue(self) -> asyncio.Queue:
+        return self._streamed_events
diff --git a/llama-index-core/llama_index/core/workflow/workflow.py b/llama-index-core/llama_index/core/workflow/workflow.py
index ac4cfdf8c5..80802db4c8 100644
--- a/llama-index-core/llama_index/core/workflow/workflow.py
+++ b/llama-index-core/llama_index/core/workflow/workflow.py
@@ -1,7 +1,7 @@
 import asyncio
 import functools
 import warnings
-from typing import Any, Callable, Dict, Optional, Set
+from typing import Any, Callable, Dict, Optional, AsyncGenerator, Set
 
 from llama_index.core.instrumentation import get_dispatcher
 from llama_index.core.workflow.decorators import StepConfig, step
@@ -48,6 +48,29 @@ def __init__(
         # Services management
         self._service_manager = service_manager
 
+    async def stream_events(self) -> AsyncGenerator[Event, None]:
+        # In the typical streaming use case, `run()` is not awaited but wrapped in a asyncio.Task. Since we'll be
+        # consuming events produced by `run()`, we must give its Task the chance to run before entering the dequeueing
+        # loop.
+        await asyncio.sleep(0)
+
+        if len(self._sessions) > 1:
+            # We can't possibly know from what session we should stream events, raise an error.
+            msg = (
+                "This workflow has multiple session running concurrently and cannot stream events. "
+                "To be able to stream events, make sure you call `run()` on this workflow only once."
+            )
+            raise WorkflowRuntimeError(msg)
+
+        # Enter the dequeuing loop.
+        session = next(iter(self._sessions))
+        while True:
+            ev = await session.streaming_queue.get()
+            if ev is None:  # FIXME: Should we use a special event?
+                break
+
+            yield ev
+
     @classmethod
     def add_step(cls, func: Callable) -> None:
         """Adds a free function as step for this workflow instance.
@@ -299,6 +322,8 @@ def is_done(self) -> bool:
     async def _done(self, ctx: Context, ev: StopEvent) -> None:
         """Tears down the whole workflow and stop execution."""
         ctx.session._retval = ev.result or None
+        ctx.session.write_stream_event(None)  # FIXME: send a special event instead?
+
         # Signal we want to stop the workflow
         raise WorkflowDone
 
diff --git a/llama-index-core/tests/workflow/test_streaming.py b/llama-index-core/tests/workflow/test_streaming.py
new file mode 100644
index 0000000000..423ca20639
--- /dev/null
+++ b/llama-index-core/tests/workflow/test_streaming.py
@@ -0,0 +1,53 @@
+import asyncio
+
+import pytest
+
+from llama_index.llms.openai import OpenAI
+
+from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.core.workflow.context import Context
+from llama_index.core.workflow.decorators import step
+from llama_index.core.workflow.events import Event, StartEvent, StopEvent
+from llama_index.core.workflow.workflow import Workflow
+from llama_index.core.workflow.errors import WorkflowRuntimeError
+
+
+class StreamingWorkflow(Workflow):
+    @step
+    async def chat(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        llm = OpenAI()
+        messages = [
+            ChatMessage(
+                role=MessageRole.SYSTEM,
+                content="You are a pirate with a colorful personality",
+            ),
+            ChatMessage(role=MessageRole.USER, content="Who is Paul Graham?"),
+        ]
+        gen = await llm.astream_chat(messages)
+        async for resp in gen:
+            ctx.session.write_stream_event(Event(msg=resp.delta))
+
+        return StopEvent(result=None)
+
+
+@pytest.mark.asyncio()
+async def test_e2e():
+    wf = StreamingWorkflow()
+    r = asyncio.create_task(wf.run())
+
+    async for ev in wf.stream_events():
+        assert "msg" in ev
+
+    await r
+
+
+@pytest.mark.asyncio()
+async def test_too_many_runs():
+    wf = StreamingWorkflow()
+    r = asyncio.gather(wf.run(), wf.run())
+    with pytest.raises(
+        WorkflowRuntimeError,
+        match="This workflow has multiple session running concurrently",
+    ):
+        async for ev in wf.stream_events():
+            pass


Code review comments:
diff_hunk: @@ -0,0 +1,53 @@
+import asyncio
+
+import pytest
+
+from llama_index.llms.openai import OpenAI
+
+from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.core.workflow.context import Context
+from llama_index.core.workflow.decorators import step
+from llama_index.core.workflow.events import Event, StartEvent, StopEvent
+from llama_index.core.workflow.workflow import Workflow
+from llama_index.core.workflow.errors import WorkflowRuntimeError
+
+
+class StreamingWorkflow(Workflow):
+    @step
+    async def chat(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        llm = OpenAI()
+        messages = [
+            ChatMessage(
+                role=MessageRole.SYSTEM,
+                content="You are a pirate with a colorful personality",
+            ),
+            ChatMessage(role=MessageRole.USER, content="Who is Paul Graham?"),
+        ]
+        gen = await llm.astream_chat(messages)
+        async for resp in gen:
+            ctx.session.write_stream_event(Event(msg=resp.delta))
path: llama-index-core/tests/workflow/test_streaming.py
text: @user1:
curious on the decision to do `ctx.session.write_stream_event()` vs. `ctx.write_stream_event()` ? Just avoiding overloading the top-level context? Do we foresee more methods attached to the session that its worth exposing now?

@author:
session was already exposed because of `send_event` so I thought I would keep the "family feeling" but honestly I'm not sure, the first version was indeed `ctx.write_stream_event`...

@user1:
ohhh I think I missed that send_event wasn't on the context, but was on the session 👀 

@author:
I think we have two options:
- We use `Context` as a facade for everything session-related, until it's time to make session part of the public API
- We start exposing users to the concept of session very gently with only two fairly advanced use cases (manual sending events and streaming)

I would go with option n.2 (the one in this PR) if we are positive the session will be eventually made more prominent

@user1:
Yea i agree, option 2 makes sense here

diff_hunk: @@ -48,6 +48,29 @@ def __init__(
         # Services management
         self._service_manager = service_manager
 
+    async def stream_events(self) -> AsyncGenerator[Event, None]:
+        # In the typical streaming use case, `run()` is not awaited but wrapped in a asyncio.Task. Since we'll be
+        # consuming events produced by `run()`, we must give its Task the chance to run before entering the dequeueing
+        # loop.
+        await asyncio.sleep(0)
+
+        if len(self._sessions) > 1:
+            # We can't possibly know from what session we should stream events, raise an error.
+            msg = (
+                "This workflow has multiple session running concurrently and cannot stream events. "
+                "To be able to stream events, make sure you call `run()` on this workflow only once."
+            )
+            raise WorkflowRuntimeError(msg)
+
+        # Enter the dequeuing loop.
+        session = next(iter(self._sessions))
+        while True:
+            ev = await session.streaming_queue.get()
+            if ev is None:  # FIXME: Should we use a special event?
path: llama-index-core/llama_index/core/workflow/workflow.py
text: @user1:
Since we know the last event is a `StopEvent`, shouldn't we just stop then?

diff_hunk: @@ -52,6 +54,9 @@ def send_event(self, message: Event, step: Optional[str] = None) -> None:
 
         self._broker_log.append(message)
 
+    def write_stream_event(self, ev: Optional[Event]) -> None:
path: llama-index-core/llama_index/core/workflow/session.py
text: @user1:
A bit of a nit, but just wondering if others think that `write_stream_event` may confuse some to think that this method produces a "special" kind of event?

A possible alternative could be something like `publish_event_to_stream()` or `write_event_to_stream`.

And perhaps instead of `self._streamed_events` it can be `self._event_stream`.

@author:
Changed to `write_event_to_stream` and `self._streaming_queue`