iomgaa/PolyGateway
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases 9 Wiki Activity

feat: add gateway client with env-driven assembly

Browse Source 
Includes config aggregation for multi-source env keys, from_env and
from_settings factories with explicit shared-backend injection,
gather_bounded, top-level exports, tightened import-linter layers with
the gate removed from the Makefile, and the finalized .env.example.
This commit is contained in:
iomgaa
2026-07-20 07:47:05 -04:00
parent 936895919c
commit 7b9815f4bc
30 changed files with 1701 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/test_cache.py
+42 -16
View File
@@ -15,12 +15,21 @@ _MSGS = [{"role": "user", "content": "hi"}]
def _resp(content="cached", **overrides):
base = dict(
content=content, thinking="", model="m", provider="p",
prompt_tokens=1, completion_tokens=2, latency_ms=30,
ttft_ms=5.0, max_inter_token_ms=2.0, cache_hit=False, call_id="orig",
source_name="s1", usage_source="measured",
)
base = {
"content": content,
"thinking": "",
"model": "m",
"provider": "p",
"prompt_tokens": 1,
"completion_tokens": 2,
"latency_ms": 30,
"ttft_ms": 5.0,
"max_inter_token_ms": 2.0,
"cache_hit": False,
"call_id": "orig",
"source_name": "s1",
"usage_source": "measured",
}
base.update(overrides)
return LLMResponse(**base)
@@ -46,21 +55,33 @@ class TestKeyFormula:
def test_multimodal_part_digested_not_inlined(self):
big_b64 = "data:image/png;base64," + "A" * 1_000_000
messages = [{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": big_b64}},
{"type": "text", "text": "describe"},
]}]
messages = [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": big_b64}},
{"type": "text", "text": "describe"},
],
}
]
digested = digest_messages(messages)
payload = json.dumps(digested, ensure_ascii=False)
assert len(payload) < 500 # 大图不进 canonical_json
expected = hashlib.sha256(big_b64.encode()).hexdigest()
assert expected in payload # 但字节变化仍改变 key
# 图像字节变化 → key 变
messages2 = [{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": big_b64[:-1] + "B"}},
{"type": "text", "text": "describe"},
]}]
assert build_cache_key("m", messages, "p", None) != build_cache_key("m", messages2, "p", None)
messages2 = [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": big_b64[:-1] + "B"}},
{"type": "text", "text": "describe"},
],
}
]
assert build_cache_key("m", messages, "p", None) != build_cache_key(
"m", messages2, "p", None
)
class _Terminal:
@@ -76,7 +97,12 @@ class _Terminal:
def _mw(backend, **kwargs):
defaults = dict(backend=backend, model_fingerprint="m", default_namespace="proj", ttl_s=3600)
defaults = {
"backend": backend,
"model_fingerprint": "m",
"default_namespace": "proj",
"ttl_s": 3600,
}
defaults.update(kwargs)
return CacheMW(**defaults)
tests/unit/test_client.py
+256
View File
@@ -0,0 +1,256 @@
"""GatewayClient 装配与端到端(fake 后端 + MockTransport)测试(设计 §2.4)。"""
import asyncio
import json
import sys
from pathlib import Path
import httpx
import pytest
from polygateway import (
AllSourcesExhausted,
GatewayClient,
GatewaySettings,
gather_bounded,
)
from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.cache import InMemoryCache
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.sources import RoundRobinSelector
from polygateway.structured.json_repair import JsonRepairStrategy
from polygateway.transports.openai_compat import OpenAICompatTransport
from polygateway.types import (
BackpressurePolicy,
BreakerConfig,
GlobalLimits,
RetryPolicy,
SourceConfig,
)
_REPO = Path(__file__).resolve().parents[2]
_ENV = {
"LLM__QWEN__1__BASE_URL": "https://gw.example/v1",
"LLM__QWEN__1__API_KEY": "sk-a",
"LLM__QWEN__1__MODEL": "qwen-max",
"LLM__QWEN__1__TIMEOUT_S": "120",
"LLM_MAX_RETRIES": "3",
"LLM_RETRY_BASE_DELAY": "2.0",
"LLM_RETRY_MAX_DELAY": "30.0",
"LLM_CIRCUIT_BREAKER_THRESHOLD": "5",
"LLM_CIRCUIT_BREAKER_COOLDOWN": "60",
"PGW_CACHE_BACKEND": "none",
"PGW_TELEMETRY_BACKEND": "none",
}
def _sse(content='{"answer": 1}'):
chunk = json.dumps({"choices": [{"delta": {"content": content}}]})
usage = json.dumps({"choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 4}})
body = f"data: {chunk}\n\ndata: {usage}\n\ndata: [DONE]\n\n"
return httpx.Response(200, content=body.encode(), headers={"content-type": "text/event-stream"})
def _source(name="qwen_1", **overrides):
base = {
"name": name,
"provider": "qwen",
"base_url": "https://gw.example/v1",
"api_key": "sk",
"model": "qwen-max",
"timeout_s": 10.0,
}
base.update(overrides)
return SourceConfig(**base)
def _client(sources=None, handler=None, *, limiter=None, quota_full="wait", **overrides):
sources = sources or [_source()]
handler = handler or (lambda request: _sse())
transport = OpenAICompatTransport(
client_factory=lambda src: httpx.AsyncClient(transport=httpx.MockTransport(handler))
)
defaults = {
"scope": "llm",
"sources": sources,
"selector": RoundRobinSelector(),
"limiter": limiter
or InMemoryLimiter(
scope="llm",
sources={s.name: s for s in sources},
global_limits=GlobalLimits(0, 0, 0),
),
"gate": InMemoryGate(config=BreakerConfig(5, 60.0, 120.0)),
"transport": transport,
"retry": RetryPolicy(3, 2.0, 30.0),
"backpressure": BackpressurePolicy(300.0, 0.01),
"quota_full": quota_full,
"structured_strategy": JsonRepairStrategy(),
}
defaults.update(overrides)
return GatewayClient(**defaults)
class TestChatEndToEnd:
async def test_plain_chat(self):
async with _client() as client:
resp = await client.chat([{"role": "user", "content": "hi"}])
assert resp.content == '{"answer": 1}'
assert resp.source_name == "qwen_1" and resp.prompt_tokens == 3
async def test_structured_json_tier(self):
async with _client() as client:
resp = await client.chat([{"role": "user", "content": "hi"}], structured="json")
assert resp.structured_data == {"answer": 1}
async def test_cache_hit_roundtrip(self):
client = _client(cache=InMemoryCache(), cache_namespace="proj", cache_ttl_s=3600)
async with client:
first = await client.chat([{"role": "user", "content": "hi"}])
second = await client.chat([{"role": "user", "content": "hi"}])
assert first.cache_hit is False and second.cache_hit is True
async def test_structured_unavailable_fails_loudly(self):
async with _client(structured_strategy=None) as client:
with pytest.raises(ImportError, match="structured"):
await client.chat([{"role": "user", "content": "hi"}], structured="json")
class TestFactories:
def test_from_env_assembles(self):
client = GatewayClient.from_env("LLM", env=_ENV)
assert isinstance(client, GatewayClient)
def test_from_env_unknown_provider_fails_at_assembly(self):
env = dict(_ENV)
for key in list(env):
if key.startswith("LLM__QWEN__"):
env[key.replace("QWEN", "GLM")] = env.pop(key)
with pytest.raises(ValueError, match="glm"):
GatewayClient.from_env("LLM", env=env)
def test_from_settings_respects_injection(self):
settings = GatewaySettings.from_env("LLM", env=_ENV)
shared = InMemoryLimiter(
scope="shared",
sources={s.name: s for s in settings.sources},
global_limits=GlobalLimits(0, 0, 0),
)
client = GatewayClient.from_settings(settings, limiter=shared)
assert isinstance(client, GatewayClient)
class TestSharedBackend:
async def test_two_clients_share_global_concurrency_gate(self):
"""VT R5: 两个逻辑角色显式注入同一 limiter → 共享全局并发闸。"""
src_a, src_b = _source("role_a_1"), _source("role_b_1")
shared = InMemoryLimiter(
scope="shared",
sources={"role_a_1": src_a, "role_b_1": src_b}, # 共享后端持源并集
global_limits=GlobalLimits(max_concurrency=1, rpm=0, tpm=0),
)
started = asyncio.Event()
async def slow_handler(request):
started.set()
await asyncio.sleep(0.2)
return _sse()
client_a = _client([src_a], slow_handler, limiter=shared)
client_b = _client([src_b], limiter=shared, quota_full="fail_fast")
task = asyncio.ensure_future(client_a.chat([{"role": "user", "content": "x"}]))
await started.wait()
with pytest.raises(AllSourcesExhausted) as ei:
await client_b.chat([{"role": "user", "content": "y"}])
assert ei.value.reason == "quota_exhausted" # 全局闸被 A 占满 → B 立即失败
await task
async def test_aclose_idempotent(self):
client = _client()
await client.aclose()
await client.aclose()
class TestGatherBounded:
async def test_order_preserved_and_concurrency_capped(self):
peak = {"now": 0, "max": 0}
async def work(i):
peak["now"] += 1
peak["max"] = max(peak["max"], peak["now"])
await asyncio.sleep(0.01)
peak["now"] -= 1
return i
results = await gather_bounded((work(i) for i in range(10)), concurrency=3)
assert results == list(range(10))
assert peak["max"] <= 3
async def test_exception_propagates(self):
async def boom():
raise RuntimeError("x")
async def ok():
return 1
with pytest.raises(RuntimeError):
await gather_bounded([ok(), boom()], concurrency=2)
async def test_invalid_concurrency(self):
with pytest.raises(ValueError):
await gather_bounded([], concurrency=0)
def _load_reference_protocol(insert_path: str, module: str):
sys.path.insert(0, str(_REPO / insert_path))
try:
import importlib
return importlib.import_module(module).LLMProvider
finally:
sys.path.pop(0)
class TestReferenceProtocolCompat:
"""结构兼容断言(只读 import reference;失败即公共承诺破裂)。"""
def test_satisfies_govdoc_llm_provider(self):
try:
proto = _load_reference_protocol(
"reference/GovDoc-SaaS/packages/docagent-core/src", "docagent_core.protocols"
)
except ImportError:
# 兜底: 按 protocols.py:15-25 逐字复制的结构断言
from typing import Any, Protocol, runtime_checkable
@runtime_checkable
class proto(Protocol): # noqa: N801 — 复制自 GovDoc protocols.py:15-25
async def chat(
self,
messages: list[dict[str, Any]],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
): ...
assert isinstance(_client(), proto)
def test_satisfies_videotree_llm_provider(self):
try:
proto = _load_reference_protocol("reference/Video-Tree-TRM5", "core.protocols")
except ImportError:
from typing import Any, Protocol, runtime_checkable
@runtime_checkable
class proto(Protocol): # noqa: N801 — 复制自 VT core/protocols.py:18-29
async def chat(
self,
messages: list[dict[str, Any]],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
cache_salt: str | None = None,
): ...
assert isinstance(_client(), proto)
tests/unit/test_config.py
+166
View File
@@ -0,0 +1,166 @@
"""config.py 配置聚合测试(设计 §8): 多源命名、键优先级、缺失报错。"""
import pytest
from polygateway.config import GatewaySettings
_BASE_ENV = {
"LLM__QWEN__1__BASE_URL": "https://gw-a.example/v1",
"LLM__QWEN__1__API_KEY": "sk-a",
"LLM__QWEN__1__MODEL": "qwen-max",
"LLM__QWEN__1__TIMEOUT_S": "120",
"LLM_MAX_RETRIES": "3",
"LLM_RETRY_BASE_DELAY": "2.0",
"LLM_RETRY_MAX_DELAY": "30.0",
"LLM_CIRCUIT_BREAKER_THRESHOLD": "5",
"LLM_CIRCUIT_BREAKER_COOLDOWN": "60",
"PGW_CACHE_BACKEND": "none",
"PGW_TELEMETRY_BACKEND": "none",
}
def _env(**overrides):
env = dict(_BASE_ENV)
env.update({k: v for k, v in overrides.items() if v is not None})
for k, v in overrides.items():
if v is None:
env.pop(k, None)
return env
class TestSourceAggregation:
def test_single_source_parsed(self):
s = GatewaySettings.from_env("LLM", env=_env())
assert len(s.sources) == 1
src = s.sources[0]
assert src.name == "qwen_1" and src.provider == "qwen"
assert src.base_url == "https://gw-a.example/v1" and src.timeout_s == 120.0
def test_multi_source_and_optional_fields(self):
env = _env(
**{
"LLM__DEEPSEEK__2__BASE_URL": "https://gw-b.example/v1",
"LLM__DEEPSEEK__2__API_KEY": "sk-b",
"LLM__DEEPSEEK__2__MODEL": "deepseek-chat",
"LLM__DEEPSEEK__2__TIMEOUT_S": "90",
"LLM__DEEPSEEK__2__RPM": "60",
"LLM__DEEPSEEK__2__ENABLE_THINKING": "true",
"LLM__DEEPSEEK__2__MISSING_DONE": "salvage",
}
)
s = GatewaySettings.from_env("LLM", env=env)
by_name = {src.name: src for src in s.sources}
assert set(by_name) == {"qwen_1", "deepseek_2"}
ds = by_name["deepseek_2"]
assert ds.rpm == 60 and ds.enable_thinking is True and ds.missing_done == "salvage"
assert by_name["qwen_1"].enable_thinking is None # 未配置 = 三态 None
def test_other_scope_keys_ignored(self):
env = _env(
**{
"OCR__MONKEY__1__BASE_URL": "http://lan/parse",
"OCR__MONKEY__1__API_KEY": "x",
"OCR__MONKEY__1__MODEL": "monkey",
"OCR__MONKEY__1__TIMEOUT_S": "60",
}
)
s = GatewaySettings.from_env("LLM", env=env)
assert len(s.sources) == 1
def test_flat_timeout_is_source_default(self):
env = _env(LLM_TIMEOUT="300", **{"LLM__QWEN__1__TIMEOUT_S": None})
s = GatewaySettings.from_env("LLM", env=env)
assert s.sources[0].timeout_s == 300.0
@pytest.mark.parametrize("missing", ["BASE_URL", "API_KEY", "MODEL"])
def test_missing_required_source_field_fails(self, missing):
with pytest.raises(ValueError, match=missing):
GatewaySettings.from_env("LLM", env=_env(**{f"LLM__QWEN__1__{missing}": None}))
def test_unknown_field_fails_loudly(self):
with pytest.raises(ValueError, match="TEMPRATURE"):
GatewaySettings.from_env("LLM", env=_env(**{"LLM__QWEN__1__TEMPRATURE": "0.7"}))
def test_no_sources_fails(self):
env = {k: v for k, v in _BASE_ENV.items() if not k.startswith("LLM__")}
with pytest.raises(ValueError, match="源"):
GatewaySettings.from_env("LLM", env=env)
class TestResilienceKeys:
def test_flat_legacy_keys(self):
s = GatewaySettings.from_env("LLM", env=_env())
assert s.retry.max_attempts == 3 and s.retry.backoff_base_s == 2.0
assert s.breaker.fail_threshold == 5 and s.breaker.cooldown_s == 60.0
def test_scope_keys_override_flat(self):
env = _env(**{"LLM__RETRY__MAX_ATTEMPTS": "7", "LLM__BREAKER__COOLDOWN_S": "15"})
s = GatewaySettings.from_env("LLM", env=env)
assert s.retry.max_attempts == 7
assert s.breaker.cooldown_s == 15.0
assert s.breaker.fail_threshold == 5 # 未覆盖的仍取平铺键
def test_missing_retry_config_fails(self):
with pytest.raises(ValueError, match="MAX_RETRIES|MAX_ATTEMPTS"):
GatewaySettings.from_env("LLM", env=_env(LLM_MAX_RETRIES=None))
def test_probe_ttl_derived_when_absent(self):
s = GatewaySettings.from_env("LLM", env=_env())
# 派生规则: max(2 × 最大源 timeout, cooldown)
assert s.breaker.probe_ttl_s == max(2 * 120.0, 60.0)
s2 = GatewaySettings.from_env("LLM", env=_env(**{"LLM__BREAKER__PROBE_TTL_S": "45"}))
assert s2.breaker.probe_ttl_s == 45.0
def test_selector_and_quota_full(self):
s = GatewaySettings.from_env("LLM", env=_env())
assert s.selector == "round_robin" and s.quota_full == "wait"
s2 = GatewaySettings.from_env(
"LLM", env=_env(**{"LLM__SELECTOR": "least_inflight", "LLM__QUOTA_FULL": "fail_fast"})
)
assert s2.selector == "least_inflight" and s2.quota_full == "fail_fast"
with pytest.raises(ValueError):
GatewaySettings.from_env("LLM", env=_env(**{"LLM__SELECTOR": "random"}))
def test_global_limits(self):
env = _env(**{"LLM__GLOBAL__MAX_CONCURRENCY": "8", "LLM__GLOBAL__RPM": "120"})
s = GatewaySettings.from_env("LLM", env=env)
assert s.global_limits.max_concurrency == 8 and s.global_limits.rpm == 120
assert s.global_limits.tpm == 0
class TestAssemblyGuards:
def test_cache_requires_namespace_and_ttl(self):
env = _env(PGW_CACHE_BACKEND="memory")
with pytest.raises(ValueError, match="NAMESPACE"):
GatewaySettings.from_env("LLM", env=env)
env2 = _env(PGW_CACHE_BACKEND="memory", PGW_CACHE_NAMESPACE="proj", PGW_CACHE_TTL_S="0")
with pytest.raises(ValueError, match="TTL"):
GatewaySettings.from_env("LLM", env=env2)
env3 = _env(PGW_CACHE_BACKEND="memory", PGW_CACHE_NAMESPACE="proj", PGW_CACHE_TTL_S="3600")
s = GatewaySettings.from_env("LLM", env=env3)
assert s.cache_namespace == "proj" and s.cache_ttl_s == 3600
def test_redis_cache_requires_url(self):
env = _env(PGW_CACHE_BACKEND="redis", PGW_CACHE_NAMESPACE="proj", PGW_CACHE_TTL_S="3600")
with pytest.raises(ValueError, match="REDIS_URL"):
GatewaySettings.from_env("LLM", env=env)
def test_sqlite_telemetry_requires_path(self):
env = _env(PGW_TELEMETRY_BACKEND="sqlite")
with pytest.raises(ValueError, match="SQLITE_PATH"):
GatewaySettings.from_env("LLM", env=env)
def test_timeout_must_fit_lease_ttl(self):
env = _env(PGW_LEASE_TTL_S="60", **{"LLM__QWEN__1__TIMEOUT_S": "120"})
with pytest.raises(ValueError, match="租约|lease"):
GatewaySettings.from_env("LLM", env=env)
def test_effective_breaker_threshold_auto_raised(self):
env = _env(**{"LLM__QWEN__1__MAX_CONCURRENCY": "8"})
s = GatewaySettings.from_env("LLM", env=env)
# 有效阈值 = max(配置值 5, 并发 8 × 2) = 16(.env 注释约定入库)
assert s.breaker.fail_threshold == 16
def test_m1_only_memory_governance_backends(self):
with pytest.raises(ValueError, match="M2"):
GatewaySettings.from_env("LLM", env=_env(PGW_LIMITER_BACKEND="redis"))
tests/unit/test_errors.py
+9 -3
View File
@@ -34,7 +34,9 @@ class TestBaseShape:
class TestResultInvalid:
def test_carries_diagnosis(self):
exc = ResultInvalidError(
"bad json", raw_text="{oops", repair_error="unterminated",
"bad json",
raw_text="{oops",
repair_error="unterminated",
validation_errors=("field x missing",),
)
assert exc.raw_text == "{oops"
@@ -45,7 +47,9 @@ class TestResultInvalid:
class TestGatewayUnavailable:
def test_fields_and_inheritance(self):
exc = AllSourcesExhausted(
scope="LLM", reason="retry_exhausted", retry_after_s=4.0,
scope="LLM",
reason="retry_exhausted",
retry_after_s=4.0,
per_source_reasons={"qwen_1": "timeout"},
)
assert isinstance(exc, GatewayUnavailableError)
@@ -66,7 +70,9 @@ class TestGatewayUnavailable:
def test_per_source_reason_domain_enforced(self):
with pytest.raises(ValueError):
AllSourcesExhausted(
scope="LLM", reason="no_sources", retry_after_s=0.0,
scope="LLM",
reason="no_sources",
retry_after_s=0.0,
per_source_reasons={"qwen_1": "weird"},
)
tests/unit/test_openai_compat.py
+39 -20
View File
@@ -22,10 +22,14 @@ 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 = {
"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)
@@ -47,9 +51,7 @@ _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"}
)
return httpx.Response(200, content=text.encode(), headers={"content-type": "text/event-stream"})
def _transport_for(handler):
@@ -61,8 +63,11 @@ def _transport_for(handler):
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",
messages=[{"role": "user", "content": "hi"}],
source=source,
stream=stream,
overlay=overlay or {},
call_id="cid-1",
)
@@ -76,8 +81,13 @@ class TestSsePureFunctions:
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 raw in [
_chunk(content="he"),
": ping",
_chunk(reasoning="think"),
_chunk(usage=_USAGE),
"data: [DONE]",
]:
for line in raw.splitlines():
yield line
@@ -97,8 +107,12 @@ class TestSsePureFunctions:
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))
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"
@@ -152,10 +166,13 @@ class TestNonStreamFastPath:
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,
})
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"
@@ -189,7 +206,8 @@ class TestRequestShaping:
return _sse_stream(_chunk(content="x"), _chunk(usage=_USAGE))
await _complete(
_transport_for(handler), _source(),
_transport_for(handler),
_source(),
overlay={"response_format": {"type": "json_object"}},
)
assert seen["response_format"] == {"type": "json_object"}
@@ -222,8 +240,9 @@ class TestErrorTranslation:
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"})
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())
tests/unit/test_ports.py
+77 -25
View File
@@ -31,48 +31,89 @@ class _DummyPermit:
class _DummyLimiter:
async def try_acquire(self, source_key: str, est_tokens: int): return _DummyPermit()
async def acquire(self, source_key: str, est_tokens: int): return _DummyPermit()
async def source_stats(self, source_key: str): return SourceStats(0, 0, 0)
async def try_acquire(self, source_key: str, est_tokens: int):
return _DummyPermit()
async def acquire(self, source_key: str, est_tokens: int):
return _DummyPermit()
async def source_stats(self, source_key: str):
return SourceStats(0, 0, 0)
async def mark_progress(self) -> None: ...
async def progress_age_s(self) -> float: return 0.0
async def progress_age_s(self) -> float:
return 0.0
class _DummyGate:
async def try_enter(self, source_name: str, owner: str): raise NotImplementedError
async def record_success(self, entry): raise NotImplementedError
async def record_failure(self, entry, reason: str, force_open: bool): raise NotImplementedError
async def release_probe(self, entry): raise NotImplementedError
async def retry_after_s(self, sources): return 0.0
async def try_enter(self, source_name: str, owner: str):
raise NotImplementedError
async def record_success(self, entry):
raise NotImplementedError
async def record_failure(self, entry, reason: str, force_open: bool):
raise NotImplementedError
async def release_probe(self, entry):
raise NotImplementedError
async def retry_after_s(self, sources):
return 0.0
class _DummyMw:
async def __call__(self, request, call_next): return await call_next(request)
async def __call__(self, request, call_next):
return await call_next(request)
class _DummyTransport:
async def complete(self, *, messages, source, stream, overlay, call_id): raise NotImplementedError
async def complete(self, *, messages, source, stream, overlay, call_id):
raise NotImplementedError
class _DummyCache:
async def get(self, key: str): return None
async def get(self, key: str):
return None
async def set(self, key: str, value: str, ttl_s: int) -> None: ...
class _DummySelector:
def order(self, sources, stats): return list(sources)
def order(self, sources, stats):
return list(sources)
class _DummyStrategy:
def request_overlay(self, schema): return {}
def parse(self, text: str) -> Any: return {}
def request_overlay(self, schema):
return {}
def parse(self, text: str) -> Any:
return {}
class _DummyRecorder:
async def record_llm_call(
self, *, call_id, parent_call_id, session_id, model, provider, source_name,
messages, response, thinking, prompt_tokens, completion_tokens, usage_source,
latency_ms, ttft_ms, max_inter_token_ms, cache_hit, error, cost,
self,
*,
call_id,
parent_call_id,
session_id,
model,
provider,
source_name,
messages,
response,
thinking,
prompt_tokens,
completion_tokens,
usage_source,
latency_ms,
ttft_ms,
max_inter_token_ms,
cache_hit,
error,
cost,
) -> None: ...
@@ -95,10 +136,15 @@ def test_protocols_are_runtime_checkable(impl, protocol):
def _decision(**overrides) -> GateDecision:
base = dict(
source_name="qwen_1", allowed=True, state=GateState.CLOSED,
epoch=0, is_probe=False, probe_owner=None, retry_after_s=0.0,
)
base = {
"source_name": "qwen_1",
"allowed": True,
"state": GateState.CLOSED,
"epoch": 0,
"is_probe": False,
"probe_owner": None,
"retry_after_s": 0.0,
}
base.update(overrides)
return GateDecision(**base)
@@ -141,9 +187,15 @@ class TestGateDecisionInvariants:
class TestGateUpdate:
def test_bounds(self):
u = GateUpdate(applied=True, state=GateState.CLOSED, epoch=0, failure_count=0, retry_after_s=0.0)
u = GateUpdate(
applied=True, state=GateState.CLOSED, epoch=0, failure_count=0, retry_after_s=0.0
)
assert u.applied
with pytest.raises(ValueError):
GateUpdate(applied=True, state=GateState.CLOSED, epoch=-1, failure_count=0, retry_after_s=0.0)
GateUpdate(
applied=True, state=GateState.CLOSED, epoch=-1, failure_count=0, retry_after_s=0.0
)
with pytest.raises(ValueError):
GateUpdate(applied=True, state=GateState.CLOSED, epoch=0, failure_count=-1, retry_after_s=0.0)
GateUpdate(
applied=True, state=GateState.CLOSED, epoch=0, failure_count=-1, retry_after_s=0.0
)
tests/unit/test_retry.py
+65 -26
View File
@@ -19,7 +19,6 @@ from polygateway.errors import (
TransientError,
)
from polygateway.middleware.retry import RetryMW
from polygateway.ports import GateState
from polygateway.sources import RoundRobinSelector, SourceCooldownMemo
from polygateway.types import (
BackpressurePolicy,
@@ -37,18 +36,28 @@ _NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
def _src(name, **overrides):
base = dict(
name=name, provider="openai", base_url="https://gw.example/v1",
api_key="sk", model="m", timeout_s=10.0,
)
base = {
"name": name,
"provider": "openai",
"base_url": "https://gw.example/v1",
"api_key": "sk",
"model": "m",
"timeout_s": 10.0,
}
base.update(overrides)
return SourceConfig(**base)
def _ok(content="ok"):
return TransportResult(
content=content, thinking="", prompt_tokens=10, completion_tokens=5,
usage_source="measured", ttft_ms=12.0, max_inter_token_ms=3.0, raw={},
content=content,
thinking="",
prompt_tokens=10,
completion_tokens=5,
usage_source="measured",
ttft_ms=12.0,
max_inter_token_ms=3.0,
raw={},
)
@@ -79,23 +88,42 @@ class FakeSleep:
self.delays.append(seconds)
def _harness(sources, script, *, clock=None, max_attempts=3, quota_full="wait",
global_limits=_NO_GLOBAL, rng=lambda: 0.0):
def _harness(
sources,
script,
*,
clock=None,
max_attempts=3,
quota_full="wait",
global_limits=_NO_GLOBAL,
rng=lambda: 0.0,
):
clock = clock or FakeClock()
limiter = InMemoryLimiter(
scope="llm", sources={s.name: s for s in sources},
global_limits=global_limits, lease_ttl_s=100.0, now=clock,
scope="llm",
sources={s.name: s for s in sources},
global_limits=global_limits,
lease_ttl_s=100.0,
now=clock,
)
gate = InMemoryGate(config=_BREAKER, now=clock)
transport = FakeTransport(script)
sleep = FakeSleep()
mw = RetryMW(
scope="llm", sources=sources, selector=RoundRobinSelector(),
limiter=limiter, gate=gate, transport=transport,
scope="llm",
sources=sources,
selector=RoundRobinSelector(),
limiter=limiter,
gate=gate,
transport=transport,
retry=RetryPolicy(max_attempts=max_attempts, backoff_base_s=2.0, backoff_max_s=30.0),
backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.01),
quota_full=quota_full, cooldown_memo=SourceCooldownMemo(now=clock),
emitter=None, now=clock, sleep=sleep, rng=rng,
quota_full=quota_full,
cooldown_memo=SourceCooldownMemo(now=clock),
emitter=None,
now=clock,
sleep=sleep,
rng=rng,
)
return mw, limiter, gate, transport, sleep, clock
@@ -140,7 +168,8 @@ class TestRetryAndFailover:
async def test_max_attempts_is_total_attempts(self):
mw, _, _, transport, _, _ = _harness(
[_src("a")], [TransientError("1"), TransientError("2"), TransientError("3")],
[_src("a")],
[TransientError("1"), TransientError("2"), TransientError("3")],
max_attempts=3,
)
with pytest.raises(AllSourcesExhausted) as ei:
@@ -196,9 +225,7 @@ class TestScopeUnavailable:
async def test_all_sources_circuit_open(self):
clock = FakeClock()
script = [TransientError(str(i)) for i in range(9)]
mw, _, gate, _, _, _ = _harness(
[_src("a")], script, clock=clock, max_attempts=99
)
mw, _, gate, _, _, _ = _harness([_src("a")], script, clock=clock, max_attempts=99)
# 3 次失败后 a 开路 → 第 4 次尝试选不到源且 gate_rejections==全部 → CircuitOpen
with pytest.raises(CircuitOpenError) as ei:
await mw(_REQ)
@@ -225,8 +252,11 @@ class TestScopeUnavailable:
src = _src("a", max_concurrency=1)
clock = FakeClock()
limiter = InMemoryLimiter(
scope="llm", sources={"a": src}, global_limits=_NO_GLOBAL,
lease_ttl_s=100.0, now=clock,
scope="llm",
sources={"a": src},
global_limits=_NO_GLOBAL,
lease_ttl_s=100.0,
now=clock,
)
held = await limiter.try_acquire("a", 0)
released = {"done": False}
@@ -239,12 +269,20 @@ class TestScopeUnavailable:
gate = InMemoryGate(config=_BREAKER, now=clock)
transport = FakeTransport([_ok()])
mw = RetryMW(
scope="llm", sources=[src], selector=RoundRobinSelector(),
limiter=limiter, gate=gate, transport=transport,
scope="llm",
sources=[src],
selector=RoundRobinSelector(),
limiter=limiter,
gate=gate,
transport=transport,
retry=RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0),
backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.01),
quota_full="wait", cooldown_memo=SourceCooldownMemo(now=clock),
emitter=None, now=clock, sleep=sleep_and_release, rng=lambda: 0.0,
quota_full="wait",
cooldown_memo=SourceCooldownMemo(now=clock),
emitter=None,
now=clock,
sleep=sleep_and_release,
rng=lambda: 0.0,
)
resp = await mw(_REQ)
assert resp.content == "ok" and released["done"]
@@ -265,7 +303,8 @@ class TestCancellation:
mw, _, gate, _, _, _ = _harness(
[_src("a")],
[TransientError("1"), TransientError("2"), TransientError("3"), "hang"],
clock=clock, max_attempts=99,
clock=clock,
max_attempts=99,
)
# 三连失败开路
with pytest.raises(CircuitOpenError):
tests/unit/test_sources.py
+6 -2
View File
@@ -6,8 +6,12 @@ from polygateway.types import SourceConfig, SourceStats
def _src(name):
return SourceConfig(
name=name, provider="openai", base_url="https://gw.example/v1",
api_key="sk", model="m", timeout_s=10.0,
name=name,
provider="openai",
base_url="https://gw.example/v1",
api_key="sk",
model="m",
timeout_s=10.0,
)
tests/unit/test_streaming.py
+3 -1
View File
@@ -47,7 +47,9 @@ class TestThreeLayers:
async def test_inter_token_timeout(self):
wrapped = stream_with_liveness_timeouts(
_emit(["a", "b"], delay_s=0.2, first_delay_s=0.0),
ttft_s=1.0, inter_token_s=0.05, total_s=5.0,
ttft_s=1.0,
inter_token_s=0.05,
total_s=5.0,
)
with pytest.raises(StreamLivenessTimeout) as ei:
await _collect(wrapped)
tests/unit/test_structured.py
+18 -8
View File
@@ -21,9 +21,19 @@ class Verdict(BaseModel):
def _resp(content):
return LLMResponse(
content=content, thinking="", model="m", provider="p", prompt_tokens=1,
completion_tokens=2, latency_ms=10, ttft_ms=None, max_inter_token_ms=None,
cache_hit=False, call_id="cid", source_name="s1", usage_source="measured",
content=content,
thinking="",
model="m",
provider="p",
prompt_tokens=1,
completion_tokens=2,
latency_ms=10,
ttft_ms=None,
max_inter_token_ms=None,
cache_hit=False,
call_id="cid",
source_name="s1",
usage_source="measured",
)
@@ -43,10 +53,10 @@ class TestJsonRepairStrategy:
@pytest.mark.parametrize(
"dirty",
[
'```json\n{"answer": 1, "reason": "ok"}\n```', # 围栏
'{"answer": 1, "reason": "ok",}', # 尾逗号
"{'answer': 1, 'reason': 'ok'}", # 单引号
'{"answer": 1, "reason": "ok"', # 缺右括号
'```json\n{"answer": 1, "reason": "ok"}\n```', # 围栏
'{"answer": 1, "reason": "ok",}', # 尾逗号
"{'answer': 1, 'reason': 'ok'}", # 单引号
'{"answer": 1, "reason": "ok"', # 缺右括号
],
)
def test_repairs_real_world_dirt(self, dirty):
@@ -79,7 +89,7 @@ class TestNativeSchemaStrategy:
def _mw(**kwargs):
defaults = dict(strategy=JsonRepairStrategy(), max_retries=1, escalation=None)
defaults = {"strategy": JsonRepairStrategy(), "max_retries": 1, "escalation": None}
defaults.update(kwargs)
return StructuredMW(**defaults)
tests/unit/test_telemetry.py
+109 -42
View File
@@ -15,37 +15,80 @@ from polygateway.types import ChatRequest, LLMResponse, SourceConfig
_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}], session_id="sess-1")
_EXPECTED_COLUMNS = [
"call_id", "parent_call_id", "session_id", "model", "provider", "source_name",
"messages", "response", "thinking", "prompt_tokens", "completion_tokens",
"usage_source", "latency_ms", "ttft_ms", "max_inter_token_ms", "cache_hit",
"error", "cost", "created_at",
"call_id",
"parent_call_id",
"session_id",
"model",
"provider",
"source_name",
"messages",
"response",
"thinking",
"prompt_tokens",
"completion_tokens",
"usage_source",
"latency_ms",
"ttft_ms",
"max_inter_token_ms",
"cache_hit",
"error",
"cost",
"created_at",
]
def _resp(**overrides):
base = dict(
content="ok", thinking="", model="m", provider="p", prompt_tokens=1,
completion_tokens=2, latency_ms=30, ttft_ms=None, max_inter_token_ms=None,
cache_hit=False, call_id="cid-1", source_name="s1", usage_source="measured",
)
base = {
"content": "ok",
"thinking": "",
"model": "m",
"provider": "p",
"prompt_tokens": 1,
"completion_tokens": 2,
"latency_ms": 30,
"ttft_ms": None,
"max_inter_token_ms": None,
"cache_hit": False,
"call_id": "cid-1",
"source_name": "s1",
"usage_source": "measured",
}
base.update(overrides)
return LLMResponse(**base)
def _source():
return SourceConfig(
name="s1", provider="p", base_url="https://gw.example/v1",
api_key="sk", model="m", timeout_s=10.0,
name="s1",
provider="p",
base_url="https://gw.example/v1",
api_key="sk",
model="m",
timeout_s=10.0,
)
async def _record_minimal(recorder, call_id="c1", **overrides):
fields = dict(
call_id=call_id, parent_call_id=None, session_id="sess-1", model="m",
provider="p", source_name="s1", messages="[]", response="ok", thinking="",
prompt_tokens=1, completion_tokens=2, usage_source="measured", latency_ms=10,
ttft_ms=None, max_inter_token_ms=None, cache_hit=False, error=None, cost=None,
)
fields = {
"call_id": call_id,
"parent_call_id": None,
"session_id": "sess-1",
"model": "m",
"provider": "p",
"source_name": "s1",
"messages": "[]",
"response": "ok",
"thinking": "",
"prompt_tokens": 1,
"completion_tokens": 2,
"usage_source": "measured",
"latency_ms": 10,
"ttft_ms": None,
"max_inter_token_ms": None,
"cache_hit": False,
"error": None,
"cost": None,
}
fields.update(overrides)
await recorder.record_llm_call(**fields)
@@ -55,9 +98,9 @@ class TestSQLiteRecorder:
recorder = SQLiteRecorder(tmp_path / "t.db")
await _record_minimal(recorder)
recorder.close()
cols = [r[1] for r in sqlite3.connect(tmp_path / "t.db").execute(
"PRAGMA table_info(llm_calls)"
)]
cols = [
r[1] for r in sqlite3.connect(tmp_path / "t.db").execute("PRAGMA table_info(llm_calls)")
]
assert cols == _EXPECTED_COLUMNS
async def test_call_id_idempotent(self, tmp_path):
@@ -65,18 +108,20 @@ class TestSQLiteRecorder:
await _record_minimal(recorder, call_id="dup")
await _record_minimal(recorder, call_id="dup", response="second")
recorder.close()
rows = sqlite3.connect(tmp_path / "t.db").execute(
"SELECT response FROM llm_calls WHERE call_id='dup'"
).fetchall()
rows = (
sqlite3.connect(tmp_path / "t.db")
.execute("SELECT response FROM llm_calls WHERE call_id='dup'")
.fetchall()
)
assert rows == [("ok",)] # INSERT OR IGNORE: 第二次静默忽略
async def test_concurrent_writes_all_land(self, tmp_path):
recorder = SQLiteRecorder(tmp_path / "t.db")
await asyncio.gather(*(_record_minimal(recorder, call_id=f"c{i}") for i in range(50)))
recorder.close()
(count,) = sqlite3.connect(tmp_path / "t.db").execute(
"SELECT COUNT(*) FROM llm_calls"
).fetchone()
(count,) = (
sqlite3.connect(tmp_path / "t.db").execute("SELECT COUNT(*) FROM llm_calls").fetchone()
)
assert count == 50
async def test_unwritable_path_degrades_silently(self):
@@ -98,8 +143,12 @@ class TestEmitter:
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
await emitter.emit_attempt(
request=_REQ, source=_source(), call_id="cid-1", latency_ms=42,
response=_resp(), error=None,
request=_REQ,
source=_source(),
call_id="cid-1",
latency_ms=42,
response=_resp(),
error=None,
)
row = rec.rows[0]
assert row["call_id"] == "cid-1" and row["error"] is None
@@ -110,8 +159,12 @@ class TestEmitter:
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
await emitter.emit_attempt(
request=_REQ, source=_source(), call_id="cid-2", latency_ms=7,
response=None, error="TransientError: boom",
request=_REQ,
source=_source(),
call_id="cid-2",
latency_ms=7,
response=None,
error="TransientError: boom",
)
row = rec.rows[0]
assert row["error"].startswith("TransientError")
@@ -121,12 +174,23 @@ class TestEmitter:
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
big = "data:image/png;base64," + "A" * 100_000
req = ChatRequest(messages=[{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": big}},
]}])
req = ChatRequest(
messages=[
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": big}},
],
}
]
)
await emitter.emit_attempt(
request=req, source=_source(), call_id="c", latency_ms=1,
response=None, error="x",
request=req,
source=_source(),
call_id="c",
latency_ms=1,
response=None,
error="x",
)
assert len(rec.rows[0]["messages"]) < 500 # base64 不整段进库(VT R12)
@@ -137,8 +201,12 @@ class TestEmitter:
emitter = TelemetryEmitter(Broken())
await emitter.emit_attempt(
request=_REQ, source=_source(), call_id="c", latency_ms=1,
response=_resp(), error=None,
request=_REQ,
source=_source(),
call_id="c",
latency_ms=1,
response=_resp(),
error=None,
) # 不抛(降级不冒泡)
@@ -194,10 +262,9 @@ def test_single_emitter_discipline():
"""铁律执法: record_llm_call 在 src/ 的调用点只允许出现在 telemetry emitter。"""
out = subprocess.run(
["grep", "-rln", "record_llm_call(", "src/polygateway"],
capture_output=True, text=True, cwd=Path(__file__).resolve().parents[2],
capture_output=True,
text=True,
cwd=Path(__file__).resolve().parents[2],
).stdout.splitlines()
callers = [
p for p in out
if not p.endswith(("ports.py", "telemetry/sqlite.py"))
]
callers = [p for p in out if not p.endswith(("ports.py", "telemetry/sqlite.py"))]
assert callers == ["src/polygateway/middleware/telemetry.py"]
tests/unit/test_types.py
+16 -10
View File
@@ -19,14 +19,14 @@ from polygateway.types import (
def _make_source(**overrides):
"""构造最小合法 SourceConfig,单点覆盖便于逐条触发不变式。"""
base = dict(
name="qwen_1",
provider="qwen",
base_url="https://gw.example/v1",
api_key="sk-test",
model="qwen-max",
timeout_s=120.0,
)
base = {
"name": "qwen_1",
"provider": "qwen",
"base_url": "https://gw.example/v1",
"api_key": "sk-test",
"model": "qwen-max",
"timeout_s": 120.0,
}
base.update(overrides)
return SourceConfig(**base)
@@ -144,7 +144,13 @@ class TestAuxTypes:
u = Usage(prompt_tokens=10, completion_tokens=20, usage_source="estimated")
assert u.prompt_tokens == 10
s = TransportResult(
content="c", thinking="", prompt_tokens=1, completion_tokens=2,
usage_source="measured", ttft_ms=12.5, max_inter_token_ms=30.0, raw={"id": "x"},
content="c",
thinking="",
prompt_tokens=1,
completion_tokens=2,
usage_source="measured",
ttft_ms=12.5,
max_inter_token_ms=30.0,
raw={"id": "x"},
)
assert s.raw["id"] == "x"