Skip to content

38. FastAPI WebSocket,流式响应与完整示例 ​

两种持续通信,各做一个完整案例 ​

WebSocket 建立长连接,客户端和服务器可随时互发消息,适合聊天室。HTTP 流式响应由服务器逐段向客户端发送内容,适合模型逐字生成。SSE 在文本流上规定事件格式,通常单向传输。

本篇拆成两个独立练习目录:chat_demo 和 llm_demo,各放自己的 main.py,不要把两个 app 无意合并到同一文件。均先安装 uv add "fastapi[standard]" httpx。

案例一:按房间广播的聊天室 ​

消息协议使用 JSON:type 表示消息类型,nickname 表示显示昵称,text 表示消息文本。本例昵称只作显示,不代表经过认证的身份。

服务端 main.py ​

python
import asyncio
from pathlib import Path
from typing import Annotated
from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse

app = FastAPI()
ALLOWED_ORIGINS = {"http://127.0.0.1:8000", "http://localhost:8000"}

class ChatHub:
    def __init__(self):
        self.rooms: dict[str, set[WebSocket]] = {}
        self.lock = asyncio.Lock()

    def remove(self, room: str, socket: WebSocket):
        peers = self.rooms.get(room)
        if peers is not None:
            peers.discard(socket)
            if not peers:
                self.rooms.pop(room, None)

    async def broadcast(self, room: str, message: dict):
        # 教学版串行发送,防止多个广播同时写同一连接。
        async with self.lock:
            for socket in tuple(self.rooms.get(room, ())):
                try:
                    await asyncio.wait_for(socket.send_json(message), timeout=2)
                except (RuntimeError, OSError, WebSocketDisconnect, asyncio.TimeoutError):
                    self.remove(room, socket)
                    try:
                        await socket.close(code=1013)
                    except (RuntimeError, OSError):
                        pass

hub = ChatHub()

@app.get("/", response_class=HTMLResponse)
def page():
    return Path(__file__).with_name("chat.html").read_text(encoding="utf-8")

@app.websocket("/ws/{room}")
async def chat(
    websocket: WebSocket,
    room: str,
    nickname: Annotated[str, Query(min_length=1, max_length=20)] = "访客",
):
    if websocket.headers.get("origin") not in ALLOWED_ORIGINS or len(room) > 50:
        await websocket.close(code=1008)
        return
    await websocket.accept()
    hub.rooms.setdefault(room, set()).add(websocket)
    try:
        await hub.broadcast(room, {"type": "system", "text": f"{nickname} 进入房间"})
        while True:
            text = await websocket.receive_text()
            if len(text) > 2000:
                await websocket.close(code=1009)
                break
            if text.strip():
                await hub.broadcast(room, {"type": "message",
                                          "nickname": nickname, "text": text})
    except WebSocketDisconnect:
        pass
    finally:
        hub.remove(room, websocket)
        await hub.broadcast(room, {"type": "system", "text": f"{nickname} 离开房间"})

客户端 chat.html ​

html
<!doctype html>
<meta charset="utf-8" />
<title>FastAPI 聊天室</title>
<input id="room" value="study" aria-label="房间" />
<input id="nickname" value="小明" aria-label="昵称" />
<button id="connect">连接</button>
<p id="state">未连接</p>
<pre id="log"></pre>
<input id="text" placeholder="输入消息" maxlength="2000" />
<button id="send">发送</button>
<script>
  let socket;
  const el = (id) => document.getElementById(id);
  el("connect").onclick = () => {
    if (socket) socket.close();
    const protocol = location.protocol === "https:" ? "wss" : "ws";
    const room = encodeURIComponent(el("room").value || "study");
    const name = encodeURIComponent(el("nickname").value || "访客");
    socket = new WebSocket(
      `${protocol}://${location.host}/ws/${room}?nickname=${name}`,
    );
    socket.onopen = () => (el("state").textContent = "已连接");
    socket.onclose = () => (el("state").textContent = "连接已关闭");
    socket.onerror = () => (el("state").textContent = "连接出错");
    socket.onmessage = (event) => {
      const data = JSON.parse(event.data);
      el("log").textContent += `${data.nickname || "系统"}:${data.text}\n`;
    };
  };
  el("send").onclick = () => {
    if (socket && socket.readyState === WebSocket.OPEN) {
      socket.send(el("text").value);
      el("text").value = "";
    }
  };
</script>

执行 uv run uvicorn main:app --host 127.0.0.1 --port 8000。打开两个浏览器标签访问该地址,输入不同昵称、同一房间,点击连接后互发消息;改用不同房间应互不收到消息。通过 FastAPI 地址加载页面,不是直接双击 HTML 文件,否则 Origin 不在允许列表。

从教学版走向真实聊天室 ​

