diff --git a/gateway/model_pool.py b/gateway/model_pool.py index 9c73b75..c5508b2 100644 --- a/gateway/model_pool.py +++ b/gateway/model_pool.py @@ -21,11 +21,13 @@ _POOL_PATH = Path(__file__).resolve().parent.parent / "config" / "model_pool.jso TIERS = ("local", "budget", "premium") BACKENDS = ("mock", "llama_server", "openai") ROLES = ("architect", "worker", "agent") +PROVIDERS = ("deepseek", "openai", "anthropic") # 代理层 usage 归一化用(D-P3) # 池条目允许的字段(其余字段拒绝写入) ENTRY_FIELDS = { "id", "name", "tier", "backend", "base_url", "model", "api_key", "price_in", "price_out", "temperature", "max_tokens", "enabled", + "provider", "in_hit_price", # 代理层扩展(D-P3):usage 归一化 / 按命中价选上游 } # 单价默认值($/1M tokens);local 档为 0 @@ -177,6 +179,8 @@ class PoolStore: backend = entry.get("backend", "openai") if backend not in BACKENDS: raise PoolError(f"backend 必须是 {BACKENDS} 之一") + if entry.get("provider") and entry["provider"] not in PROVIDERS: + raise PoolError(f"provider 必须是 {PROVIDERS} 之一") tier = entry.get("tier", "budget") if tier not in TIERS: raise PoolError(f"tier 必须是 {TIERS} 之一") @@ -212,6 +216,10 @@ class PoolStore: "temperature": temperature, "max_tokens": max_tokens, "enabled": bool(entry.get("enabled", True)), + # 代理层扩展(D-P3):provider 缺省 openai;in_hit_price 缺省 = price_in × 1/30 + "provider": (str(entry["provider"]) if entry.get("provider") else "openai"), + "in_hit_price": (float(entry["in_hit_price"]) if entry.get("in_hit_price") is not None + else round(price_in / 30.0, 6)), } @staticmethod diff --git a/gateway/proxy/upstream.py b/gateway/proxy/upstream.py index c6370cc..3ce91f3 100644 --- a/gateway/proxy/upstream.py +++ b/gateway/proxy/upstream.py @@ -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}") diff --git a/tests/test_proxy_upstream.py b/tests/test_proxy_upstream.py new file mode 100644 index 0000000..591ad20 --- /dev/null +++ b/tests/test_proxy_upstream.py @@ -0,0 +1,206 @@ +"""上游客户端测试(T-P2):SSE 透传 / usage 注入过滤 / 三家归一化 / failover / aborted。""" +import asyncio +import json + +import httpx +import pytest + +from gateway.proxy.errors import UpstreamError +from gateway.proxy.upstream import ( + UpstreamAborted, + filter_usage_chunk, + normalize_usage, + stream, +) + + +def asyncio_run(coro): + return asyncio.run(coremaybe(coro)) + + +async def coremaybe(coro): + return await coro + + +def _sse(lines) -> bytes: + return ("\n\n".join(lines) + "\n\n").encode("utf-8") + + +def test_stream_passthrough_and_usage_sink(): + """SSE 逐块 yield;usage chunk 进 sink;始终注入 include_usage。""" + body = _sse([ + 'data: {"choices":[{"delta":{"content":"你"}}]}', + 'data: {"choices":[{"delta":{"content":"好"}}]}', + 'data: {"choices":[{"delta":{},"finish_reason":"stop"}],' + '"usage":{"prompt_tokens":100,"prompt_cache_hit_tokens":60,"completion_tokens":20}}', + "data: [DONE]", + ]) + seen_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen_bodies.append(json.loads(request.read())) + return httpx.Response(200, content=body) + + transport = httpx.MockTransport(handler) + import gateway.proxy.upstream as up + client = httpx.AsyncClient(transport=transport) + orig = up._client + up._client = client + try: + entry = {"base_url": "http://up", "api_key": "k", "provider": "deepseek"} + + async def run(): + sink = {} + chunks = [] + async for b in stream({"model": "deepseek-chat", "messages": []}, entry, sink): + chunks.append(b) + return sink, chunks + + sink, chunks = asyncio_run(run()) + finally: + up._client = orig + asyncio.get_event_loop_policy() + client.__asyncio_run = None + try: + asyncio.get_event_loop().run_until_complete(client.aclose()) + except Exception: + pass + text = b"".join(chunks).decode("utf-8") + assert "\u4f60" in text and "\u597d" in text and "data: [DONE]" in text + assert sink["usage"] == {"in_miss": 40, "in_hit": 60, "out": 20} + assert isinstance(sink["ttfb_ms"], int) + assert seen_bodies[0]["stream_options"] == {"include_usage": True} + + +def test_failover_before_first_token(): + """首条目连不上 -> 切换备用 -> 正常输出。""" + ok_body = _sse(['data: {"choices":[{"delta":{"content":"ok"}}]}', "data: [DONE]"]) + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.host == "bad": + raise httpx.ConnectError("拒连") + return httpx.Response(200, content=ok_body) + + transport = httpx.MockTransport(handler) + import gateway.proxy.upstream as up + client = httpx.AsyncClient(transport=transport) + orig = up._client + up._client = client + try: + async def run(): + sink = {} + return [b async for b in stream( + {"model": "m", "messages": []}, + {"base_url": "http://bad", "api_key": "", "provider": "openai"}, + sink, [{"base_url": "http://good", "api_key": "", "provider": "openai"}])] + + chunks = asyncio_run(run()) + assert b"ok" in b"".join(chunks) + finally: + up._client = orig + try: + asyncio.run(client.aclose()) + except Exception: + pass + + +def test_abort_after_first_token_no_failover(): + """首 token 已 yield 后断流 -> UpstreamAborted(不切换,D-P4)。""" + class StreamResp: + def raise_for_status(self): + pass + + async def aiter_lines(self): + yield 'data: {"choices":[{"delta":{"content":"前半"}}]}' + raise httpx.ReadError("中途断") + + class CM: + async def __aenter__(self): + return StreamResp() + + async def __aexit__(self, *a): + return False + + class FakeClient: + def stream(self, *a, **k): + return CM() + + import gateway.proxy.upstream as up + orig = up._client + up._client = FakeClient() + try: + async def run(): + sink = {} + gen = stream({"model": "m", "messages": []}, + {"base_url": "http://x", "api_key": "", "provider": "openai"}, sink) + first = await gen.__anext__() + assert "\u524d\u534a".encode("utf-8") in first # "前半" + rest = [] + async for b in gen: + rest.append(b) + return rest + + try: + asyncio_run(run()) + raised = False + except UpstreamAborted: + raised = True + assert raised, "应抛 UpstreamAborted" + finally: + up._client = orig + + +def test_all_candidates_fail_raises_upstream_error(): + """全部候选首 token 前失败 -> UpstreamError(502)。""" + import gateway.proxy.upstream as up + + class FakeClient: + def stream(self, *a, **k): + raise httpx.ConnectError("全挂") + + orig = up._client + up._client = FakeClient() + try: + async def run(): + async for _ in stream({"model": "m", "messages": []}, + {"base_url": "http://a", "api_key": "", "provider": "openai"}, + {}, [{"base_url": "http://b", "api_key": "", + "provider": "openai"}]): + pass + + try: + asyncio_run(run()) + raised = False + except UpstreamError: + raised = True + assert raised + finally: + up._client = orig + + +def test_normalize_usage_three_providers(): + """deepseek / openai 兼容 / anthropic 三家字段归一;命中数钳制。""" + assert normalize_usage("deepseek", { + "prompt_tokens": 100, "prompt_cache_hit_tokens": 60, "completion_tokens": 20 + }) == {"in_miss": 40, "in_hit": 60, "out": 20} + assert normalize_usage("openai", { + "prompt_tokens": 100, + "prompt_tokens_details": {"cached_tokens": 30}, + "completion_tokens": 20, + }) == {"in_miss": 70, "in_hit": 30, "out": 20} + assert normalize_usage("anthropic", { + "input_tokens": 100, "cache_read_input_tokens": 90, "output_tokens": 5 + }) == {"in_miss": 10, "in_hit": 90, "out": 5} + assert normalize_usage("deepseek", { + "prompt_tokens": 10, "prompt_cache_hit_tokens": 99, "completion_tokens": 1 + })["in_miss"] == 0 + + +def test_filter_usage_chunk(): + """客户端未要求 usage 时剥除 usage chunk;要求时透传;[DONE] 恒透传。""" + usage_line = 'data: {"choices":[],"usage":{"prompt_tokens":1}}' + normal_line = 'data: {"choices":[{"delta":{"content":"x"}}]}' + assert filter_usage_chunk(usage_line, client_wants_usage=False) is None + assert filter_usage_chunk(usage_line, client_wants_usage=True) == usage_line + assert filter_usage_chunk(normal_line, client_wants_usage=False) == normal_line + assert filter_usage_chunk("data: [DONE]", False) == "data: [DONE]" diff --git a/任务拆解与执行计划.md b/任务拆解与执行计划.md index 0ee0c68..403e820 100644 --- a/任务拆解与执行计划.md +++ b/任务拆解与执行计划.md @@ -137,7 +137,7 @@ P0 完成后的能力:干净的后端抽象 + 可量化的评测 + 可追溯 |---|------|------|--------| | T-P0 | 骨架:gateway/proxy/ 包 + SQLite DDL + enabled 门控挂路由 | ✅ 完成 | T-P0 | | T-P1 | 鉴权+账本:key 签发/令牌桶/四表/request_id 幂等/日限额 | ✅ 完成 | T-P1 | -| T-P2 | 上游客户端:流式派发+三家 usage 归一化+首 token 前 failover | ⬜ 待办 | | +| T-P2 | 上游客户端:流式派发+三家 usage 归一化+首 token 前 failover | ✅ 完成 | T-P2 | | T-P3 | 计价+结算:峰谷窗口/毫元整数/黄金用例 ≥10 组 | ⬜ 待办 | | | T-P4 | 路由端到端:/proxy/v1 非流式+流式+402/429/413 语义 | ⬜ 待办 | | | T-P5 | 规范化+桶:稳定哈希/整形顺序/doc_version 失效/多轮不缓存 | ⬜ 待办 | |