Files
meida_backend/packages/llm/budget.py
T
xsl 9e7a886863 chore(prep): shared groundwork for the parallel task batch
多 agent 并行批次的协调提交(消除全部共享文件冲突面):
- 迁移 0004:id_session + ob_event 两表(T3/T4 用,任务不再写迁移)
- ORM:UserSession / ObEvent;错误码:九个新码(40xxx/60006,含 meta)
- 配置:四档 tiers / refresh TTL / 任务奖励注册表 / 埋点上限 / COS 直传全套
- 骨架:packages/analytics + packages/storage + events/uploads 空路由(已注册)
- deps:AnalyticsServiceDep / StorageServiceDep;DailyBudgetGuard 接收 tiers
- uv add cos-python-sdk-v5;import-linter/CI/Makefile 纳入两个新包
2026-08-14 16:02:54 +08:00

125 lines
4.6 KiB
Python

"""Global daily cost circuit breaker (架构文档 §13「必须设置的三道成本闸门」).
Gate 1 in §13 is described as "某功能超预算自动降级,并告警". Alerting alone is
not enough: the failure modes that blow up AI spend (retry storms, a cache-key
bug, someone scripting the API) all produce cost faster than a human can react
to a page. So this is an *enforcement* point, not a notification: once the
platform-wide daily spend crosses the threshold, non-critical scenes are
rejected outright.
Critical scenes (e.g. content moderation follow-ups, paid-path completions)
opt out via ``SceneConfig.critical`` so that a budget overrun degrades the
product instead of breaking it.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta, timezone
from decimal import Decimal
from typing import Protocol, runtime_checkable
from packages.llm.errors import DailyBudgetExceededError
_DAY_SECONDS = 60 * 60 * 24
@runtime_checkable
class CostCounter(Protocol):
"""Atomic accumulating counter. Redis in production, dict in tests."""
async def incr(self, key: str, amount: Decimal, ttl_seconds: int) -> Decimal: ...
async def get(self, key: str) -> Decimal: ...
class InMemoryCostCounter:
"""Process-local counter. Correct only for a single process — tests only."""
def __init__(self) -> None:
self._values: dict[str, Decimal] = {}
async def incr(self, key: str, amount: Decimal, ttl_seconds: int) -> Decimal:
self._values[key] = self._values.get(key, Decimal("0")) + amount
return self._values[key]
async def get(self, key: str) -> Decimal:
return self._values.get(key, Decimal("0"))
class RedisCostCounter:
"""Cross-process counter backed by ``INCRBYFLOAT`` + a rolling TTL."""
def __init__(self, redis: object) -> None:
self._redis = redis
async def incr(self, key: str, amount: Decimal, ttl_seconds: int) -> Decimal:
pipe = self._redis.pipeline() # type: ignore[attr-defined]
pipe.incrbyfloat(key, float(amount))
pipe.expire(key, ttl_seconds)
results = await pipe.execute()
return Decimal(str(results[0]))
async def get(self, key: str) -> Decimal:
value = await self._redis.get(key) # type: ignore[attr-defined]
return Decimal(str(value)) if value is not None else Decimal("0")
class DailyBudgetGuard:
"""Hard platform-wide daily AI budget."""
def __init__(
self,
counter: CostCounter,
*,
daily_budget_cents: Decimal,
enabled: bool = True,
key_prefix: str = "llm:cost:daily",
tz_offset_hours: int = 8,
clock: object | None = None,
tiers: list[dict[str, object]] | None = None,
) -> None:
self._counter = counter
self._budget = daily_budget_cents
self._enabled = enabled
self._key_prefix = key_prefix
self._tz = timezone(timedelta(hours=tz_offset_hours))
self._clock = clock
# 四档分级降级的配置(架构 §13.3 / T-P1-25b)。T1 任务实现判定逻辑;
# 这里先接收配置,使 deps 装配层不必随后再改。None = 使用内置默认四档。
self._tiers_config = tiers
@property
def enabled(self) -> bool:
return self._enabled
@property
def budget_cents(self) -> Decimal:
return self._budget
def _now(self) -> datetime:
if self._clock is not None:
return self._clock() # type: ignore[operator,no-any-return]
return datetime.now(UTC)
def current_key(self) -> str:
"""Budget day boundary follows UTC+8, consistent with FR-BILL-01."""
local = self._now().astimezone(self._tz)
return f"{self._key_prefix}:{local:%Y-%m-%d}"
async def spent_today(self) -> Decimal:
return await self._counter.get(self.current_key())
async def ensure_within_budget(self, *, scene: str, critical: bool = False) -> None:
"""Raise :class:`DailyBudgetExceededError` for non-critical scenes over budget."""
if not self._enabled or critical:
return
spent = await self.spent_today()
if spent >= self._budget:
raise DailyBudgetExceededError(str(spent), str(self._budget), scene)
async def add_cost(self, cost_cents: Decimal) -> Decimal:
"""Accumulate spend. Called after every provider attempt, success or not."""
if not self._enabled or cost_cents <= 0:
return await self.spent_today() if self._enabled else Decimal("0")
return await self._counter.incr(self.current_key(), cost_cents, _DAY_SECONDS)