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:
@@ -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
|
||||
|
||||
+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}")
|
||||
|
||||
@@ -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]"
|
||||
+1
-1
@@ -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 失效/多轮不缓存 | ⬜ 待办 | |
|
||||
|
||||
Reference in New Issue
Block a user