AI API 服务开发

subtitle: FastAPI + LLM 流式响应 SSE / WebSocket、请求队列与限流、Token 用量计量

AI API 服务概述

与传统 API 的差异

AI 推理服务与传统 HTTP API 在交互模式上存在本质区别，主要体现在以下几个方面：

维度	传统 API	AI 推理 API
响应时间	毫秒级	秒到分钟级
输出方式	一次性返回	流式逐 token 返回
请求密度	稳定可控	突发性强，显存受限
计费粒度	按请求次数	按 token 数量
状态管理	无状态为主	需管理推理上下文

架构总览

一个典型的 AI API 服务架构包含以下层次：

1. 网关层：负责认证、限流、请求路由。常用组件有 Nginx、Kong、Cloudflare 等。
2. 调度层：管理请求队列，协调 GPU 资源分配，实现优先级调度与动态批处理。
3. 推理层：运行 LLM 模型推理，常见引擎包括 vLLM、TensorRT-LLM、TGI 等。
4. 计量层：统计每次请求的 token 用量，输出计费所需的数据。

各层通过异步消息或 HTTP 回调解耦，整体架构强调高吞吐、低延迟和资源利用率的最大化。

FastAPI + LLM 流式响应

SSE（Server-Sent Events）协议

SSE 是一种基于 HTTP 的单向推送协议，客户端通过 EventSource 接口监听服务端推送的事件流。

SSE 数据格式遵循以下约定：

• 每个事件以 data: 开头，后接 JSON 数据
• 事件之间以空行分隔
• 可选字段：event: 指定事件类型，id: 指定事件 ID，retry: 指定重连间隔

流式接口有两种推送方式：

推送方式	特点	适用场景
一次推送	生成完毕后一次性返回完整结果	非流式接口，兼容传统客户端
逐 token 推送	每生成一个 token 立即推送	流式聊天，提升用户感知速度

FastAPI StreamingResponse

FastAPI 提供了 StreamingResponse 支持流式输出，配合 Python 的 async generator 可以实现高效的逐 token 推送。

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import json
from typing import AsyncGenerator

app = FastAPI()

async def token_generator(prompt: str) -> AsyncGenerator[str, None]:
    """模拟 LLM 逐 token 生成的异步生成器。"""
    # 实际场景中此处替换为模型推理调用
    result = f"这是对 '{prompt}' 的回复内容"
    for char in result:
        yield f"data: {json.dumps({'token': char, 'finished': False})}\n\n"
        await asyncio.sleep(0.05)  # 模拟推理延迟
    # 发送结束标记
    yield f"data: {json.dumps({'token': '', 'finished': True})}\n\n"

@app.post("/v1/chat/completions")
async def chat_completion(prompt: str):
    return StreamingResponse(
        token_generator(prompt),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )

关键点说明：

• media_type="text/event-stream" 设置正确的 SSE MIME 类型
• X-Accel-Buffering: no 禁用 Nginx 缓冲，确保流式数据即时转发
• 生成器必须为异步（async def），避免阻塞事件循环

WebSocket 方案

当需要双向通信时（如客户端发送取消信号），WebSocket 是比 SSE 更合适的选择。

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_json()
            prompt = data["prompt"]
            # 逐 token 发送
            async for token in token_generator(prompt):
                await websocket.send_json({"token": token})
                # 检查客户端取消信号
                try:
                    control = await asyncio.wait_for(
                        websocket.receive_text(), timeout=0.01
                    )
                    if control == "cancel":
                        await websocket.send_json({"event": "cancelled"})
                        break
                except asyncio.TimeoutError:
                    pass
    except WebSocketDisconnect:
        # 清理推理任务
        pass

WebSocket 的 cancel 机制允许客户端在推理过程中主动中止请求，释放 GPU 资源，避免无效计算。

请求队列与限流

为什么需要请求队列

LLM 推理服务面临以下挑战：

• GPU 显存限制：单张 GPU 能同时承载的推理实例数有限（通常为 1-8 个），超出后会导致 OOM。
• 推理延迟积累：当请求到达速率超过处理速率时，未处理的请求会持续堆积，导致所有请求的响应时间恶化。
• 价格波动：无限制的并发会导致 GPU 利用率波动剧烈，难以预测成本。

引入请求队列可以将突发的请求削峰填谷，让后端在稳定的负载下运行。

