"""上游客户端（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}")