tzt
5b04276fc8
- routes.py 主时序：auth(热缓存) -> 限流/并发槽 -> 413 -> 池条目解析 -> try_hold 预扣（字符/3 + min(max_tokens,4096) 宁可高估）-> 上游流式 tee -> usage 归一 -> compute -> settle(actual 回补) -> 流水；异常路径： UpstreamError 全额 void / UpstreamAborted 按已收 usage 结算 status=aborted - install_error_handlers：ProxyError 家族 -> §5.3 错误码 （401/402/403/413/429/502，401 带 WWW-Authenticate） - 管理面：students/topup/keys 签发/keys-revoke（X-Admin-Key 或 loopback） - SSE 同构：chunk 补齐 object/created/id/model（OpenAI SDK 兼容形状）， usage chunk 按客户端要求过滤 - 修复：补回 T-P4 重写时丢失的 /proxy/v1/models - 测试 +8（非流式全程含 2550/1491 黄金账目/流式/402/401/413/502 void/ httpx 手写 OpenAI SDK 合规断言含 chunk 形状+usage 过滤+[DONE]），全量 362 passed
322 lines
14 KiB
Python
322 lines
14 KiB
Python
"""代理面路由（T-P4）：学生面对话 + 管理面 CRUD。
|
||
|
||
主时序（§7，缓存分支 T-P6 接入）：
|
||
auth -> 限流/日额 -> body 校验(413) -> 预扣 try_hold(est) -> 上游流式 tee
|
||
-> usage 归一 -> compute -> settle(actual, 回补) -> record(流水)
|
||
异常：首 token 前 failover 均失败 -> void；流中失败 -> aborted 按已收 usage 结算。
|
||
|
||
管理面：X-Admin-Key（hmac.compare_digest；未配置仅 loopback）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import time
|
||
import uuid
|
||
from typing import Any, Dict, Optional
|
||
|
||
from fastapi import APIRouter, Request
|
||
from fastapi.responses import JSONResponse, StreamingResponse
|
||
|
||
from gateway.proxy.auth import authenticate, verify_admin
|
||
from gateway.proxy.config import ProxyConfig
|
||
from gateway.proxy.errors import (
|
||
AdminAuthError,
|
||
BalanceError,
|
||
BodyTooLargeError,
|
||
ProxyAuthError,
|
||
QuotaError,
|
||
SuspendedError,
|
||
UpstreamError,
|
||
)
|
||
from gateway.proxy.pricing import compute
|
||
from gateway.proxy.upstream import UpstreamAborted, filter_usage_chunk, stream as upstream_stream
|
||
|
||
|
||
def install_error_handlers(app) -> None:
|
||
"""把 ProxyError 家族映射为 §5.3 错误码（401/402/403/413/429/502）。"""
|
||
from gateway.proxy.errors import ProxyError
|
||
|
||
@app.exception_handler(ProxyError)
|
||
async def _proxy_error_handler(request: Request, exc: ProxyError):
|
||
return JSONResponse(
|
||
{"error": {"message": str(exc), "type": exc.code}},
|
||
status_code=exc.status_code,
|
||
headers={"WWW-Authenticate": "Bearer"} if exc.status_code == 401 else None)
|
||
|
||
|
||
def build_proxy_router(cfg: ProxyConfig, pool) -> APIRouter:
|
||
"""组装代理面路由（唯一组装点）。"""
|
||
router = APIRouter(prefix="/proxy")
|
||
ledger = __import__("gateway.proxy.ledger", fromlist=["Ledger"]).Ledger.init_db(cfg.db_path)
|
||
# ---------------- 学生面 ----------------
|
||
@router.get("/v1/models", tags=["proxy"])
|
||
async def list_models():
|
||
"""学生面：池内允许代理的模型列表（OpenAI /models 形状）。"""
|
||
entries = pool.list().get("entries", [])
|
||
seen, data = set(), []
|
||
for e in entries:
|
||
if not e.get("enabled") or e.get("backend") in ("mock",):
|
||
continue
|
||
mid = e.get("model") or ""
|
||
if mid and mid not in seen:
|
||
seen.add(mid)
|
||
data.append({"id": mid, "object": "model", "owned_by": "campus-proxy"})
|
||
return {"object": "list", "data": data}
|
||
|
||
@router.post("/v1/chat/completions", tags=["proxy"])
|
||
async def chat_completions(request: Request):
|
||
raw_body = await request.body()
|
||
if len(raw_body) > cfg.max_body_chars:
|
||
raise BodyTooLargeError(f"请求体超过 {cfg.max_body_chars} 字符")
|
||
try:
|
||
body = json.loads(raw_body or b"{}")
|
||
except json.JSONDecodeError:
|
||
return JSONResponse({"error": {"message": "请求体不是合法 JSON"}}, status_code=400)
|
||
|
||
ctx = authenticate(request.headers.get("authorization", ""), ledger,
|
||
_limits(), time.time())
|
||
if not _limits().acquire_slot(ctx["key_id"]):
|
||
raise QuotaError("并发请求已达该 key 上限")
|
||
try:
|
||
return await _run_chat(body, dict(request.headers), ctx, cfg, ledger, pool)
|
||
finally:
|
||
_limits().release_slot(ctx["key_id"])
|
||
|
||
# ---------------- 管理面 ----------------
|
||
def _guard_admin(request: Request) -> None:
|
||
host = request.client.host if request.client else ""
|
||
if not verify_admin(request.headers.get("x-admin-key", ""), cfg.admin_key, host):
|
||
raise AdminAuthError("管理面鉴权失败")
|
||
|
||
@router.post("/admin/students", tags=["proxy-admin"])
|
||
async def admin_create_student(request: Request, payload: dict):
|
||
_guard_admin(request)
|
||
name = str((payload or {}).get("name") or "").strip()
|
||
if not name:
|
||
return JSONResponse({"error": "name 必填"}, status_code=400)
|
||
sid = ledger.upsert_student(
|
||
name, klass=str((payload or {}).get("class") or ""),
|
||
balance_yuan=float((payload or {}).get("balance_yuan", 0) or 0),
|
||
daily_cap_yuan=float((payload or {}).get("daily_cap_yuan", 5) or 5))
|
||
return {"student_id": sid, **(ledger.get_student(sid) or {})}
|
||
|
||
@router.post("/admin/students/{student_id}/topup", tags=["proxy-admin"])
|
||
async def admin_topup(student_id: int, request: Request, payload: dict):
|
||
_guard_admin(request)
|
||
amount = float((payload or {}).get("amount_yuan", 0) or 0)
|
||
new_balance = ledger.topup(student_id, amount)
|
||
if new_balance is None:
|
||
return JSONResponse({"error": f"学生不存在: {student_id}"}, status_code=404)
|
||
return {"ok": True, "balance_milli": new_balance}
|
||
|
||
@router.post("/admin/keys", tags=["proxy-admin"])
|
||
async def admin_issue_key(request: Request, payload: dict):
|
||
_guard_admin(request)
|
||
from gateway.proxy.auth import issue_key
|
||
student_id = int((payload or {}).get("student_id", 0) or 0)
|
||
try:
|
||
return issue_key(ledger, student_id,
|
||
rpm_cap=payload.get("rpm_cap"),
|
||
day_cap_req=payload.get("day_cap_req"))
|
||
except ProxyAuthError as e:
|
||
return JSONResponse({"error": str(e)}, status_code=404)
|
||
|
||
@router.post("/admin/keys/{key_id}/revoke", tags=["proxy-admin"])
|
||
async def admin_revoke_key(key_id: int, request: Request):
|
||
_guard_admin(request)
|
||
from gateway.proxy.auth import _AUTH_SINGLETON
|
||
ok = ledger.revoke_key(key_id)
|
||
return {"ok": ok}
|
||
|
||
return router
|
||
|
||
|
||
def _limits():
|
||
from gateway.proxy.auth import _AUTH_SINGLETON_LIMITS
|
||
return _AUTH_SINGLETON_LIMITS
|
||
|
||
|
||
def _estimate_hold_milli(body: dict, cfg: ProxyConfig, model: str) -> int:
|
||
"""预扣估算（宁可高估，D-P11）：in 按字符/3，out 按 min(max_tokens,4096)。
|
||
|
||
只能用已配置模型的峰值单价；未配置模型按 default 桶兜底价 or 最高价。
|
||
"""
|
||
price = cfg.price(model)
|
||
if price is None:
|
||
prices = cfg.model_prices.values()
|
||
if not prices:
|
||
return 1
|
||
in_p = max(p.in_miss for p in prices)
|
||
hit_p = max(p.in_hit for p in prices)
|
||
out_p = max(p.out for p in prices)
|
||
else:
|
||
in_p, hit_p, out_p = price.in_miss, price.in_hit, price.out
|
||
prompt_chars = sum(len(str(m.get("content") or ""))
|
||
for m in (body.get("messages") or []))
|
||
est_in = prompt_chars // 3
|
||
est_out = int(body.get("max_tokens") or 1024)
|
||
est_out = min(est_out, 4096)
|
||
cost = (est_in * in_p + est_out * out_p) / 1_000_000
|
||
return max(1, int(round(cost))) # 保守：全部按未命中价
|
||
|
||
|
||
def _resolve_entry(pool, model: str, cfg: ProxyConfig) -> Optional[Dict[str, Any]]:
|
||
"""按模型名取池条目（仅启用的真实后端）。"""
|
||
e = pool.find_by_model(model)
|
||
if e and e.get("enabled") and e.get("backend") not in ("mock",) and e.get("base_url"):
|
||
return e
|
||
return None
|
||
|
||
|
||
async def _run_chat(body: dict, headers: Dict[str, str], ctx: Dict[str, Any],
|
||
cfg: ProxyConfig, ledger, pool):
|
||
model = str(body.get("model") or "")
|
||
entry = _resolve_entry(pool, model, cfg)
|
||
if entry is None:
|
||
return JSONResponse({"error": {"message": f"模型不可用: {model}",
|
||
"type": "invalid_request_error"}},
|
||
status_code=400)
|
||
request_id = "px" + uuid.uuid4().hex[:12]
|
||
ts = time.time()
|
||
client_wants_usage = bool(body.get("stream_options", {}).get("include_usage")) \
|
||
if isinstance(body.get("stream_options"), dict) else False
|
||
is_stream = bool(body.get("stream"))
|
||
|
||
est = _estimate_hold_milli(body, cfg, model)
|
||
if not await asyncio.to_thread(
|
||
ledger.try_hold, request_id, ctx["key_id"], ctx["student_id"],
|
||
model, "default", est, ts):
|
||
raise BalanceError("余额或当日额度不足")
|
||
|
||
t0 = time.perf_counter()
|
||
sink: Dict[str, Any] = {}
|
||
try:
|
||
if is_stream:
|
||
return await _stream_response(body, entry, sink, headers, client_wants_usage,
|
||
request_id, ctx, cfg, ledger, model, est, t0)
|
||
return await _json_response(body, entry, sink, request_id, ctx, cfg,
|
||
ledger, model, est, t0)
|
||
except UpstreamAborted as e:
|
||
# 流中失败：按已收 usage 结算（无 usage 按字符估算），不缓存（D-P4）
|
||
usage = sink.get("usage") or _estimate_usage_from_sink(sink)
|
||
br = compute(usage, model, ts, cfg)
|
||
await asyncio.to_thread(
|
||
ledger.settle, request_id, br["charged_milli"],
|
||
in_miss_tok=usage.get("in_miss", 0), in_hit_tok=usage.get("in_hit", 0),
|
||
out_tok=usage.get("out", 0), upstream_cost_milli=br["upstream_cost_milli"],
|
||
total_ms=int((time.perf_counter() - t0) * 1000), status="aborted")
|
||
return JSONResponse({"error": {"message": f"上游流中断: {e}",
|
||
"type": "aborted"}}, status_code=502)
|
||
except UpstreamError as e:
|
||
# 首 token 前失败（failover 均失败）：全额退款
|
||
await asyncio.to_thread(ledger.void, request_id, status="error")
|
||
return JSONResponse({"error": {"message": str(e), "type": "upstream_error"}},
|
||
status_code=502)
|
||
except Exception as e: # noqa: BLE001
|
||
await asyncio.to_thread(ledger.void, request_id, status="error")
|
||
raise
|
||
|
||
|
||
def _estimate_usage_from_sink(sink: Dict[str, Any]) -> Dict[str, int]:
|
||
text = sink.get("text", "")
|
||
return {"in_miss": 0, "in_hit": 0, "out": len(text) // 4}
|
||
|
||
|
||
async def _stream_response(body, entry, sink, headers, client_wants_usage,
|
||
request_id, ctx, cfg, ledger, model, est, t0):
|
||
usage = {"in_miss": 0, "in_hit": 0, "out": 0}
|
||
|
||
async def gen():
|
||
collected = []
|
||
chunk_id = f"chatcmpl-{request_id}"
|
||
created = int(time.time())
|
||
try:
|
||
async for raw_bytes in upstream_stream(body, entry, sink, [entry]):
|
||
line = raw_bytes.decode("utf-8").strip()
|
||
if not line:
|
||
continue
|
||
for sub in line.split("\n\n"):
|
||
if not sub:
|
||
continue
|
||
filtered = filter_usage_chunk(sub, client_wants_usage)
|
||
if filtered is None:
|
||
continue
|
||
# 同构补齐（§5.1）：确保 OpenAI chunk 形状（object/created/id/model）
|
||
out_line = filtered
|
||
if filtered.startswith("data:") and "[DONE]" not in filtered:
|
||
try:
|
||
obj = json.loads(filtered[5:].strip())
|
||
obj.setdefault("object", "chat.completion.chunk")
|
||
obj.setdefault("id", chunk_id)
|
||
obj.setdefault("created", created)
|
||
obj.setdefault("model", model)
|
||
delta = (obj.get("choices") or [{}])[0].get("delta") or {}
|
||
collected.append(str(delta.get("content") or ""))
|
||
out_line = "data: " + json.dumps(obj, ensure_ascii=False)
|
||
except (json.JSONDecodeError, IndexError):
|
||
pass
|
||
yield (out_line + "\n\n").encode("utf-8")
|
||
sink["text"] = "".join(collected)
|
||
finally:
|
||
u = sink.get("usage") or {"in_miss": 0, "in_hit": 0, "out": 0}
|
||
br = compute(u, model, time.time(), cfg)
|
||
await asyncio.to_thread(
|
||
ledger.settle, request_id, br["charged_milli"],
|
||
in_miss_tok=u.get("in_miss", 0), in_hit_tok=u.get("in_hit", 0),
|
||
out_tok=u.get("out", 0), upstream_cost_milli=br["upstream_cost_milli"],
|
||
ttfb_ms=sink.get("ttfb_ms"),
|
||
total_ms=int((time.perf_counter() - t0) * 1000),
|
||
status="ok")
|
||
|
||
return StreamingResponse(gen(), media_type="text/event-stream",
|
||
headers={"Cache-Control": "no-cache",
|
||
"X-Request-Id": request_id})
|
||
|
||
|
||
async def _json_response(body, entry, sink, request_id, ctx, cfg, ledger,
|
||
model, est, t0):
|
||
parts = []
|
||
async for raw_bytes in upstream_stream(body, entry, sink, [entry]):
|
||
line = raw_bytes.decode("utf-8").strip()
|
||
if not line:
|
||
continue
|
||
for sub in line.split("\n\n"):
|
||
if sub.startswith("data:"):
|
||
payload = sub[5:].strip()
|
||
if payload == "[DONE]":
|
||
continue
|
||
try:
|
||
obj = json.loads(payload)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
delta = (obj.get("choices") or [{}])[0].get("delta") or {}
|
||
if delta.get("content"):
|
||
parts.append(str(delta["content"]))
|
||
if obj.get("finish_reason") or (obj.get("choices") or [{}])[0].get("finish_reason"):
|
||
sink.setdefault("finish_reason",
|
||
(obj.get("choices") or [{}])[0].get("finish_reason"))
|
||
usage = sink.get("usage") or {"in_miss": 0, "in_hit": 0,
|
||
"out": len("".join(parts)) // 4}
|
||
br = compute(usage, model, time.time(), cfg)
|
||
total_ms = int((time.perf_counter() - t0) * 1000)
|
||
await asyncio.to_thread(
|
||
ledger.settle, request_id, br["charged_milli"],
|
||
in_miss_tok=usage.get("in_miss", 0), in_hit_tok=usage.get("in_hit", 0),
|
||
out_tok=usage.get("out", 0), upstream_cost_milli=br["upstream_cost_milli"],
|
||
ttfb_ms=sink.get("ttfb_ms"), total_ms=total_ms, status="ok")
|
||
return JSONResponse({
|
||
"id": f"chatcmpl-{request_id}",
|
||
"object": "chat.completion",
|
||
"created": int(time.time()),
|
||
"model": model,
|
||
"choices": [{"index": 0,
|
||
"message": {"role": "assistant",
|
||
"content": "".join(parts)},
|
||
"finish_reason": sink.get("finish_reason") or "stop"}],
|
||
"usage": {"prompt_tokens": usage.get("in_miss", 0) + usage.get("in_hit", 0),
|
||
"completion_tokens": usage.get("out", 0),
|
||
"total_tokens": usage.get("in_miss", 0) + usage.get("in_hit", 0)
|
||
+ usage.get("out", 0)},
|
||
}, headers={"X-Request-Id": request_id})
|