Stream response with AsyncGenerator

[LorenzoVaccher01]

I would need to stream data to the user via an endpoint through a function that returns an AsyncGenerator

The idea is as follows:

message_route = Blueprint('message_route', __name__, template_folder='templates')

@message_route.route('/message/<string:key>/<string:chat_uuid>', methods=['POST'])
async def message(key, chat_uuid):
    
   # [...] Other endpoint logic not helpful in resolving the issue

    async def event_stream():
        agent_handler = agent.run(input=content['message'])

        # stream_events() return an AsyncGenerator[Event, None], so async is necessary
        async for ev in agent_handler.stream_events():
            if isinstance(ev, MessageOutputEvent):
                print(ev.message)
                yield ev.message

       # wait for the process to finish before continuing
        response = await agent_handler
        print(response)

    resp = make_response(event_stream())
    resp.mimetype = "text/event-stream"
    resp.set_cookie(key='user_uuid', samesite="none", value="user_123")
    return resp, 200

To be able to handle requests with async/await I installed flask[async] (3.0.3) but when I add async to the message() function I get the error:

TypeError: The view function did not return a valid response. The return type must be a string, dict, list, tuple with headers or status, Response instance, or WSGI callable, but it was a coroutine.

[CNFeffery]

@LorenzoVaccher01 Hello, is there any progress regarding this issue? I encountered a similar problem when trying to integrate the asynchronous streaming agent execution in pydantic-ai through the asynchronous version of flask:

Debugging middleware caught exception in streamed response at a point where response headers were already sent.
Traceback (most recent call last):
  File "C:\Users\pengz\AppData\Roaming\Python\Python310\site-packages\werkzeug\wsgi.py", line 256, in __next__
    return self._next()
  File "C:\Users\pengz\AppData\Roaming\Python\Python310\site-packages\werkzeug\wrappers\response.py", line 32, in _iter_encoded
    for item in iterable:
TypeError: 'async_generator' object is not iterable

[LorenzoVaccher01]

Hi @CNFeffery,
Unfortunately I didn't find a solution and I migrated to FastAPI for my project 🥲

[CNFeffery]

Hi @CNFeffery, Unfortunately I didn't find a solution and I migrated to FastAPI for my project 🥲

It seems that for this part of the service, I also need to find a way to migrate to FastAPI 😰

[davidism]

Quart is a drop-in ASGI replacement for Flask, maintained by Pallets.

[CNFeffery]

Quart is a drop-in ASGI replacement for Flask, maintained by Pallets.

But currently, I want to directly integrate the sse service for the asynchronous streaming agent into the Dash application. The Dash application has a built-in Flask instance.

[davidism]

You can also use Gevent to serve your Flask app. Gevent has been supported since the beginning, and has equivalent concurrent connection support to ASGI, so you can have as many SSE connections as you like without running out of workers. https://flask.palletsprojects.com/en/stable/deploying/gevent/

[avi2413]

I ran into a similar issue a while back. I was trying to stream from an async generator via flask.Response, but Flask still expect a sync iterable, so you can’t just pass them an async_generator. As a workaround, I wrapped my async generator in a sync generator and drove it with an event loop....it works, but because Flask is still running under WSGI, the worker is blocked while that loop runs, so you don’t really get the benefits of async. Use this only if you have to!

import asyncio
from typing import AsyncGenerator, Callable

def stream_helper(async_gen_factory: Callable[[], AsyncGenerator[str, None]]):
    def _stream():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        agen = async_gen_factory()
        try:
            while True:
                try:
                    chunk = loop.run_until_complete(agen.__anext__())
                except StopAsyncIteration:
                    break
                yield chunk
        finally:
            loop.run_until_complete(agen.aclose())
            loop.close()

    return _stream()

Drop in for your case

from Flask import Response

@message_route.route('/message/<string:key>/<string:chat_uuid>', methods=['POST'])
async def message(key, chat_uuid):
    # [...]
    
    return Response(
        stream_helper(event_stream()),
        content_type="text/event-stream",
    )