tzt
a043493548
- routes._fallback_chain：主条目之外的启用真实后端按（档位升序，单价和升序） 构成有序候选链（≤3），排除 mock/停用/同模型，仅作 D-P4 首 token 前 failover - _stream_response/_json_response 接入链式候选（原 [entry,entry] 同条目重试升级为真降级链） - upstream.stream：failover 发生即记 sink[upstream_fallback] 与模块级统计； JSON 路径经 X-Upstream-Fallback/Original/Used/Reason 四头如实上报 （流式路径头已发出不可追溯，由统计聚合暴露） - /proxy/admin/stats 新增 upstream_failover 块（total + recent 20 条） pytest 447 passed（T-X1 后 443 + 4）
185 lines
7.6 KiB
Python
185 lines
7.6 KiB
Python
"""上游客户端（T-P2）：流式派发 + usage 注入/过滤 + 三家归一化 + 首 token 前 failover。
|
||
|
||
锁定决策：
|
||
- D-P4：仅在上游**首 token 返回前**允许切换备用条目；流中失败 = aborted（由调用方
|
||
按已收 usage 计费），本模块以 UpstreamAborted 标记。
|
||
- §5.1：代理向上游始终注入 stream_options.include_usage（计量不依赖客户端行为）；
|
||
客户端未要求 usage 时，透传层过滤该 chunk 不下发（routes 的 tee 负责，见 filter_usage_chunk）。
|
||
- httpx.AsyncClient 模块级单例（keepalive；limits.max_connections=100），
|
||
超时 connect=10s / read=120s / write=10s / pool=30s。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from collections import deque
|
||
from typing import Any, AsyncIterator, Deque, Dict, List, Optional
|
||
|
||
import httpx
|
||
|
||
from gateway.proxy.errors import UpstreamError
|
||
|
||
_TIMEOUT = httpx.Timeout(connect=10.0, read=120.0, write=10.0, pool=30.0)
|
||
_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=20)
|
||
|
||
_client: Optional[httpx.AsyncClient] = None
|
||
|
||
|
||
def get_client() -> httpx.AsyncClient:
|
||
"""模块级单例（keepalive 连接池）。"""
|
||
global _client
|
||
if _client is None:
|
||
_client = httpx.AsyncClient(timeout=_TIMEOUT, limits=_LIMITS)
|
||
return _client
|
||
|
||
|
||
async def close_client() -> None:
|
||
global _client
|
||
if _client is not None:
|
||
await _client.aclose()
|
||
_client = None
|
||
|
||
|
||
class UpstreamAborted(Exception):
|
||
"""首 token 已下发后上游失败（D-P4：不可 failover，按已收 usage 计费）。"""
|
||
|
||
|
||
# ---------------- failover 统计（T-X2：如实上报，进程内，D-P9 单进程） ----------------
|
||
_FAILOVER_EVENTS_MAX = 100
|
||
_failover_events: Deque[Dict[str, Any]] = deque(maxlen=_FAILOVER_EVENTS_MAX)
|
||
_failover_total = 0
|
||
|
||
|
||
def _record_failover(original: str, failed_model: str, reason: str) -> None:
|
||
"""记录一次「首 token 前失败 -> 切换候选」事件。"""
|
||
global _failover_total
|
||
_failover_total += 1
|
||
_failover_events.append({
|
||
"ts": int(time.time()), "original": original,
|
||
"failed_model": failed_model, "reason": reason,
|
||
})
|
||
|
||
|
||
def failover_stats(recent: int = 20) -> Dict[str, Any]:
|
||
"""failover 统计（/proxy/admin/stats 聚合暴露）。"""
|
||
events = list(_failover_events)
|
||
return {"total": _failover_total,
|
||
"recent": events[-max(0, int(recent)):][::-1]}
|
||
|
||
|
||
def reset_failover_stats() -> None:
|
||
"""测试用：清空 failover 统计。"""
|
||
global _failover_total
|
||
_failover_events.clear()
|
||
_failover_total = 0
|
||
|
||
|
||
def normalize_usage(provider: str, usage_dict: Dict[str, Any]) -> Dict[str, int]:
|
||
"""三家 usage 字段 -> 统一 {in_miss, in_hit, out}。
|
||
|
||
- deepseek：prompt_cache_hit_tokens / prompt_tokens
|
||
- openai 兼容：prompt_tokens_details.cached_tokens
|
||
- anthropic：cache_read_input_tokens / input_tokens
|
||
"""
|
||
usage = usage_dict or {}
|
||
out = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
|
||
total_in = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
|
||
if provider == "deepseek":
|
||
hit = int(usage.get("prompt_cache_hit_tokens") or 0)
|
||
elif provider == "anthropic":
|
||
hit = int(usage.get("cache_read_input_tokens") or 0)
|
||
else: # openai / 兼容端点
|
||
details = usage.get("prompt_tokens_details") or {}
|
||
hit = int(details.get("cached_tokens") or 0) if isinstance(details, dict) else 0
|
||
hit = min(hit, total_in)
|
||
return {"in_miss": total_in - hit, "in_hit": hit, "out": out}
|
||
|
||
|
||
def _inject_usage_option(body: dict) -> dict:
|
||
"""浅拷贝注入 stream_options.include_usage（§5.1 计量不依赖客户端行为）。"""
|
||
shaped = dict(body)
|
||
shaped["stream_options"] = {"include_usage": True}
|
||
return shaped
|
||
|
||
|
||
def filter_usage_chunk(raw_line: str, client_wants_usage: bool) -> Optional[str]:
|
||
"""透传过滤：客户端未要求 usage 时剥除 usage 字段所在 chunk（返回 None = 丢弃）。"""
|
||
if client_wants_usage or not raw_line.startswith("data:"):
|
||
return raw_line
|
||
payload = raw_line[5:].strip()
|
||
if payload == "[DONE]":
|
||
return raw_line
|
||
try:
|
||
obj = json.loads(payload)
|
||
except json.JSONDecodeError:
|
||
return raw_line
|
||
if obj.get("usage"):
|
||
return None
|
||
return raw_line
|
||
|
||
|
||
async def stream(body: dict, entry: Dict[str, Any], usage_sink: Dict[str, Any],
|
||
fallback_entries: Optional[List[Dict[str, Any]]] = None
|
||
) -> AsyncIterator[bytes]:
|
||
"""流式派发（§6 签名扩展 failover 链；T-X2 扩展降级上报）。
|
||
|
||
- 始终注入 include_usage；usage chunk 交给 usage_sink（透传过滤由调用方用
|
||
filter_usage_chunk 决定）。
|
||
- D-P4：首 token 前（未 yield 任何字节）失败 -> 切换 fallback_entries；
|
||
已 yield 后失败 -> 抛 UpstreamAborted。
|
||
- usage_sink["usage"] 收敛为归一化 dict；sink["ttfb_ms"] 记录首字节耗时。
|
||
- T-X2 如实上报：failover 发生时记 sink["upstream_fallback"]
|
||
{used, original, used_model, reason}，并计入模块级 failover 统计
|
||
（failover_stats / reset_failover_stats），/proxy/admin/stats 聚合暴露。
|
||
"""
|
||
import time as _time
|
||
candidates: List[Dict[str, Any]] = [entry] + list(fallback_entries or [])
|
||
last_err: Optional[Exception] = None
|
||
head_model = str(candidates[0].get("model") or "")
|
||
first_failure: Optional[Dict[str, Any]] = None
|
||
|
||
for cand in candidates:
|
||
provider = str(cand.get("provider") or "openai")
|
||
url = cand["base_url"].rstrip("/") + "/chat/completions"
|
||
headers = {"Authorization": f"Bearer {cand['api_key']}"} if cand.get("api_key") else {}
|
||
payload = _inject_usage_option(body)
|
||
t0 = _time.perf_counter()
|
||
first = False
|
||
try:
|
||
client = get_client()
|
||
async with client.stream("POST", url, headers=headers, json=payload) as resp:
|
||
resp.raise_for_status()
|
||
async for line in resp.aiter_lines():
|
||
if not line.startswith("data:"):
|
||
continue
|
||
raw = line
|
||
payload_txt = line[5:].strip()
|
||
if payload_txt == "[DONE]":
|
||
yield (raw + "\n\n").encode("utf-8")
|
||
continue
|
||
try:
|
||
obj = json.loads(payload_txt)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if obj.get("usage"):
|
||
usage_sink["usage"] = normalize_usage(provider, obj["usage"])
|
||
if not first:
|
||
first = True
|
||
usage_sink["ttfb_ms"] = int((_time.perf_counter() - t0) * 1000)
|
||
yield (raw + "\n\n").encode("utf-8")
|
||
if first_failure is not None:
|
||
first_failure["used_model"] = str(cand.get("model") or "")
|
||
usage_sink["upstream_fallback"] = dict(first_failure)
|
||
return # 正常结束
|
||
except Exception as exc: # noqa: BLE001
|
||
if first:
|
||
raise UpstreamAborted(str(exc)) from exc
|
||
last_err = exc
|
||
if first_failure is None:
|
||
first_failure = {"used": True, "original": head_model,
|
||
"used_model": "", "reason": f"{type(exc).__name__}: {exc}"}
|
||
_record_failover(head_model, str(cand.get("model") or ""),
|
||
f"{type(exc).__name__}: {exc}")
|
||
continue # 首 token 前失败 -> failover
|
||
raise UpstreamError(f"上游均不可用: {type(last_err).__name__}: {last_err}")
|