153 lines
5.6 KiB
Python
153 lines
5.6 KiB
Python
# requires: fastapi, httpx
|
|
"""流式对话 API：鉴权、参数校验、请求 ID、统一错误格式、SSE 流式输出、依赖注入替换模型客户端。
|
|
|
|
运行：uvicorn app:app --reload（或 python app.py 打印一次演示）
|
|
测试：python -m pytest -q test_app.py
|
|
模型客户端是假的（逐字返回固定文本），换成真实 SDK 时只需要实现同样的 stream 方法。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from typing import Annotated, Literal, Protocol
|
|
|
|
from fastapi import Depends, FastAPI, Header, HTTPException, Request
|
|
from fastapi.exceptions import RequestValidationError
|
|
from fastapi.responses import JSONResponse, StreamingResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
log = logging.getLogger("chat-api")
|
|
|
|
|
|
# ---------- 模型 ----------
|
|
class Message(BaseModel):
|
|
role: Literal["system", "user", "assistant"]
|
|
content: str = Field(min_length=1, max_length=4000)
|
|
|
|
|
|
class ChatRequest(BaseModel):
|
|
messages: list[Message] = Field(min_length=1, max_length=50)
|
|
temperature: float = Field(default=0.3, ge=0, le=2)
|
|
stream: bool = True
|
|
|
|
|
|
class ChatResponse(BaseModel):
|
|
request_id: str
|
|
content: str
|
|
|
|
|
|
# ---------- 依赖 ----------
|
|
class LLMClient(Protocol):
|
|
def stream(self, messages: list[Message], temperature: float) -> AsyncIterator[str]: ...
|
|
|
|
|
|
class FakeLLM:
|
|
async def stream(self, messages: list[Message], temperature: float) -> AsyncIterator[str]:
|
|
answer = f"你问的是：{messages[-1].content}。这是一个演示回答。"
|
|
for ch in answer:
|
|
await asyncio.sleep(0.005)
|
|
yield ch
|
|
|
|
|
|
def get_llm() -> LLMClient:
|
|
return FakeLLM()
|
|
|
|
|
|
API_KEYS = {"demo-key": "tenant-a"} # 演示用；真实项目从数据库或密钥服务读取，并做哈希比对
|
|
|
|
|
|
def current_tenant(x_api_key: Annotated[str | None, Header()] = None) -> str:
|
|
tenant = API_KEYS.get(x_api_key or "")
|
|
if tenant is None:
|
|
raise HTTPException(status_code=401, detail="API Key 无效")
|
|
return tenant
|
|
|
|
|
|
# ---------- 应用 ----------
|
|
app = FastAPI(title="chat-api")
|
|
|
|
|
|
@app.middleware("http")
|
|
async def request_id_and_timing(request: Request, call_next):
|
|
rid = request.headers.get("X-Request-ID") or uuid.uuid4().hex[:16]
|
|
request.state.request_id = rid
|
|
start = time.perf_counter()
|
|
response = await call_next(request)
|
|
response.headers["X-Request-ID"] = rid
|
|
log.info("%s %s %s %.1fms", request.method, request.url.path, response.status_code,
|
|
(time.perf_counter() - start) * 1000)
|
|
return response
|
|
|
|
|
|
def error_body(request: Request, code: str, message: str, details=None) -> dict:
|
|
body = {"code": code, "message": message, "request_id": getattr(request.state, "request_id", "")}
|
|
if details is not None:
|
|
body["details"] = details
|
|
return body
|
|
|
|
|
|
@app.exception_handler(HTTPException)
|
|
async def http_error(request: Request, exc: HTTPException):
|
|
return JSONResponse(error_body(request, f"HTTP_{exc.status_code}", str(exc.detail)), status_code=exc.status_code)
|
|
|
|
|
|
@app.exception_handler(RequestValidationError)
|
|
async def validation_error(request: Request, exc: RequestValidationError):
|
|
details = [{"field": ".".join(str(p) for p in e["loc"][1:]), "message": e["msg"]} for e in exc.errors()]
|
|
return JSONResponse(error_body(request, "INVALID_PARAMS", "参数校验失败", details), status_code=422)
|
|
|
|
|
|
@app.exception_handler(Exception)
|
|
async def unhandled_error(request: Request, exc: Exception):
|
|
log.exception("未处理的异常") # 详细信息只进日志，不返回给调用方
|
|
return JSONResponse(error_body(request, "INTERNAL", "服务器内部错误"), status_code=500)
|
|
|
|
|
|
@app.get("/healthz")
|
|
async def healthz() -> dict:
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/v1/chat", response_model=None)
|
|
async def chat(req: ChatRequest, request: Request,
|
|
tenant: Annotated[str, Depends(current_tenant)],
|
|
llm: Annotated[LLMClient, Depends(get_llm)]):
|
|
rid = request.state.request_id
|
|
if not req.stream:
|
|
content = "".join([t async for t in llm.stream(req.messages, req.temperature)])
|
|
return ChatResponse(request_id=rid, content=content)
|
|
|
|
async def events() -> AsyncIterator[str]:
|
|
try:
|
|
async for token in llm.stream(req.messages, req.temperature):
|
|
if await request.is_disconnected(): # 客户端断开就停止生成，节省模型费用
|
|
log.info("客户端已断开，停止生成 request_id=%s", rid)
|
|
return
|
|
yield f"data: {json.dumps({'delta': token}, ensure_ascii=False)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
except Exception:
|
|
log.exception("生成失败 request_id=%s", rid)
|
|
# 响应头已经发出，不能再改状态码，只能在流里发送错误事件
|
|
yield f"event: error\ndata: {json.dumps({'message': '生成失败', 'request_id': rid}, ensure_ascii=False)}\n\n"
|
|
|
|
return StreamingResponse(events(), media_type="text/event-stream",
|
|
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
|
|
|
|
|
|
@app.get("/boom")
|
|
async def boom():
|
|
raise RuntimeError("数据库连接串 postgres://secret@...") # 演示：内部细节不会泄露给调用方
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from fastapi.testclient import TestClient
|
|
|
|
with TestClient(app) as c:
|
|
r = c.post("/v1/chat", headers={"X-API-Key": "demo-key"},
|
|
json={"messages": [{"role": "user", "content": "年假怎么算"}], "stream": False})
|
|
print(r.status_code, r.json())
|