projectAIpopular/gateway/proxy/upstream.py

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


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}。

    - 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}")