本例房间保存在单进程内存,多 worker 间不共享。扩展时用 Redis Pub/Sub 等消息分发,历史记录另做持久化;Pub/Sub 本身不等于可靠历史存储。

串行广播便于理解,但慢连接会拖延广播。高并发时为每个连接建立有上限的发送队列和独立发送协程,定义队列满时断开或丢弃策略。加入连接数限制、心跳、限流、重连退避和消息 ID。

WebSocket 不自动受 HTTP CORS 中间件保护,应验证 Origin、握手身份和房间权限。浏览器 WebSocket API 不能像 fetch 那样任意设置 Authorization 头,可结合安全 Cookie 或短期一次性握手凭证设计;避免把长期令牌放查询串而进入日志。

案例二:接入大模型并转发 SSE ​

接口契约约定:上游支持 POST {LLM_BASE_URL}/chat/completions、Bearer Key、stream=true;返回 SSE,其中文本在 choices[0].delta.content,结束标识为 data: [DONE]。这是一种常见兼容协议,不是所有厂商都遵循。厂商协议不同时只修改上游请求和解析适配层。

Key 留在服务端环境变量,浏览器不接触。未配置时接口返回 503,不会伪造模型回答,也不要求你现在提供 Key。

llm_demo/main.py ​

python
import json
import os
from contextlib import asynccontextmanager
from pathlib import Path
import httpx
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from pydantic import BaseModel, Field

@asynccontextmanager
async def lifespan(app: FastAPI):
    async with httpx.AsyncClient(
        timeout=httpx.Timeout(connect=10, read=60, write=30, pool=10),
        limits=httpx.Limits(max_connections=50, max_keepalive_connections=20),
    ) as client:
        app.state.http = client
        yield

app = FastAPI(lifespan=lifespan)

@app.get("/", response_class=HTMLResponse)
def index():
    return Path(__file__).with_name("llm.html").read_text(encoding="utf-8")

class PromptIn(BaseModel):
    prompt: str = Field(min_length=1, max_length=8000)

def event(name: str, payload: dict):
    return f"event: {name}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"

async def read_sse_data(lines):
    data = []
    async for line in lines:
        if line == "":
            if data:
                yield "\n".join(data)
                data.clear()
        elif line.startswith("data:"):
            value = line[5:]
            data.append(value[1:] if value.startswith(" ") else value)
        # 注释、event、id 等字段在此适配层无需使用。
    if data:
        yield "\n".join(data)

@app.post("/chat/stream")
async def chat_stream(data: PromptIn, request: Request):
    key = os.getenv("LLM_API_KEY")
    base = os.getenv("LLM_BASE_URL", "").rstrip("/")
    model = os.getenv("LLM_MODEL")
    if not key or not base or not model:
        raise HTTPException(503, "请先配置模型地址、名称和 API Key")
    client = request.app.state.http
    upstream_request = client.build_request(
        "POST", f"{base}/chat/completions",
        headers={"Authorization": f"Bearer {key}"},
        json={"model": model, "stream": True,
              "messages": [{"role": "user", "content": data.prompt}]},
    )
    # 在返回流之前检查上游状态,此时仍可返回普通 HTTP 错误。
    try:
        upstream = await client.send(upstream_request, stream=True)
    except httpx.TimeoutException:
        raise HTTPException(504, "模型服务连接超时")
    except httpx.HTTPError:
        raise HTTPException(502, "无法连接模型服务")
    if upstream.is_error:
        await upstream.aclose()
        raise HTTPException(502, "模型服务拒绝了请求,请检查服务端配置和额度")

    async def generate():
        try:
            async for raw in read_sse_data(upstream.aiter_lines()):
                if raw == "[DONE]":
                    yield event("done", {})
                    return
                payload = json.loads(raw)
                if payload.get("error"):
                    raise ValueError("上游流中报告错误")
                choices = payload.get("choices") or []
                if choices:
                    content = (choices[0].get("delta") or {}).get("content")
                    if isinstance(content, str) and content:
                        yield event("token", {"text": content})
            raise ValueError("上游连接结束但缺少完成标记")
        except (httpx.HTTPError, ValueError, TypeError, AttributeError, KeyError, IndexError):
            yield event("error", {"message": "模型输出中断,请稍后重试"})
        finally:
            await upstream.aclose()

    return StreamingResponse(
        generate(), media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )

HTTPX 按行迭代后,再按 SSE 的空行拼成事件,避免把网络分块误当一条完整 JSON。网络块、UTF-8 字符、JSON 文档和 SSE 事件不是一回事。

配置、启动与调用 ​

在当前终端设置以下变量,实际值由你以后提供:

