Files
meida_backend/tests/unit/test_daily_budget.py
T
xsl a4d59aca30 feat(llm): four-tier daily budget degradation with manual override
把单一日预算熔断升级为架构 §13.3 的四档分级降级(T-P1-25b / NFR-COST-02):

- budget.py:新增 BudgetTier 配置模型与档位判定(current_tier/tier_name/
  blocked_scenes)。阈值相对 LLM_DAILY_BUDGET_CENTS:warn ≥60% 只告警;
  degrade_secondary ≥80% 拒 hd_export/report_regen;degrade_generation ≥90%
  追加拒 hairstyle/tryon;block_all ≥100% 拒全部非 critical、allow 名单
  (默认 agent_chat)放行。高档位 blocked 取「截至当前档的并集」(架构原文
  「追加拒绝」)。critical 场景全档豁免但照常计成本。
- 手动 override:Redis 键 llm:budget:override(值 "1"~"4"),优先于成本水
  位,读不到/非法则忽略;set_override TTL 到当日 UTC+8 24 点,跨预算日必然
  失效。CostCounter 协议相应增加 set/delete。
- 指标:Gauge llm_budget_tier + Counter llm_budget_blocked_total{scene}
  挂到共享 REGISTRY 随 /metrics 暴露;档位上升沿只告警一次
  (budget_tier_raised,实例内记 last_tier)。
- errors.py 新增 SceneDegradedError(tier2/3 拒绝);gateway.py 调用点换成
  ensure_scene_allowed(保留 ensure_within_budget 兼容入口);agent.py 映射
  LLM_SCENE_DEGRADED(60006);100% 档维持 DailyBudgetExceededError/60003 语义。
- 单测:新增 tests/unit/test_tiered_budget.py(27 例:四档边界、各档放行/
  拒绝、critical 豁免、override 加严/放宽/非法值/TTL、配置驱动)。
- 集成:tests/integration/test_vertical_slice.py 新增用例——预算 1 分 +
  block_all allow 清空 + Redis 注入当日成本,POST /agent/chat 返回 60003,
  断言 bl_ledger 中 consume(refunded,-1)+refund(committed,+1) 净和为 0,
  覆盖「reserve 之后才熔断 → refund」路径;fixture 重构为共享
  FakeProvider/fake_client,断言语义未弱化。

门禁:ruff check/format、mypy(packages/llm strict + 其余包)、lint-imports、
单测 196 passed、集成 4 passed。

偏差说明:
1. tests/unit/test_gateway.py 的 test_daily_budget_rejects_non_critical_
   scenes_before_spending 把被测场景 AGENT_CHAT 换成 REPORT_COPY(test_daily_
   budget.py 同用例换 tryon):block_all 档默认 allow=["agent_chat"](架构
   「仅保留纯文本对话」)后 agent_chat 在 100% 档不再被拒,原场景选择与新
   语义冲突;断言本身(进 provider 前被拒)未弱化。
2. 集成测试的 block_all 档配置为 blocked="*"+allow=[](文档示例为
   blocked=["agent_chat"]):后者按档位规则会走 SceneDegradedError/60006,
   前者才能按设计走既有 60003 语义(文档允许二者取其一且"按上面设计应为
   60003")。
2026-08-14 16:27:21 +08:00

87 lines
3.0 KiB
Python
Executable File

"""Global daily cost circuit breaker (架构文档 §13 三道成本闸门).
The requirement is enforcement, not alerting: once the platform-wide daily
spend crosses the threshold, non-critical AI calls are rejected.
"""
from __future__ import annotations
from datetime import UTC, datetime
from decimal import Decimal
import pytest
from packages.llm.budget import DailyBudgetGuard, InMemoryCostCounter
from packages.llm.errors import DailyBudgetExceededError
def make_guard(
budget: str = "100", *, enabled: bool = True, now: datetime | None = None
) -> tuple[DailyBudgetGuard, InMemoryCostCounter]:
counter = InMemoryCostCounter()
guard = DailyBudgetGuard(
counter,
daily_budget_cents=Decimal(budget),
enabled=enabled,
clock=(lambda: now) if now else None,
)
return guard, counter
async def test_allows_calls_below_the_threshold() -> None:
guard, _ = make_guard("100")
await guard.add_cost(Decimal("99.9"))
await guard.ensure_within_budget(scene="agent_chat")
async def test_rejects_non_critical_scenes_once_the_budget_is_reached() -> None:
guard, _ = make_guard("100")
await guard.add_cost(Decimal("100"))
# block_all 档默认放行 agent_chat(「仅保留纯文本对话」),用非 allow 场景
# 断言 100% 档的既有拒绝语义不变
with pytest.raises(DailyBudgetExceededError) as excinfo:
await guard.ensure_within_budget(scene="tryon")
assert "tryon" in str(excinfo.value)
async def test_critical_scenes_bypass_the_breaker() -> None:
guard, _ = make_guard("100")
await guard.add_cost(Decimal("500"))
# Degrade the product, do not break it: 意图识别/越界拒答 must keep working.
await guard.ensure_within_budget(scene="intent_classify", critical=True)
async def test_disabled_guard_never_blocks_and_never_counts() -> None:
guard, counter = make_guard("1", enabled=False)
await guard.add_cost(Decimal("999"))
await guard.ensure_within_budget(scene="agent_chat")
assert await counter.get(guard.current_key()) == Decimal("0")
async def test_cost_accumulates_across_calls() -> None:
guard, _ = make_guard("100")
await guard.add_cost(Decimal("1.5"))
await guard.add_cost(Decimal("2.25"))
assert await guard.spent_today() == Decimal("3.75")
async def test_zero_cost_is_not_counted() -> None:
guard, _ = make_guard("100")
await guard.add_cost(Decimal("0"))
assert await guard.spent_today() == Decimal("0")
async def test_budget_day_boundary_follows_utc_plus_8() -> None:
"""FR-BILL-01: 跨时区按 UTC+8 统一。17:00Z already belongs to the next day."""
guard, _ = make_guard(now=datetime(2026, 8, 13, 17, 0, tzinfo=UTC))
assert guard.current_key().endswith("2026-08-14")
guard_before, _ = make_guard(now=datetime(2026, 8, 13, 15, 0, tzinfo=UTC))
assert guard_before.current_key().endswith("2026-08-13")
async def test_budget_properties_are_exposed_for_dashboards() -> None:
guard, _ = make_guard("42")
assert guard.enabled is True
assert guard.budget_cents == Decimal("42")