调用模型与流式代理

学习目标

这一课实现 Python 版的 AI 流式代理，协议逐点对齐本站 /api/ai/chat：密钥只留服务端、SSE 转发增量、错误按 category 分类、客户端断开时取消上游、已生成的部分内容不丢。

问题场景

前端要接 AI 助手。第一反应是把密钥塞进前端——那就等于把密钥公开了。本站的解法是 Next.js API 路由当代理：密钥在服务端 .env.local，浏览器只看到 SSE 流。现在用 FastAPI 做等价物，顺便把流式代理最难的三件事（错误、取消、部分响应）一次讲透。

核心解释

SSE（Server-Sent Events）的格式朴素到不像协议：data: 前缀 + JSON + 两个换行。事件序列与本站完全一致：start（会话与消息标识）→ text（增量 delta）→ done（正常收尾），出错时发 error。

import json
import os
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import httpx

app = FastAPI()

CATEGORY_OF_STATUS = {401: "not_configured", 402: "quota", 429: "rate_limit"}

@app.post("/ai/chat")
async def ai_chat(req: Request):
    api_key = os.environ.get("AI_API_KEY", "")
    if not api_key:                                   # 可选能力：服务不崩，分类明确
        return sse_error("not_configured", "AI 服务未配置：缺少服务端密钥")

    upstream = httpx.AsyncClient(timeout=httpx.Timeout(60, connect=10))
    queue: asyncio.Queue = asyncio.Queue()

    async def pump():
        try:
            async with upstream.stream("POST", "https://api.example.com/v1/chat/completions",
                                       headers={"Authorization": f"Bearer {api_key}"},
                                       json=build_body(req)) as resp:
                if resp.status_code != 200:
                    category = CATEGORY_OF_STATUS.get(resp.status_code, "model_error")
                    await queue.put(error_event(category, f"上游返回 {resp.status_code}"))
                    return
                async for line in resp.aiter_lines():     # 逐行读上游 SSE
                    if line.startswith("data: ") and line != "data: [DONE]":
                        delta = json.loads(line[6:])["choices"][0]["delta"].get("content")
                        if delta:
                            await queue.put(sse({"type": "text", "delta": delta}))
            await queue.put(sse({"type": "done"}))
        except httpx.TimeoutException:
            await queue.put(error_event("unreachable", "连接超时"))
        except asyncio.CancelledError:                # 客户端断开 → 取消传播
            await queue.put(error_event("cancelled", "请求已取消"))
            raise
        finally:
            await upstream.aclose()

    return StreamingResponse(consume(queue, pump()), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache"})

四条设计决策，全部能在本站 app/api/ai/chat/route.ts 里找到原型：

1. 密钥只在服务端。浏览器请求你的 /ai/chat，你带上密钥请求上游；密钥永远不会出现在响应里。
2. 错误必须分类。照搬本站的 category 枚举：not_configured（没配密钥）、unreachable（连不上/超时）、quota（欠费）、rate_limit（限流）、model_error（上游 5xx）、cancelled（用户取消）、unknown。前端拿 category 决定提示语和重试策略，比解析自由文本可靠得多。
3. 取消要传播到上游。客户端断开时 FastAPI 取消协程，CancelledError 触发后要主动断开上游请求——不然上游继续烧 token。本站用 req.signal 监听 abort 后 upstreamController.abort()，同一个意图。
4. 部分响应要保留。流到一半失败时，错误事件里带上已收到的 partial 文本（本站还会把非空部分持久化、只删掉完全空的占位消息），前端已经渲染的字不闪没。

示例拆解

用命令行验证事件序列（不依赖前端）：

curl -N -X POST http://127.0.0.1:8100/ai/chat \
  -H "Content-Type: application/json" -d "{\"message\":\"解释一下 SSE\"}"
# data: {"type":"start","conversationId":"conv_1"}
# data: {"type":"text","delta":"SSE 是"}
# data: {"type":"text","delta":"服务器单向推送…"}
# data: {"type":"done"}

三个失败场景各验一遍：拔掉 AI_API_KEY 环境变量启动 → 第一条就是 not_configured 错误事件；把上游地址改成不可达端口 → unreachable；请求发出后 Ctrl+C 断开 curl → 服务端日志出现 cancelled。前端 EventSource（或 fetch + ReadableStream）按 type 字段分支渲染，逻辑与本站 assistant 功能一一对应。

动手任务

1. 在 pump 里累计已发送的文本长度，出错时把已生成内容放进 error 事件的 partial 字段。
2. 加一个「会话占位消息」：收到 start 时在内存里登记 message_id，流正常结束改状态 done，出错改 error——对齐本站 ai_message.status 的三态。
3. 把上游换成一个会中途断开的假服务（返回几个 delta 后直接关连接），验证客户端收到的是错误事件而不是挂死。

常见错误

• 把密钥放进响应或前端 env：NEXT_PUBLIC 前缀的变量会进浏览器包；密钥只读服务端环境变量。
• 忘了 flush/缓冲：代理层攒够一定字节才发，前端「逐字」变成「整段蹦」；StreamingResponse 关闭缓冲，必要时逐行 yield。
• 吞掉 CancelledError：except asyncio.CancelledError: pass 会破坏取消机制；清理完必须 raise。
• 错误只打日志不进流：前端收不到 error 事件就一直转圈等；上游失败也要以 SSE error 事件收尾。

检查与复盘

自查清单：

• 密钥只出现在服务端环境变量，代码里搜不到明文
• 事件序列 start/text/done/error 与本站 /api/ai/chat 一致，错误带 category
• 客户端中途断开时上游请求被取消，已生成文本以 partial 保留

参考资料

• FastAPI 进阶：Custom Response（StreamingResponse） —— 流式响应实现
• MDN：Server-sent events —— SSE 协议与浏览器端 API