powershell
$env:LLM_BASE_URL = 'https://你的模型网关地址/v1'
$env:LLM_MODEL = '你的模型名称'
$env:LLM_API_KEY = '你的模型Key'
uv run uvicorn main:app --port 8001

新终端执行:

powershell
curl.exe -N http://127.0.0.1:8001/chat/stream -H "Content-Type: application/json" --data-raw '{"prompt":"用通俗语言解释依赖注入"}'

部分旧版 PowerShell 处理原生命令引号不同,遇到正文编码问题可把 JSON 保存为 UTF-8 的 prompt.json,再使用 --data-binary "@prompt.json"。

响应协议如下,text 的颗粒大小由上游决定,不保证逐字:

text
event: token
data: {"text":"依赖注入"}

event: token
data: {"text":"可以理解为……"}

event: done
data: {}

浏览器客户端 llm.html ​

将以下文件放在模型示例的 main.py 旁边,启动后访问 http://127.0.0.1:8001/。页面不需要也不接收模型 Key。它只调用自己的服务端接口。

html
<!doctype html>
<meta charset="utf-8" />
<title>学习助手</title>
<textarea id="prompt" rows="4" cols="60">用通俗语言解释依赖注入</textarea>
<p><button id="run">开始</button><button id="stop">取消</button></p>
<p id="state">等待提问</p>
<pre id="answer" style="white-space:pre-wrap"></pre>
<script>
  const el = (id) => document.getElementById(id);
  let controller = null;
  el("stop").onclick = () => {
    if (controller) controller.abort();
  };
  el("run").onclick = async () => {
    controller = new AbortController();
    el("run").disabled = true;
    el("answer").textContent = "";
    el("state").textContent = "正在生成";
    let reader;
    try {
      const response = await fetch("/chat/stream", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ prompt: el("prompt").value }),
        signal: controller.signal,
      });
      if (!response.ok) {
        const error = await response.json();
        throw new Error(
          typeof error.detail === "string" ? error.detail : "请求参数不合法",
        );
      }
      reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = "",
        completed = false;
      while (!completed) {
        const { value, done } = await reader.read();
        buffer += decoder.decode(value, { stream: !done });
        let boundary;
        while ((boundary = buffer.indexOf("\n\n")) >= 0) {
          const frame = buffer.slice(0, boundary);
          buffer = buffer.slice(boundary + 2);
          const lines = frame.split("\n");
          const type = (
            lines.find((x) => x.startsWith("event:")) || "event: message"
          )
            .slice(6)
            .trim();
          const raw = lines
            .filter((x) => x.startsWith("data:"))
            .map((x) => x.slice(5).trimStart())
            .join("\n");
          if (!raw) continue;
          const data = JSON.parse(raw);
          if (type === "token") el("answer").textContent += data.text;
          if (type === "error") throw new Error(data.message);
          if (type === "done") completed = true;
        }
        if (done && !completed) throw new Error("连接提前结束,回答未完成");
      }
      el("state").textContent = "完成";
    } catch (error) {
      el("state").textContent =
        error.name === "AbortError" ? "已取消" : error.message;
    } finally {
      if (reader) await reader.cancel().catch(() => {});
      controller = null;
      el("run").disabled = false;
    }
  };
</script>

本客户端针对本地服务端统一生成的 LF 换行 SSE;若直接连接任意 SSE 服务,应使用支持 CRLF、id、retry 等完整协议的解析器。

浏览器消费流与错误边界 ​

原生 EventSource 适合 GET,不能直接用它 POST 本例 JSON。浏览器应使用 fetch 发 POST,先检查 HTTP 状态,再通过 response.body.getReader() 读取,使用 TextDecoder 的流模式保留跨块字符,并缓存到 \n\n 后解析事件。也可以选用支持 fetch 的成熟 SSE 解析器。

响应开始前的配置、认证和连接错误可用 HTTP 状态表达;开始后的错误通过 event: error 发送。客户端收到 done 才视为完成,流异常关闭应视为未完成。取消按钮使用 AbortController;服务端生成器的 finally 关闭上游连接。

实际接入还需补齐什么 ​

上线时给本地接口增加第 36 篇认证、用户级并发和频率限制、输入输出上限、费用记录和超时策略;Key 只在服务器保存,日志不记录完整 Authorization。不要让浏览器传任意上游地址,以免将接口变成任意网络代理。

代理需关闭相关流缓冲并设置合理空闲超时;长时间无输出时可设计 SSE 心跳。工具调用、图片、推理字段、计费用量和其他结束协议需要按厂商格式扩展,不能直接当成 delta.content。

练习:聊天室验证同房、异房、主动断开和超长消息;模型接口先验证无配置返回 503,再用本地模拟 SSE 验证 token/done/error。真实 Key 配置后,另测实际模型与网络,这一步不能由离线模拟替代。