38. FastAPI WebSocket,流式响应与完整示例
两种持续通信,各做一个完整案例
WebSocket 建立长连接,客户端和服务器可随时互发消息,适合聊天室。HTTP 流式响应由服务器逐段向客户端发送内容,适合模型逐字生成。SSE 在文本流上规定事件格式,通常单向传输。
本篇拆成两个独立练习目录:chat_demo 和 llm_demo,各放自己的 main.py,不要把两个 app 无意合并到同一文件。均先安装 uv add "fastapi[standard]" httpx。
案例一:按房间广播的聊天室
消息协议使用 JSON:type 表示消息类型,nickname 表示显示昵称,text 表示消息文本。本例昵称只作显示,不代表经过认证的身份。
服务端 main.py
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
<!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
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 事件不是一回事。
配置、启动与调用
在当前终端设置以下变量,实际值由你以后提供:
$env:LLM_BASE_URL = 'https://你的模型网关地址/v1'
$env:LLM_MODEL = '你的模型名称'
$env:LLM_API_KEY = '你的模型Key'
uv run uvicorn main:app --port 8001新终端执行:
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 的颗粒大小由上游决定,不保证逐字:
event: token
data: {"text":"依赖注入"}
event: token
data: {"text":"可以理解为……"}
event: done
data: {}浏览器客户端 llm.html
将以下文件放在模型示例的 main.py 旁边,启动后访问 http://127.0.0.1:8001/。页面不需要也不接收模型 Key。它只调用自己的服务端接口。
<!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 配置后,另测实际模型与网络,这一步不能由离线模拟替代。