队列策略

策略	描述	适用场景
FIFO（先进先出）	按到达顺序处理	通用场景，公平性最佳
优先级队列	高优先级请求插队	付费用户优先、实时交互
动态批处理	合并多个请求为一次推理	短请求密集场景，吞吐量最高

动态批处理是 vLLM 等引擎的核心优化手段，它将等待队列中的多个请求合并为一个 batch 送入 GPU 推理，能显著提升吞吐量。但批处理会增加首个 token 的等待时间，需要在延迟和吞吐之间权衡。

限流策略

限流用于防止单个客户端或 IP 过度占用服务资源。

令牌桶（Token Bucket）

令牌桶算法以固定速率向桶中放入令牌，每次请求消耗一个令牌。桶的容量限制了突发流量的大小。

import time
import asyncio
from functools import wraps

class TokenBucket:
    """令牌桶限流器"""

    def __init__(self, rate: float, capacity: int):
        self.rate = rate          # 令牌产生速率（个/秒）
        self.capacity = capacity  # 桶容量
        self.tokens = capacity    # 当前令牌数
        self.last_refill = time.monotonic()

    def _refill(self):
        now = time.monotonic()
        elapsed = now - self.last_refill
        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
        self.last_refill = now

    def consume(self, count: int = 1) -> bool:
        self._refill()
        if self.tokens >= count:
            self.tokens -= count
            return True
        return False


