xsl
93b88db2f5
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.
470 lines
17 KiB
Python
Executable File
470 lines
17 KiB
Python
Executable File
"""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
|