feat: add gateway client with env-driven assembly
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:
@@ -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)
|
||||
Reference in New Issue
Block a user