def rate_limit(bucket: TokenBucket):
    """限流装饰器"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            if not bucket.consume():
                from fastapi import HTTPException
                raise HTTPException(status_code=429, detail="请求过于频繁，请稍后重试")
            return await func(*args, **kwargs)
        return wrapper
    return decorator


# 使用示例：每秒 10 个请求，突发上限 20
global_bucket = TokenBucket(rate=10, capacity=20)

@app.post("/v1/completions")
@rate_limit(global_bucket)
async def completions(prompt: str):
    ...

滑动窗口

滑动窗口算法维护一个固定时间窗口内的请求计数，窗口随时间滑动。相比令牌桶，滑动窗口的限流边界更精确，适合对突发有严格限制的场景。

算法	优点	缺点
令牌桶	允许突发，实现简单	突发边界不够精确
滑动窗口	边界精确，控制严格	需要维护时间戳列表，内存开销略高
漏桶	输出速率恒定	不能应对突发，可能导致请求大量丢弃

Token 用量计量

统计维度

每次推理请求需要统计以下指标：

• input_tokens：用户输入经过 tokenizer 后得到的 token 数量
• output_tokens：模型生成结果中的 token 数量
• total_tokens = input_tokens + output_tokens

from pydantic import BaseModel

class TokenUsage(BaseModel):
    input_tokens: int
    output_tokens: int
    total_tokens: int

class CompletionResponse(BaseModel):
    id: str
    object: str = "text_completion"
    choices: list
    usage: TokenUsage

统计 token 用量的关键在于使用与模型相同的 tokenizer（如 HuggingFace 的 AutoTokenizer）进行计数，确保计数与实际推理时的切分方式一致。

计费模型

大多数 LLM API 服务采用按 token 计费的模式，且 input 和 output 定价不同：

模型	输入价格（每千 token）	输出价格（每千 token）
GPT-4o	$0.005	$0.015
Claude 3.5 Sonnet	$0.003	$0.015
自部署 LLaMA 3 70B	按 GPU 成本折算	同上

计费系统通常在计量层记录原始用量数据，然后异步聚合到账单系统，避免在推理路径中引入额外延迟。

统计并返回 token 用量

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

@app.post("/v1/chat/completions")
async def chat_completion(messages: list):
    # 计算输入 token
    input_text = " ".join(m["content"] for m in messages)
    input_tokens = len(tokenizer.encode(input_text))

    # 流式生成并累计输出 token
    output_tokens = 0
    async def generate():
        nonlocal output_tokens
        # 实际调用模型生成
        full_output = ""
        async for chunk in model.generate_stream(messages):
            output_tokens += 1
            full_output += chunk
            yield f"data: {json.dumps({'content': chunk})}\n\n"
        # 最后一条消息包含用量统计
        usage = TokenUsage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=input_tokens + output_tokens,
        )
        yield f"data: {json.dumps({'usage': usage.dict(), 'finish_reason': 'stop'})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

完整示例：FastAPI + vLLM 后端 + SSE 流式接口

以下代码整合了上述所有概念，实现一个可直接运行的 AI API 服务：

import asyncio
import json
import time
from functools import wraps
from typing import AsyncGenerator

from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

app = FastAPI(title="AI API Service")

# ---------------------------------------------------------------------------
# 限流器
# ---------------------------------------------------------------------------

class TokenBucket:
    def __init__(self, rate: float, capacity: int):
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.last_refill = time.monotonic()

    def _refill(self):
        now = time.monotonic()
        elapsed = now - self.last_refill
        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
        self.last_refill = now

    def consume(self, count: int = 1) -> bool:
        self._refill()
        if self.tokens >= count:
            self.tokens -= count
            return True
        return False


bucket = TokenBucket(rate=20, capacity=40)


def rate_limit(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        if not bucket.consume():
            raise HTTPException(status_code=429, detail="速率限制已触发")
        return await func(*args, **kwargs)
    return wrapper


# ---------------------------------------------------------------------------
# 请求模型
# ---------------------------------------------------------------------------

class ChatRequest(BaseModel):
    messages: list
    stream: bool = True


# ---------------------------------------------------------------------------
# 模拟 vLLM 异步生成器（实际使用时替换为 vLLM 的 AsyncLLMEngine）
# ---------------------------------------------------------------------------

async def mock_vllm_generate(prompt: str) -> AsyncGenerator[str, None]:
    """模拟 vLLM 逐 token 生成。"""
    response = f"这是对'{prompt[:20]}...'的回复"
    for char in response:
        await asyncio.sleep(0.03)
        yield char


# ---------------------------------------------------------------------------
# SSE 流式接口
# ---------------------------------------------------------------------------

@app.post("/v1/chat/completions")
@rate_limit
async def chat_completions(req: ChatRequest):
    prompt = req.messages[-1]["content"] if req.messages else ""

    async def event_stream() -> AsyncGenerator[str, None]:
        input_tokens = len(prompt)
        output_tokens = 0
        full_response = ""

        async for token in mock_vllm_generate(prompt):
            output_tokens += 1
            full_response += token
            yield f"data: {json.dumps({'choices': [{'delta': {'content': token}}]})}\n\n"

        # 发送用量统计
        usage = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
        }
        yield f"data: {json.dumps({'choices': [{'delta': {}}], 'usage': usage})}\n\n"
        yield "data: [DONE]\n\n"

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


# ---------------------------------------------------------------------------
# WebSocket 双向接口
# ---------------------------------------------------------------------------

@app.websocket("/ws/chat")
async def ws_chat(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_json()
            prompt = data.get("prompt", "")
            input_tokens = len(prompt)
            output_tokens = 0

            async for token in mock_vllm_generate(prompt):
                output_tokens += 1
                await websocket.send_json({"token": token})
                # 检查取消信号
                try:
                    msg = await asyncio.wait_for(
                        websocket.receive_text(), timeout=0.005
                    )
                    if msg == "cancel":
                        await websocket.send_json({"event": "cancelled"})
                        break
                except asyncio.TimeoutError:
                    pass

            await websocket.send_json({
                "event": "done",
                "usage": {
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "total_tokens": input_tokens + output_tokens,
                },
            })
    except WebSocketDisconnect:
        pass


# ---------------------------------------------------------------------------
# 启动入口
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

启动与验证

# 安装依赖
pip install fastapi uvicorn transformers

# 启动服务
python main.py

# 测试流式接口
curl -N -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"你好，请介绍一下你自己"}],"stream":true}'

与 vLLM 集成时，只需将 mock_vllm_generate 替换为 AsyncLLMEngine 的 generate 方法调用，并在服务启动时加载模型。vLLM 本身已内置 PagedAttention 高效显存管理和动态批处理，与上述队列、限流机制配合可构建生产级别的 AI API 服务。

总结

本文介绍了基于 FastAPI 构建 AI API 服务的核心技术方案：

• SSE 和 WebSocket 分别适用于单向推送和双向通信场景
• 请求队列与令牌桶限流 保护后端资源，确保服务稳定性
• Token 用量计量 为按量计费提供数据基础

实际生产部署还需额外考虑身份认证、模型热加载、健康检查、日志追踪等基础设施，但上述核心模式构成了 AI API 服务的骨架，可在此基础上按需扩展。