🚨 Websocket connection does not send response to client till all messages are received from generator function

[asifrahaman13]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google "How to X in FastAPI" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse

app = FastAPI()


class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_json(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()

from functions import get_query


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    try:
        while True:
            try:
                data = await websocket.receive_text()
                print(data)
                response = get_query(data)

                generator = iter(response)

                while True:
                    try:
                        response = next(generator)

                        print(
                            f"########################################################## Received data is  {response}"
                        )

                        await manager.send_personal_message(response, websocket)
                    except StopIteration:
                        break
            except WebSocketDisconnect:
                manager.disconnect(websocket)
                break

    except WebSocketDisconnect:
        
        await manager.disconnect(websocket)

Thje query function file

import logging
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

# from whitelist import get_query
from multiprocessing import Pool
from duckduckgo_search import DDGS  # Assuming you're using a DuckDuckGo search library
from openai import OpenAI
import os

OpenAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OpenAI_API_KEY)

whitelist = ["en.wikipedia.org"]


def search_domain(domain, search):
    import time

    start_time = time.time()
    result = DDGS().text(f"site:{domain} {search}", max_results=2)
    print(f"Time taken for {domain} is {time.time() - start_time}")
    print(f"Results for {domain}: {result}")  # Print results for debugging
    return result


def search_result(search):
    with Pool(processes=len(whitelist)) as pool:
        results = pool.starmap(
            search_domain, [(domain, search) for domain in whitelist]
        )
    flattened_results = [result for sublist in results for result in sublist]
    return flattened_results


def get_query(query: str):
    knowledge_base = search_result(query)
    print(f"Knowledge base: {knowledge_base}")  # Print knowledge base for debugging
    import time

    start_time = time.time()

    knowledge_source = [{"body": item["body"]} for item in knowledge_base]

    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": f"You are a poetic assistant. You are a helpful assistant. Your goal is to answer the user's query. The data from which you can frame your answer is: {str(knowledge_source)}. Feel free to add your own answers too.",
            },
            {
                "role": "user",
                "content": query,
            },
        ],
        max_tokens=300,
        temperature=0.7,
        stream=True,
    )

    extracted_data = [
        {"title": item["title"], "href": item["href"], "content": item["body"][:40]} for item in knowledge_base
    ]

    # Create a completion
    # Initialize a buffer to store the sentence
    sentence_buffer = ""

    total_text = ""
    # Iterate over the stream of chunks
    for chunk in completion:

        # Check if the completion is a message
        if chunk.choices[0].delta.content is not None:

            # Append the chunk to the buffer
            sentence_buffer += chunk.choices[0].delta.content

            total_text += chunk.choices[0].delta.content

            # Check if the sentence is complete
            if sentence_buffer.endswith((".", "!", "?")):
                # Record the end time
                end_time = time.time()

                # Calculate the elapsed time
                elapsed_time = end_time - start_time

                logging.info(f"Elapsed time for open ai: {elapsed_time} seconds")

                # Yield the sentence
                yield {"response": sentence_buffer.strip(),"finished": False}
                sentence_buffer = ""
    yield { "sources": extracted_data, "finished": True}

Description

Hello I am trying to stream response from Open AI. I am using the yield method to do so. However even though the response is getting the streaming data it is not sending the data to client immediately. It is waiting for all the generator to complete first and then it is sending them in one go.

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse

app = FastAPI()


class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_json(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()

from functions import get_query


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    try:
        while True:
            try:
                data = await websocket.receive_text()
                print(data)
                response = get_query(data)

                generator = iter(response)

                while True:
                    try:
                        response = next(generator)

                        print(
                            f"########################################################## Received data is  {response}"
                        )

                        await manager.send_personal_message(response, websocket)
                    except StopIteration:
                        break
            except WebSocketDisconnect:
                manager.disconnect(websocket)
                break

    except WebSocketDisconnect:

The lines on which it is blocking is:

                        print(
                            f"########################################################## Received data is  {response}"
                        )

                        await manager.send_personal_message(response, websocket)

Note that the response is being received as stream only. But it is not sending anything to the client unless all the messages are being received from the get_query method.

The get_query function:

def get_query(query: str):
    knowledge_base = search_result(query)
    print(f"Knowledge base: {knowledge_base}")  # Print knowledge base for debugging
    import time

    start_time = time.time()

    knowledge_source = [{"body": item["body"]} for item in knowledge_base]

    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": f"You are a poetic assistant. You are a helpful assistant. Your goal is to answer the user's query. The data from which you can frame your answer is: {str(knowledge_source)}. Feel free to add your own answers too.",
            },
            {
                "role": "user",
                "content": query,
            },
        ],
        max_tokens=300,
        temperature=0.7,
        stream=True,
    )

    extracted_data = [
        {"title": item["title"], "href": item["href"], "content": item["body"][:40]} for item in knowledge_base
    ]

    # Create a completion
    # Initialize a buffer to store the sentence
    sentence_buffer = ""

    total_text = ""
    # Iterate over the stream of chunks
    for chunk in completion:

        # Check if the completion is a message
        if chunk.choices[0].delta.content is not None:

            # Append the chunk to the buffer
            sentence_buffer += chunk.choices[0].delta.content

            total_text += chunk.choices[0].delta.content

            # Check if the sentence is complete
            if sentence_buffer.endswith((".", "!", "?")):
                # Record the end time
                end_time = time.time()

                # Calculate the elapsed time
                elapsed_time = end_time - start_time

                logging.info(f"Elapsed time for open ai: {elapsed_time} seconds")

                # Yield the sentence
                yield {"response": sentence_buffer.strip(),"finished": False}
                sentence_buffer = ""
    yield { "sources": extracted_data, "finished": True}

Interestingly if I use some other websocket library and create a echo server it works perfectly.

import asyncio
import logging
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import websockets

app = FastAPI()


from functions import get_query


async def websocket_endpoint(websocket: WebSocket, client_id: str):

    try:
        async for message in websocket:
            try:
                data = message
                print(data)
                response = get_query(data)

                generator = iter(response)

                while True:
                    try:
                        response = next(generator)

                        print(
                            f"########################################################## Received data is  {response}"
                        )

                        await websocket.send(str(response))
                    except StopIteration:
                        break
            except WebSocketDisconnect:

                break

    except WebSocketDisconnect:
        # print("ok")
        websocket.disconnect(websocket)


async def main():
    try:
        async with websockets.serve(websocket_endpoint, "0.0.0.0", 8000):
            logging.info("Server started")
            await asyncio.Future()  # run forever
    except Exception as e:
        logging.error(f"Server error: {e}")
        # await main()


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

Operating System

Linux

Operating System Details

Linux OS
UBUNTU 24 LTS version

FastAPI Version

0.111.0

Pydantic Version

2.7.3

Python Version

Python 3.12.3

Additional Context

No response