test: unit suite for the gateway invariants plus an integration slice
Covers routing, failover, the circuit breaker, structured-output retries, the cache-prefix validator, cost computation including cached-token pricing, the daily budget breaker, the quota protocol, the task state machine and log redaction. HTTP is mocked -- no unit test spends money. The integration test runs the real slice against PostgreSQL and Redis with a faked provider.
This commit is contained in:
Executable
+469
@@ -0,0 +1,469 @@
|
||||
"""LLM gateway: routing, failover, circuit breaking, structured output, cost.
|
||||
|
||||
No test here touches the network — providers are fakes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from collections.abc import Sequence
|
||||
from decimal import Decimal
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from packages.llm.budget import DailyBudgetGuard, InMemoryCostCounter
|
||||
from packages.llm.circuit import CircuitBreakerRegistry, CircuitState
|
||||
from packages.llm.cost import InMemoryAICallRecorder, ModelRate, RateTable
|
||||
from packages.llm.errors import (
|
||||
AllProvidersFailedError,
|
||||
DailyBudgetExceededError,
|
||||
ProviderBadRequestError,
|
||||
ProviderRateLimitedError,
|
||||
ProviderServerError,
|
||||
ProviderTimeoutError,
|
||||
SceneNotRoutedError,
|
||||
StructuredOutputError,
|
||||
)
|
||||
from packages.llm.gateway import LLMGateway, extract_json
|
||||
from packages.llm.prompts import PromptRegistry
|
||||
from packages.llm.providers.zhipu import ZhipuGLMProvider
|
||||
from packages.llm.routing import SceneRouter, build_scene_configs
|
||||
from packages.llm.types import Message, ProviderRequest, ProviderResponse, Scene, Usage
|
||||
|
||||
RATE = ModelRate(
|
||||
input_cents_per_mtok=Decimal("100"),
|
||||
cached_input_cents_per_mtok=Decimal("10"),
|
||||
output_cents_per_mtok=Decimal("200"),
|
||||
)
|
||||
|
||||
|
||||
class FakeProvider:
|
||||
"""Returns queued responses/exceptions in order; repeats the last forever."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
outcomes: Sequence[ProviderResponse | Exception],
|
||||
default_model: str = "fake-model",
|
||||
) -> None:
|
||||
self._name = name
|
||||
self._outcomes = list(outcomes)
|
||||
self._default_model = default_model
|
||||
self.requests: list[ProviderRequest] = []
|
||||
self.closed = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def default_model(self) -> str:
|
||||
return self._default_model
|
||||
|
||||
async def complete(self, request: ProviderRequest) -> ProviderResponse:
|
||||
self.requests.append(request)
|
||||
outcome = self._outcomes[min(len(self.requests) - 1, len(self._outcomes) - 1)]
|
||||
if isinstance(outcome, Exception):
|
||||
raise outcome
|
||||
return outcome
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def response(text: str, model: str = "fake-model") -> ProviderResponse:
|
||||
return ProviderResponse(
|
||||
text=text,
|
||||
model=model,
|
||||
usage=Usage(tokens_in=1000, tokens_out=200, tokens_cached=800),
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
|
||||
def build_gateway(
|
||||
prompts_dir: Path,
|
||||
providers: dict[str, FakeProvider],
|
||||
*,
|
||||
primary: str = "p1",
|
||||
backup: str | None = "p2",
|
||||
budget_cents: str = "100000",
|
||||
budget_enabled: bool = True,
|
||||
failure_threshold: int = 3,
|
||||
structured_max_retries: int = 2,
|
||||
recorder: InMemoryAICallRecorder | None = None,
|
||||
) -> tuple[LLMGateway, InMemoryAICallRecorder, DailyBudgetGuard]:
|
||||
recorder = recorder or InMemoryAICallRecorder()
|
||||
guard = DailyBudgetGuard(
|
||||
InMemoryCostCounter(),
|
||||
daily_budget_cents=Decimal(budget_cents),
|
||||
enabled=budget_enabled,
|
||||
)
|
||||
gateway = LLMGateway(
|
||||
providers=dict(providers),
|
||||
router=SceneRouter(
|
||||
build_scene_configs(
|
||||
primary_provider=primary,
|
||||
backup_provider=backup,
|
||||
structured_max_retries=structured_max_retries,
|
||||
)
|
||||
),
|
||||
prompts=PromptRegistry(prompts_dir),
|
||||
rates=RateTable({"fake-model": RATE}, default=RATE),
|
||||
recorder=recorder,
|
||||
budget=guard,
|
||||
breakers=CircuitBreakerRegistry(failure_threshold=failure_threshold),
|
||||
id_factory=itertools.count(1).__next__,
|
||||
)
|
||||
return gateway, recorder, guard
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# routing
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_routes_to_the_primary_provider(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("好的")])
|
||||
backup = FakeProvider("p2", [response("backup")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})
|
||||
|
||||
result = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "面试"})
|
||||
|
||||
assert result.provider == "p1"
|
||||
assert result.text == "好的"
|
||||
assert backup.requests == []
|
||||
|
||||
|
||||
async def test_scene_config_drives_model_and_sampling(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")], default_model="deepseek-v4-flash")
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
|
||||
request = primary.requests[0]
|
||||
assert request.model == "deepseek-v4-flash"
|
||||
assert request.max_tokens == 800 # 架构文档 §7.2 输出预算
|
||||
assert request.timeout_seconds == 30.0
|
||||
|
||||
|
||||
async def test_cacheable_system_prefix_is_the_first_message(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
await gateway.complete(
|
||||
scene=Scene.AGENT_CHAT,
|
||||
messages=[Message(role="user", content="历史消息")],
|
||||
prompt_vars={"user_input": "现在的问题"},
|
||||
)
|
||||
|
||||
messages = primary.requests[0].messages
|
||||
assert messages[0].role == "system"
|
||||
assert "穿搭顾问助手" in messages[0].content
|
||||
assert messages[-1].content == "现在的问题"
|
||||
|
||||
|
||||
async def test_unrouted_scene_raises(prompts_dir: Path) -> None:
|
||||
gateway = LLMGateway(
|
||||
providers={},
|
||||
router=SceneRouter({}),
|
||||
prompts=PromptRegistry(prompts_dir),
|
||||
rates=RateTable({}, default=RATE),
|
||||
recorder=InMemoryAICallRecorder(),
|
||||
budget=DailyBudgetGuard(
|
||||
InMemoryCostCounter(), daily_budget_cents=Decimal("1"), enabled=False
|
||||
),
|
||||
)
|
||||
with pytest.raises(SceneNotRoutedError):
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT)
|
||||
|
||||
|
||||
async def test_unregistered_provider_is_reported_as_a_failure(prompts_dir: Path) -> None:
|
||||
gateway, _, _ = build_gateway(prompts_dir, {}, primary="ghost", backup=None)
|
||||
with pytest.raises(AllProvidersFailedError):
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# failover
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"failure",
|
||||
[
|
||||
ProviderTimeoutError("p1", "timed out"),
|
||||
ProviderServerError("p1", "502", status_code=502),
|
||||
ProviderRateLimitedError("p1", "429", status_code=429),
|
||||
],
|
||||
)
|
||||
async def test_timeout_5xx_and_rate_limit_fail_over_to_the_backup(
|
||||
prompts_dir: Path, failure: Exception
|
||||
) -> None:
|
||||
primary = FakeProvider("p1", [failure])
|
||||
backup = FakeProvider("p2", [response("备用回答")])
|
||||
gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})
|
||||
|
||||
result = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
|
||||
assert result.provider == "p2"
|
||||
assert result.text == "备用回答"
|
||||
# Both the failed and the successful attempt are attributable.
|
||||
assert [(r.provider, r.success) for r in recorder.records] == [("p1", False), ("p2", True)]
|
||||
|
||||
|
||||
async def test_bad_request_does_not_fail_over(prompts_dir: Path) -> None:
|
||||
"""A malformed request would fail on the backup too — burning money twice."""
|
||||
primary = FakeProvider("p1", [ProviderBadRequestError("p1", "400", status_code=400)])
|
||||
backup = FakeProvider("p2", [response("should not be used")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})
|
||||
|
||||
with pytest.raises(ProviderBadRequestError):
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
assert backup.requests == []
|
||||
|
||||
|
||||
async def test_all_providers_failing_raises_with_every_cause(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [ProviderTimeoutError("p1", "timeout")])
|
||||
backup = FakeProvider("p2", [ProviderServerError("p2", "500", status_code=500)])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})
|
||||
|
||||
with pytest.raises(AllProvidersFailedError) as excinfo:
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
assert len(excinfo.value.failures) == 2
|
||||
|
||||
|
||||
async def test_zhipu_placeholder_participates_in_the_real_routing_path(
|
||||
prompts_dir: Path,
|
||||
) -> None:
|
||||
"""The backup slot is structural, but the code path through it is real."""
|
||||
primary = FakeProvider("p1", [ProviderTimeoutError("p1", "timeout")])
|
||||
gateway, recorder, _ = build_gateway(
|
||||
prompts_dir, {"p1": primary, "zhipu": ZhipuGLMProvider()}, backup="zhipu"
|
||||
)
|
||||
|
||||
with pytest.raises(AllProvidersFailedError) as excinfo:
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
|
||||
assert isinstance(excinfo.value.failures[1], NotImplementedError)
|
||||
assert {r.provider for r in recorder.records} == {"p1", "zhipu"}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# circuit breaker
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_open_circuit_skips_the_primary_entirely(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [ProviderServerError("p1", "500", status_code=500)])
|
||||
backup = FakeProvider("p2", [response("备用")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup}, failure_threshold=1)
|
||||
|
||||
first = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "a"})
|
||||
assert first.provider == "p2"
|
||||
assert gateway.breaker_states()["p1"] == CircuitState.OPEN.value
|
||||
|
||||
second = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "b"})
|
||||
assert second.provider == "p2"
|
||||
assert len(primary.requests) == 1 # not attempted the second time
|
||||
|
||||
|
||||
async def test_success_keeps_the_circuit_closed(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "a"})
|
||||
assert gateway.breaker_states()["p1"] == CircuitState.CLOSED.value
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# structured output
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Outfit(BaseModel):
|
||||
top: str
|
||||
bottom: str
|
||||
|
||||
|
||||
async def test_structured_output_is_parsed_and_validated(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response('{"top": "白衬衫", "bottom": "西裤"}')])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
result = await gateway.complete(
|
||||
scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
|
||||
)
|
||||
|
||||
assert result.parsed_as(Outfit).top == "白衬衫"
|
||||
assert result.attempts == 1
|
||||
assert primary.requests[0].json_mode is True
|
||||
assert "JSON Schema" in primary.requests[0].messages[-1].content
|
||||
|
||||
|
||||
async def test_invalid_structured_output_is_retried_with_the_validation_error(
|
||||
prompts_dir: Path,
|
||||
) -> None:
|
||||
primary = FakeProvider(
|
||||
"p1",
|
||||
[
|
||||
response("这不是 JSON"),
|
||||
response('{"top": "针织衫"}'), # missing `bottom`
|
||||
response('{"top": "针织衫", "bottom": "阔腿裤"}'),
|
||||
],
|
||||
)
|
||||
gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
result = await gateway.complete(
|
||||
scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
|
||||
)
|
||||
|
||||
assert result.attempts == 3
|
||||
assert result.parsed_as(Outfit).bottom == "阔腿裤"
|
||||
# Every attempt costs money and every attempt is recorded.
|
||||
assert len(recorder.records) == 3
|
||||
assert [r.success for r in recorder.records] == [False, False, True]
|
||||
assert recorder.records[0].error_code == "llm_structured_output_invalid"
|
||||
# The repair turn carries the validation error back to the model.
|
||||
assert "校验错误" in primary.requests[1].messages[-1].content
|
||||
|
||||
|
||||
async def test_structured_output_gives_up_and_fails_over(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("never json")])
|
||||
backup = FakeProvider("p2", [response('{"top": "T恤", "bottom": "牛仔裤"}')])
|
||||
gateway, _, _ = build_gateway(
|
||||
prompts_dir, {"p1": primary, "p2": backup}, structured_max_retries=1
|
||||
)
|
||||
|
||||
result = await gateway.complete(
|
||||
scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
|
||||
)
|
||||
|
||||
assert result.provider == "p2"
|
||||
assert len(primary.requests) == 2 # 1 + structured_max_retries
|
||||
|
||||
|
||||
async def test_structured_failure_on_the_only_provider_surfaces(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("never json")])
|
||||
gateway, _, _ = build_gateway(
|
||||
prompts_dir, {"p1": primary}, backup=None, structured_max_retries=0
|
||||
)
|
||||
with pytest.raises(AllProvidersFailedError) as excinfo:
|
||||
await gateway.complete(
|
||||
scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
|
||||
)
|
||||
assert isinstance(excinfo.value.failures[0], StructuredOutputError)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[
|
||||
('{"a": 1}', '{"a": 1}'),
|
||||
('```json\n{"a": 1}\n```', '{"a": 1}'),
|
||||
('好的：\n{"a": 1}\n以上。', '{"a": 1}'),
|
||||
("no json here", "no json here"),
|
||||
],
|
||||
)
|
||||
def test_extract_json_handles_fenced_and_chatty_replies(raw: str, expected: str) -> None:
|
||||
assert extract_json(raw) == expected
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# cost attribution + daily budget
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_every_call_records_a_cost_row(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
result = await gateway.complete(
|
||||
scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, user_id=42, ref_id="req-1"
|
||||
)
|
||||
|
||||
record = recorder.records[0]
|
||||
assert record.scene == "agent_chat"
|
||||
assert record.provider == "p1"
|
||||
assert record.model == "fake-model"
|
||||
assert record.prompt_id == "agent_chat"
|
||||
assert record.prompt_version == "v1" # 可复现: 记录 prompt 版本
|
||||
assert record.user_id == 42
|
||||
assert record.ref_id == "req-1"
|
||||
assert record.usage.tokens_cached == 800
|
||||
assert record.cost_cents == result.cost_cents > Decimal("0")
|
||||
assert record.created_at is not None
|
||||
|
||||
|
||||
async def test_failed_calls_are_recorded_too(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [ProviderTimeoutError("p1", "timed out")])
|
||||
gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
with pytest.raises(AllProvidersFailedError):
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
|
||||
record = recorder.records[0]
|
||||
assert record.success is False
|
||||
assert record.error_code == "llm_timeout"
|
||||
assert record.cost_cents == Decimal("0")
|
||||
|
||||
|
||||
async def test_a_broken_recorder_never_breaks_the_user_request(prompts_dir: Path) -> None:
|
||||
class ExplodingRecorder:
|
||||
async def record(self, record: object) -> None:
|
||||
msg = "database is down"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
gateway._recorder = ExplodingRecorder()
|
||||
|
||||
result = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
assert result.text == "ok"
|
||||
|
||||
|
||||
async def test_spend_accumulates_into_the_daily_budget(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
gateway, _, guard = build_gateway(prompts_dir, {"p1": primary}, backup=None)
|
||||
|
||||
first = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "a"})
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "b"})
|
||||
|
||||
assert await guard.spent_today() == first.cost_cents * 2
|
||||
|
||||
|
||||
async def test_daily_budget_rejects_non_critical_scenes_before_spending(
|
||||
prompts_dir: Path,
|
||||
) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
gateway, _, guard = build_gateway(
|
||||
prompts_dir, {"p1": primary}, backup=None, budget_cents="0.00001"
|
||||
)
|
||||
await guard.add_cost(Decimal("1"))
|
||||
|
||||
with pytest.raises(DailyBudgetExceededError):
|
||||
await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
|
||||
assert primary.requests == [] # rejected, not merely alerted
|
||||
|
||||
|
||||
async def test_daily_budget_lets_critical_scenes_through(prompts_dir: Path, tmp_path: Path) -> None:
|
||||
(prompts_dir / "intent_classify").mkdir()
|
||||
(prompts_dir / "intent_classify" / "v1.txt").write_text(
|
||||
"[system]\n你是意图分类器。\n\n[user]\n{user_input}\n", encoding="utf-8"
|
||||
)
|
||||
primary = FakeProvider("p1", [response("scene_consult")])
|
||||
gateway, _, guard = build_gateway(
|
||||
prompts_dir, {"p1": primary}, backup=None, budget_cents="0.00001"
|
||||
)
|
||||
await guard.add_cost(Decimal("1"))
|
||||
|
||||
result = await gateway.complete(
|
||||
scene=Scene.INTENT_CLASSIFY, prompt_vars={"user_input": "面试穿什么"}
|
||||
)
|
||||
assert result.text == "scene_consult"
|
||||
|
||||
|
||||
async def test_aclose_closes_every_provider(prompts_dir: Path) -> None:
|
||||
primary = FakeProvider("p1", [response("ok")])
|
||||
backup = FakeProvider("p2", [response("ok")])
|
||||
gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})
|
||||
await gateway.aclose()
|
||||
assert primary.closed and backup.closed
|
||||
Reference in New Issue
Block a user