Streaming response over FastAPI

[ChiIIBiII]

Hi, I'm wondering what the best way is to stream LLM responses over a custom fastapi server. I'm running into a bunch of asyncio related issues. My current code is the following, but it doesn't work:

def done_callback(x):
    try:
        print(x.result)
    except Exception as e:
        print(f"Bad code {e}")


@app.post("/chat")
async def chat(request: ChatRequest):
    assistant = get_assistant("default")

    def Mygen() -> Generator[Any, str, None]:
        try:
            chunk = yield
            yield chunk
            chunk = yield
            while chunk is not None and chunk != "":
                yield chunk
                chunk = yield
        except asyncio.CancelledError:
            print("Streaming was cancelled")
            return
        return

    generator = Mygen()
    next(generator)

    def assistant_callback(chunk: StreamingChunk) -> None:
        if (
            chunk.meta.get("tool_calls") is None
            and chunk.meta.get("function_calls") is None
            and chunk.meta.get("finish_reason") != "tool_calls"
        ):
            generator.send(chunk.content)

    assistant_task = asyncio.create_task(
        assistant.run(
            request.query,
            request.context,
            streaming_callback=assistant_callback,
        )
    )
    assistant_task.add_done_callback(done_callback)

    return StreamingResponse(generator)

where assistant is my haystack pipeline. I'm honestly kind of surprised that I can't find sample code to do that online. I also looked at the hayhooks repo but it doesn't seem they implemented streaming. Any help greatly appreciated!

[ggdupont]

I'm stuck on the exact same problem (well given that I have a Pipeline instead of a raw Generator)

[ChiIIBiII]

So as far as I can tell the problem is, that you can't return something from an async function and have that returned object in turn get data from a running function, if that function is not declared async down to the callback. That means I'll have to make the whole pipeline run function (and the function it calls down to the component run functions) async. At that point you probably might as well just rewrite your whole pipeline as a generator.
I have now settled for the ugly workaround of starting a new thread, and waiting for that thread to call the callback inside the generator that I'm returning.

@app.get("/chat")
async def chat(request: ChatRequest):
    q = queue.Queue()

    def assistant_callback(chunk: StreamingChunk) -> None:
        if (
            chunk.meta.get("tool_calls") is None
            and chunk.meta.get("function_calls") is None
            and chunk.meta.get("finish_reason") != "tool_calls"
        ):
            q.put(chunk)

    assistant = get_assistant("default")
    t = threading.Thread(
        target=assistant.run,
        args=[request.query, request.context, assistant_callback],
        daemon=True,
    )

    def MyGen(thread: threading.Thread) -> Generator[str, Any, None]:
        chunk = StreamingChunk("")
        while (
            "finish_reason" not in chunk.meta or chunk.meta["finish_reason"] != "stop"
        ):
            chunk = q.get()
            yield chunk.content
            q.task_done()
        thread.join()
        return

    generator = MyGen(t)
    t.start()

    return StreamingResponse(generator)

I'm still kind of surprised that this use case is neither documented nor really expected apparently.

[ggdupont]

There a discord discussion where they propse to use the callback to write directly to the response stream to another API by posting on an http request object:
https://discord.com/channels/993534733298450452/1285319730948997182/1285402858673405962
Maybe there is an idea here to have the callback writing in the StreamingResponse?

[PaulEibensteiner]

Hey, thanks for the suggestion, I don't have access to the discord channel but I doubt that it solves my problem from your description. I know I could post to some api, I just can't yield from a generator that I returned from the same function that also runs the callbacks using the async library without making all my code async.