FastAPI WebSockets for real-time LLM chat

Bridging asynchronous LLM token streams with a browser connection is the core of any fastapi websockets llm chat service. This tutorial builds a minimal but production-shaped backend that streams model output token-by-token over a WebSocket, plus a tiny frontend to consume it.

Prerequisites

Python 3.11+ (asyncio improvements help, but 3.10 works)
pip install fastapi uvicorn websockets openai python-dotenv
An OpenAI-compatible endpoint. Export LLM_API_KEY. Optionally set LLM_BASE_URL (default is OpenAI) and LLM_MODEL (default gpt-4o-mini).

If you run a local venv:

python -m venv .venv
source .venv/bin/activate
pip install fastapi uvicorn websockets openai python-dotenv

Server skeleton

We start with a single main.py. The key piece is a WebSocket route that receives a JSON message, appends it to a per-connection conversation list, calls the LLM with stream=True, and forwards each delta to the client.

import os
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from openai import AsyncOpenAI

app = FastAPI()
client = AsyncOpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.openai.com/v1"),
)
MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")

INDEX_HTML = """
<!doctype html><html><body>
<input id="msg" placeholder="Type a message"/><button onclick="send()">Send</button>
<pre id="out"></pre>
<script>
const ws = new WebSocket(`ws://${location.host}/ws/chat`);
const out = document.getElementById("out");
ws.onmessage = (e) => {
  const d = JSON.parse(e.data);
  if (d.type === "token") out.textContent += d.content;
  if (d.type === "error") out.textContent += "\\nERROR: " + d.message;
};
function send() {
  const v = document.getElementById("msg").value;
  out.textContent += "\\nYou: " + v + "\\nAssistant: ";
  ws.send(JSON.stringify({role: "user", content: v}));
  document.getElementById("msg").value = "";
}
</script></body></html>
"""

@app.get("/")
async def index():
    return HTMLResponse(INDEX_HTML)

@app.websocket("/ws/chat")
async def chat_ws(websocket: WebSocket):
    await websocket.accept()
    messages = []
    try:
        while True:
            data = await websocket.receive_json()
            messages.append({"role": data.get("role", "user"), "content": data["content"]})
            stream = await client.chat.completions.create(
                model=MODEL,
                messages=messages,
                stream=True,
            )
            collected = []
            async for chunk in stream:
                delta = chunk.choices[0].delta.content or ""
                if delta:
                    collected.append(delta)
                    await websocket.send_json({"type": "token", "content": delta})
            messages.append({"role": "assistant", "content": "".join(collected)})
    except Exception as e:
        await websocket.send_json({"type": "error", "message": str(e)})

This is a complete fastapi websockets llm chat loop. Each connection keeps its own messages list, so conversations are isolated. The assistant’s full reply is stored after streaming so the next turn has context.

Streaming mechanics

The OpenAI Python client returns an async iterator when stream=True. Each chunk conforms to the SSE protocol but is already parsed into objects. We extract delta.content and immediately push it over the socket. Doing this per-token keeps latency low; the browser paints as soon as the packet arrives.

One subtlety: chunk.choices[0].delta.content can be None on the final chunk (which carries finish_reason). The or "" guard avoids sending nulls.

Running it

With the environment set:

export LLM_API_KEY=sk-...
uvicorn main:app --port 8000 --reload

Open http://localhost:8000. Type Explain asyncio in one sentence and click Send. Expected browser output:

You: Explain asyncio in one sentence
Assistant: asyncio lets you write concurrent code using coroutines that yield control instead of blocking threads.

Tokens appear incrementally if your network and model latency allow. On the server side, uvicorn logs the WebSocket handshake and then stays quiet—no per-token logging unless you add it.

Handling client disconnects

If the user closes the tab mid-stream, websocket.send_json raises WebSocketDisconnect. The except block catches it and ends the coroutine. You should not attempt to send the error frame after a disconnect; wrap the final send in a try or just let it bubble. For cleaner behavior:

from fastapi import WebSocketDisconnect

@app.websocket("/ws/chat")
async def chat_ws(websocket: WebSocket):
    await websocket.accept()
    messages = []
    try:
        while True:
            data = await websocket.receive_json()
            messages.append({"role": "user", "content": data["content"]})
            stream = await client.chat.completions.create(model=MODEL, messages=messages, stream=True)
            collected = []
            async for chunk in stream:
                delta = chunk.choices[0].delta.content or ""
                if delta:
                    collected.append(delta)
                    await websocket.send_json({"type": "token", "content": delta})
            messages.append({"role": "assistant", "content": "".join(collected)})
    except WebSocketDisconnect:
        return
    except Exception as e:
        try:
            await websocket.send_json({"type": "error", "message": str(e)})
        except WebSocketDisconnect:
            pass

Backpressure and concurrency

FastAPI runs each WebSocket in its own task. If the client is slow, send_json awaits until the OS socket buffer accepts the data. This naturally applies backpressure to the LLM stream consumption: you won’t accumulate unbounded memory from a stalled consumer. If you expect many simultaneous sessions, set uvicorn --workers 4 or put it behind a load balancer that supports WebSocket affinity.

Do not buffer the entire assistant message in memory and send it at once; that defeats the purpose of fastapi websockets llm chat real-time feel. The pattern above streams directly.

Using a model gateway

Hardcoding LLM_BASE_URL to one provider limits you to its model catalog and outage profile. Pointing the AsyncOpenAI client at an OpenAI-compatible gateway like n4n.ai gives one endpoint that addresses 240+ models with automatic fallback when a provider is rate-limited or degraded. Your WebSocket handler stays identical; only the environment variables change:

export LLM_BASE_URL=https://api.n4n.ai/v1
export LLM_API_KEY=your-gateway-key
export LLM_MODEL=anthropic/claude-3.5-sonnet

The same streaming code works because the gateway honors the OpenAI chat completions contract and forwards provider cache-control hints.

Production considerations

Auth: Put an auth check before websocket.accept() using a token in query params or headers.
Rate limits: The LLM call can fail with 429. Catch it and send a structured error so the UI can retry.
Timeouts: Wrap client.chat.completions.create in asyncio.wait_for with a 60s cap to avoid hung connections.
History truncation: Cap messages length to avoid token-limit errors; drop oldest turns.

A fastapi websockets llm chat endpoint is deceptively small, but the operational edges—disconnects, backpressure, model routing—decide whether it survives real traffic.

Final check

You now have a runnable server that:

Accepts a browser WebSocket.
Maintains per-session conversation history.
Streams LLM tokens with low latency.
Degrades gracefully on disconnect.

Swap the base URL to any OpenAI-compatible service and the same code serves as the real-time layer for your product.