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:
@@ -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]"
|
||||
Reference in New Issue
Block a user