feat: add openai-compatible transport with sse parsing and error translation
This commit is contained in:
@@ -0,0 +1,255 @@
|
||||
"""OpenAICompatTransport 测试(M1 设计 §4.4/§6):SSE 解析、错误翻译、缺 DONE 语义。
|
||||
|
||||
SSE 帧样本按三项目真实网关响应形态二次构造(OpenAI 兼容 chunk 结构)。
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from polygateway.errors import (
|
||||
RequestRejectedError,
|
||||
SourceDeadError,
|
||||
TransientError,
|
||||
)
|
||||
from polygateway.transports.openai_compat import (
|
||||
OpenAICompatTransport,
|
||||
_iter_sse_deltas,
|
||||
_sse_data_payload,
|
||||
)
|
||||
from polygateway.types import SourceConfig
|
||||
|
||||
|
||||
def _source(**overrides):
|
||||
base = dict(
|
||||
name="qwen_1", provider="qwen", base_url="https://gw.example/v1",
|
||||
api_key="sk-test", model="qwen-max", timeout_s=5.0,
|
||||
)
|
||||
base.update(overrides)
|
||||
return SourceConfig(**base)
|
||||
|
||||
|
||||
def _chunk(content=None, reasoning=None, usage=None):
|
||||
delta = {}
|
||||
if content is not None:
|
||||
delta["content"] = content
|
||||
if reasoning is not None:
|
||||
delta["reasoning_content"] = reasoning
|
||||
body = {"choices": [{"delta": delta}]} if (delta or usage is None) else {"choices": []}
|
||||
if usage is not None:
|
||||
body["usage"] = usage
|
||||
return f"data: {json.dumps(body)}\n\n"
|
||||
|
||||
|
||||
_USAGE = {"prompt_tokens": 11, "completion_tokens": 7}
|
||||
|
||||
|
||||
def _sse_stream(*frames, done=True):
|
||||
text = "".join(frames) + ("data: [DONE]\n\n" if done else "")
|
||||
return httpx.Response(
|
||||
200, content=text.encode(), headers={"content-type": "text/event-stream"}
|
||||
)
|
||||
|
||||
|
||||
def _transport_for(handler):
|
||||
mock = httpx.MockTransport(handler)
|
||||
return OpenAICompatTransport(
|
||||
client_factory=lambda src: httpx.AsyncClient(base_url=src.base_url, transport=mock)
|
||||
)
|
||||
|
||||
|
||||
async def _complete(transport, source, *, stream=True, overlay=None):
|
||||
return await transport.complete(
|
||||
messages=[{"role": "user", "content": "hi"}], source=source,
|
||||
stream=stream, overlay=overlay or {}, call_id="cid-1",
|
||||
)
|
||||
|
||||
|
||||
class TestSsePureFunctions:
|
||||
def test_data_payload_filters_noise(self):
|
||||
assert _sse_data_payload("") is None
|
||||
assert _sse_data_payload(": ping") is None
|
||||
assert _sse_data_payload("event: x") is None
|
||||
assert _sse_data_payload("data: {}") == "{}"
|
||||
assert _sse_data_payload("data: [DONE]") == "[DONE]"
|
||||
|
||||
async def test_iter_deltas_yields_and_flags_done(self):
|
||||
async def lines():
|
||||
for raw in [_chunk(content="he"), ": ping", _chunk(reasoning="think"),
|
||||
_chunk(usage=_USAGE), "data: [DONE]"]:
|
||||
for line in raw.splitlines():
|
||||
yield line
|
||||
|
||||
sink = {}
|
||||
out = [d async for d in _iter_sse_deltas(lines(), sink)]
|
||||
assert out == [(True, "he"), (False, "think")]
|
||||
assert sink["done"] is True and sink["usage"] == _USAGE
|
||||
|
||||
async def test_malformed_frame_raises_transient(self):
|
||||
async def lines():
|
||||
yield "data: {not-json"
|
||||
|
||||
with pytest.raises(TransientError, match="malformed"):
|
||||
_ = [d async for d in _iter_sse_deltas(lines(), {})]
|
||||
|
||||
|
||||
class TestStreamHappyPath:
|
||||
async def test_full_stream_with_usage(self):
|
||||
def handler(request):
|
||||
return _sse_stream(_chunk(reasoning="ponder"), _chunk(content="hello"),
|
||||
_chunk(content=" world"), _chunk(usage=_USAGE))
|
||||
|
||||
result = await _complete(_transport_for(handler), _source())
|
||||
assert result.content == "hello world"
|
||||
assert result.thinking == "ponder"
|
||||
assert result.prompt_tokens == 11 and result.completion_tokens == 7
|
||||
assert result.usage_source == "measured"
|
||||
assert result.ttft_ms is not None and result.ttft_ms >= 0
|
||||
|
||||
async def test_qwen_think_tag_stripped(self):
|
||||
def handler(request):
|
||||
return _sse_stream(_chunk(content="<think>hmm</think>answer"), _chunk(usage=_USAGE))
|
||||
|
||||
result = await _complete(_transport_for(handler), _source())
|
||||
assert result.content == "answer"
|
||||
assert result.thinking == "hmm"
|
||||
|
||||
async def test_usage_missing_falls_back_to_est(self):
|
||||
def handler(request):
|
||||
return _sse_stream(_chunk(content="ok"))
|
||||
|
||||
result = await _complete(_transport_for(handler), _source(tpm=1000, est_tokens=333))
|
||||
assert result.usage_source == "estimated"
|
||||
assert result.prompt_tokens == 0 and result.completion_tokens == 333
|
||||
|
||||
|
||||
class TestMissingDoneSemantics:
|
||||
def _no_done_handler(self, request):
|
||||
return _sse_stream(_chunk(content="partial"), _chunk(usage=_USAGE), done=False)
|
||||
|
||||
async def test_default_retry_policy_raises_transient(self):
|
||||
with pytest.raises(TransientError, match="missing_done|truncated"):
|
||||
await _complete(_transport_for(self._no_done_handler), _source())
|
||||
|
||||
async def test_salvage_policy_keeps_content_as_estimated(self):
|
||||
result = await _complete(
|
||||
_transport_for(self._no_done_handler), _source(missing_done="salvage")
|
||||
)
|
||||
assert result.content == "partial"
|
||||
assert result.usage_source == "estimated" # 打捞路径强制 estimated
|
||||
|
||||
async def test_early_eof_always_transient_even_under_salvage(self):
|
||||
def handler(request):
|
||||
return _sse_stream(done=False) # 零内容提前断流
|
||||
|
||||
with pytest.raises(TransientError, match="early_eof"):
|
||||
await _complete(_transport_for(handler), _source(missing_done="salvage"))
|
||||
|
||||
|
||||
class TestNonStreamFastPath:
|
||||
async def test_non_stream_parses_message(self):
|
||||
def handler(request):
|
||||
body = json.loads(request.content)
|
||||
assert body.get("stream") is False and "stream_options" not in body
|
||||
return httpx.Response(200, json={
|
||||
"choices": [{"message": {"content": "42", "reasoning_content": "count"}}],
|
||||
"usage": _USAGE,
|
||||
})
|
||||
|
||||
result = await _complete(_transport_for(handler), _source(), stream=False)
|
||||
assert result.content == "42" and result.thinking == "count"
|
||||
assert result.usage_source == "measured"
|
||||
assert result.ttft_ms is None
|
||||
|
||||
|
||||
class TestRequestShaping:
|
||||
@pytest.mark.parametrize(
|
||||
("enable_thinking", "expected"),
|
||||
[(True, {"enable_thinking": True}), (False, {"enable_thinking": False}), (None, {})],
|
||||
)
|
||||
async def test_thinking_tri_state_injection(self, enable_thinking, expected):
|
||||
seen = {}
|
||||
|
||||
def handler(request):
|
||||
seen.update(json.loads(request.content))
|
||||
return _sse_stream(_chunk(content="x"), _chunk(usage=_USAGE))
|
||||
|
||||
await _complete(_transport_for(handler), _source(enable_thinking=enable_thinking))
|
||||
assert {k: seen[k] for k in expected} == expected
|
||||
if enable_thinking is None:
|
||||
assert "enable_thinking" not in seen
|
||||
assert seen["stream_options"] == {"include_usage": True}
|
||||
|
||||
async def test_overlay_merged_into_payload(self):
|
||||
seen = {}
|
||||
|
||||
def handler(request):
|
||||
seen.update(json.loads(request.content))
|
||||
return _sse_stream(_chunk(content="x"), _chunk(usage=_USAGE))
|
||||
|
||||
await _complete(
|
||||
_transport_for(handler), _source(),
|
||||
overlay={"response_format": {"type": "json_object"}},
|
||||
)
|
||||
assert seen["response_format"] == {"type": "json_object"}
|
||||
|
||||
|
||||
class TestErrorTranslation:
|
||||
@pytest.mark.parametrize(
|
||||
("status", "body", "headers", "exc", "match"),
|
||||
[
|
||||
(401, "{}", {}, SourceDeadError, None),
|
||||
(403, "{}", {}, SourceDeadError, None),
|
||||
(400, "{}", {}, RequestRejectedError, None),
|
||||
(418, "{}", {}, RequestRejectedError, None),
|
||||
(500, "{}", {}, TransientError, None),
|
||||
(503, "{}", {}, TransientError, None),
|
||||
(429, json.dumps({"error": {"type": "insufficient_quota"}}), {}, SourceDeadError, None),
|
||||
(429, "{}", {"retry-after": "2.5"}, TransientError, None),
|
||||
],
|
||||
)
|
||||
async def test_status_translation(self, status, body, headers, exc, match):
|
||||
def handler(request):
|
||||
return httpx.Response(status, content=body.encode(), headers=headers)
|
||||
|
||||
with pytest.raises(exc) as ei:
|
||||
await _complete(_transport_for(handler), _source())
|
||||
assert ei.value.status_code == status
|
||||
assert ei.value.source_name == "qwen_1"
|
||||
if status == 429 and exc is TransientError:
|
||||
assert ei.value.retry_after_s == 2.5
|
||||
|
||||
async def test_retry_after_http_date_ignored(self):
|
||||
def handler(request):
|
||||
return httpx.Response(429, content=b"{}",
|
||||
headers={"retry-after": "Wed, 21 Oct 2026 07:28:00 GMT"})
|
||||
|
||||
with pytest.raises(TransientError) as ei:
|
||||
await _complete(_transport_for(handler), _source())
|
||||
assert ei.value.retry_after_s is None # 仅支持秒数形态(CHS 同款)
|
||||
|
||||
async def test_httpx_timeout_becomes_transient(self):
|
||||
def handler(request):
|
||||
raise httpx.ConnectTimeout("boom")
|
||||
|
||||
with pytest.raises(TransientError):
|
||||
await _complete(_transport_for(handler), _source())
|
||||
|
||||
async def test_httpx_transport_error_becomes_transient(self):
|
||||
def handler(request):
|
||||
raise httpx.RemoteProtocolError("broken pipe")
|
||||
|
||||
with pytest.raises(TransientError):
|
||||
await _complete(_transport_for(handler), _source())
|
||||
|
||||
|
||||
class TestLifecycle:
|
||||
async def test_aclose_idempotent(self):
|
||||
def handler(request):
|
||||
return _sse_stream(_chunk(content="x"), _chunk(usage=_USAGE))
|
||||
|
||||
transport = _transport_for(handler)
|
||||
await _complete(transport, _source())
|
||||
await transport.aclose()
|
||||
await transport.aclose()
|
||||
Reference in New Issue
Block a user