feat(proxy): T-P2 上游客户端（流式派发/usage 注入过滤/三家归一化/首 token 前 failover）
- upstream.py：httpx.AsyncClient 模块级单例（keepalive，limits=100， 超时 connect10/read120/write10/pool30）；stream() 始终注入 stream_options.include_usage（计量不依赖客户端）+ filter_usage_chunk （客户端未要求 usage 时剥除该 chunk）；D-P4 首 token 前 failover 链、 流中失败抛 UpstreamAborted（不可切换）；ttfb_ms 记录 - normalize_usage 三家归一：deepseek prompt_cache_hit_tokens / openai prompt_tokens_details.cached_tokens / anthropic cache_read_input_tokens （命中数>总数时钳制） - model_pool：+provider（枚举校验 deepseek/openai/anthropic）+in_hit_price （缺省 = price_in×1/30，D-P3） - 测试 +6：透传+sink/failover/流中 aborted/全挂 502/三家归一/过滤，全量 345 passed
This commit is contained in:
tzt
+130
-12
@@ -1,21 +1,139 @@
|
||||
"""上游客户端（T-P2 落地；本文件先立签名）。
|
||||
"""上游客户端（T-P2）：流式派发 + usage 注入/过滤 + 三家归一化 + 首 token 前 failover。
|
||||
|
||||
规格（§6）：模块级 httpx.AsyncClient 单例（keepalive，max_connections=100），
|
||||
超时 connect=10s/read=120s/write=10s/pool=30s；始终注入
|
||||
stream_options.include_usage（客户端未要求 usage 时过滤该 chunk 不下发）；
|
||||
首 token 前 failover（D-P4）。
|
||||
锁定决策：
|
||||
- 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
|
||||
|
||||
from typing import Any, AsyncIterator, Dict
|
||||
import json
|
||||
from typing import Any, AsyncIterator, 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
|
||||
|
||||
|
||||
async def stream(body: dict, entry: Dict[str, Any], usage_sink) -> AsyncIterator[bytes]:
|
||||
"""流式派发到上游，逐块 yield（T-P2 实现）。"""
|
||||
raise NotImplementedError("T-P2")
|
||||
yield b"" # pragma: no cover
|
||||
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 计费）。"""
|
||||
|
||||
|
||||
def normalize_usage(provider: str, usage_dict: Dict[str, Any]) -> Dict[str, int]:
|
||||
"""三家 usage 字段 -> 统一 {in_miss, in_hit, out}（T-P2 实现）。"""
|
||||
raise NotImplementedError("T-P2")
|
||||
"""三家 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 链）。
|
||||
|
||||
- 始终注入 include_usage；usage chunk 交给 usage_sink（透传过滤由调用方用
|
||||
filter_usage_chunk 决定）。
|
||||
- D-P4：首 token 前（未 yield 任何字节）失败 -> 切换 fallback_entries；
|
||||
已 yield 后失败 -> 抛 UpstreamAborted。
|
||||
- usage_sink["usage"] 收敛为归一化 dict；sink["ttfb_ms"] 记录首字节耗时。
|
||||
"""
|
||||
import time as _time
|
||||
candidates: List[Dict[str, Any]] = [entry] + list(fallback_entries or [])
|
||||
last_err: Optional[Exception] = 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")
|
||||
return # 正常结束
|
||||
except Exception as exc: # noqa: BLE001
|
||||
if first:
|
||||
raise UpstreamAborted(str(exc)) from exc
|
||||
last_err = exc
|
||||
continue # 首 token 前失败 -> failover
|
||||
raise UpstreamError(f"上游均不可用: {type(last_err).__name__}: {last_err}")
|
||||
|
||||
Reference in New Issue
Block a user