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.
97 lines
3.2 KiB
Python
Executable File
97 lines
3.2 KiB
Python
Executable File
"""Cost computation and attribution (架构文档 §13, FR-BILL-08)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from decimal import Decimal
|
|
|
|
from packages.llm.cost import (
|
|
AICallRecord,
|
|
CompositeAICallRecorder,
|
|
InMemoryAICallRecorder,
|
|
ModelRate,
|
|
NullAICallRecorder,
|
|
RateTable,
|
|
compute_cost_cents,
|
|
)
|
|
from packages.llm.recording import PrometheusAICallRecorder
|
|
from packages.llm.types import Usage
|
|
|
|
RATE = ModelRate(
|
|
input_cents_per_mtok=Decimal("100"),
|
|
cached_input_cents_per_mtok=Decimal("10"),
|
|
output_cents_per_mtok=Decimal("200"),
|
|
)
|
|
|
|
|
|
def test_cached_tokens_are_billed_at_the_cheaper_rate() -> None:
|
|
usage = Usage(tokens_in=1_000_000, tokens_out=0, tokens_cached=1_000_000)
|
|
assert compute_cost_cents(usage, RATE) == Decimal("10.000000")
|
|
|
|
|
|
def test_uncached_input_is_billed_at_the_full_rate() -> None:
|
|
usage = Usage(tokens_in=1_000_000, tokens_out=0, tokens_cached=0)
|
|
assert compute_cost_cents(usage, RATE) == Decimal("100.000000")
|
|
|
|
|
|
def test_mixed_usage_splits_cached_and_uncached_input() -> None:
|
|
# 3.7k input of which 3k cached, 800 output — the §7.2 Agent token budget.
|
|
usage = Usage(tokens_in=3700, tokens_out=800, tokens_cached=3000)
|
|
expected = (
|
|
Decimal(700) * RATE.input_cents_per_mtok
|
|
+ Decimal(3000) * RATE.cached_input_cents_per_mtok
|
|
+ Decimal(800) * RATE.output_cents_per_mtok
|
|
) / Decimal(1_000_000)
|
|
assert compute_cost_cents(usage, RATE) == expected.quantize(Decimal("0.000001"))
|
|
|
|
|
|
def test_zero_usage_costs_nothing() -> None:
|
|
assert compute_cost_cents(Usage(), RATE) == Decimal("0.000000")
|
|
|
|
|
|
def test_billable_input_never_goes_negative() -> None:
|
|
usage = Usage(tokens_in=100, tokens_cached=500)
|
|
assert usage.billable_input == 0
|
|
|
|
|
|
def test_rate_table_falls_back_so_unknown_models_are_never_free() -> None:
|
|
default = ModelRate(Decimal("1"), Decimal("1"), Decimal("1"))
|
|
table = RateTable({"deepseek-v4-flash": RATE}, default=default)
|
|
assert table.for_model("deepseek-v4-flash") is RATE
|
|
assert table.for_model("some-new-model") is default
|
|
assert table.known_models() == ("deepseek-v4-flash",)
|
|
|
|
|
|
def _record(success: bool = True) -> AICallRecord:
|
|
return AICallRecord(
|
|
call_id=1,
|
|
scene="agent_chat",
|
|
provider="deepseek",
|
|
model="deepseek-v4-flash",
|
|
success=success,
|
|
usage=Usage(tokens_in=10, tokens_out=5, tokens_cached=2),
|
|
latency_ms=120,
|
|
cost_cents=Decimal("0.5"),
|
|
)
|
|
|
|
|
|
async def test_with_defaults_stamps_created_at() -> None:
|
|
stamped = _record().with_defaults()
|
|
assert stamped.created_at is not None
|
|
assert stamped.with_defaults() is stamped
|
|
|
|
|
|
async def test_in_memory_recorder_collects_records() -> None:
|
|
recorder = InMemoryAICallRecorder()
|
|
await recorder.record(_record())
|
|
assert len(recorder.records) == 1
|
|
|
|
|
|
async def test_null_recorder_is_a_no_op() -> None:
|
|
assert await NullAICallRecorder().record(_record()) is None
|
|
|
|
|
|
async def test_composite_recorder_fans_out() -> None:
|
|
first, second = InMemoryAICallRecorder(), InMemoryAICallRecorder()
|
|
await CompositeAICallRecorder(first, second, PrometheusAICallRecorder()).record(_record())
|
|
assert len(first.records) == len(second.records) == 1
|