[Bug]: Breaking Change Occurred in Workflow Streaming b/w `0.11.7` & `0.11.20`

[brycecf]

Bug Description

Workflow streaming has a breaking change at some point between 0.11.7 & 0.11.20. Code provided in Steps to Reproduce works in 0.11.7 and produces the following output:

Running workflow
Step one is happening
0
Step 0 is happening
1
Step 1 is happening
2
Step 2 is happening
3
Step 3 is happening
4
Step 4 is happening
Second step is happening
Step three is happening
Final result Workflow complete.

Python 3.11.9 was used both times.

Version

0.11.20

Steps to Reproduce

Run the following code using Python 3.11.9:

from llama_index.core.workflow import (
    StartEvent,
    StopEvent,
    Workflow,
    step,
    Event,
    Context,
)
import asyncio
import time


class FirstEvent(Event):
    first_output: str


class SecondEvent(Event):
    second_output: str


class ProgressEvent(Event):
    msg: str


class MyWorkflow(Workflow):
    @step
    async def step_one(self, ctx: Context, ev: StartEvent) -> FirstEvent:
        ctx.write_event_to_stream(ProgressEvent(msg="Step one is happening"))
        await asyncio.sleep(0)
        for i in range(10):
            time.sleep(1)
            print(i)
            ctx.write_event_to_stream(ProgressEvent(msg=f"Step {i} is happening"))
            await asyncio.sleep(0)
        return FirstEvent(first_output="First step complete.")

    @step
    async def step_two(self, ctx: Context, ev: FirstEvent) -> SecondEvent:
        ctx.write_event_to_stream(ProgressEvent(msg=f"Second step is happening"))
        return SecondEvent(second_output="Second step complete.")

    @step
    async def step_three(self, ctx: Context, ev: SecondEvent) -> StopEvent:
        ctx.write_event_to_stream(ProgressEvent(msg="Step three is happening"))
        return StopEvent(result="Workflow complete.")


async def main():
    w = MyWorkflow(timeout=30, verbose=False)
    r = asyncio.create_task(w.run())
    print('Running workflow')

    async for ev in w.stream_events():
        if isinstance(ev, ProgressEvent):
            print(ev.msg)

    final_result = await r
    print("Final result", final_result)


if __name__ == "__main__":
    asyncio.run(main())

Relevant Logs/Tracbacks

Traceback (most recent call last):
  File "/Users/me/demo/src/test.py", line 64, in <module>
    asyncio.run(main())
  File "/Users/me/.pyenv/versions/3.11.9/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/Users/me/.pyenv/versions/3.11.9/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/.pyenv/versions/3.11.9/lib/python3.11/asyncio/base_events.py", line 654, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/Users//Users/me/demo/src/test.py", line 52, in main
    r = asyncio.create_task(w.run())
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/.pyenv/versions/3.11.9/lib/python3.11/asyncio/tasks.py", line 384, in create_task
    task = loop.create_task(coro)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/.pyenv/versions/3.11.9/lib/python3.11/asyncio/base_events.py", line 437, in create_task
    task = tasks.Task(coro, loop=self, name=name, context=context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: a coroutine was expected, got <WorkflowHandler pending cb=[Dispatcher.span.<locals>.wrapper.<locals>.handle_future_result(span_id='Workflow.run...-545c508246a8', bound_args=<BoundArguments ()>, instance=<__main__.MyW...t 0x127debf10>, context=<_contextvars...t 0x12671bd00>)() at /Users/me/Library/Caches/pypoetry/virtualenvs/demo-n-PaIroy-py3.11/lib/python3.11/site-packages/llama_index/core/instrumentation/dispatcher.py:273]>

[dosubot]

The error you're encountering seems to be related to a change in how the Workflow class handles coroutine execution between versions 0.11.7 and 0.11.20. The traceback indicates that the run method is returning a WorkflowHandler instead of a coroutine, which is causing the TypeError.

In version 0.11.20, the run method of the Workflow class might have been updated to return a WorkflowHandler directly, which is not a coroutine. This change would require you to adjust how you initiate the workflow execution. Instead of using asyncio.create_task(w.run()), you should directly call w.run() and handle the WorkflowHandler it returns.

Here's how you can modify your main function to accommodate this change:

async def main():
    w = MyWorkflow(timeout=30, verbose=False)
    handler = w.run()  # Directly call run() to get the WorkflowHandler
    print('Running workflow')

    async for ev in w.stream_events():
        if isinstance(ev, ProgressEvent):
            print(ev.msg)

    final_result = await handler  # Await the handler to get the final result
    print("Final result", final_result)

This adjustment should resolve the TypeError by correctly handling the WorkflowHandler returned by the run method.

To continue talking to Dosu, mention @dosu.

[logan-markewich]

Indeed it did -- it now hides all the asyncio machinery. Docs are updated

(and dosu is correct!)

[logan-markewich]

Now each run has its own compartmentalized context as well, which is really nice. You can also serialize and re-use the context (need to add this to docs still, but its very handy)

from llama_index.core.workflow import Context
from llama_index.core.workflow.context_serializers import JsonSerializer, JsonPickleSerializer

handler = workflow.run()
result = await handler
serialized_context = handler.ctx.to_dict(serializer=JsonSerializer())

# resume at a later point
ctx = Context.from_dict(serialized_context, serializer=JsonSerializer())
handler = workflow.run(ctx=ctx)
final_result = await handler

[brycecf]

Okay, that works for 0.11.20.