meida_backend/packages/llm/budget.py

"""Global daily cost circuit breaker — four-tier graded degradation (架构文档 §13.3).

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, and the degradation is graded so
the product winds down gracefully instead of dying all at once. Thresholds are
percentages of ``LLM_DAILY_BUDGET_CENTS``; higher tiers *add* rejected scenes on
top of the lower ones (架构原文「追加拒绝」), never un-block them::

    tier 1  warn                ≥60%    只告警（指标 + 上升沿日志），不影响请求
    tier 2  degrade_secondary   ≥80%    拒绝 hd_export / report_regen
    tier 3  degrade_generation  ≥90%    追加拒绝 hairstyle / tryon
    tier 4  block_all           ≥100%   拒绝所有非 critical；allow 名单放行
                                        （默认 agent_chat ——「仅保留在途任务与
                                        纯文本对话」）

Critical scenes (``SceneConfig.critical``, e.g. INTENT_CLASSIFY) opt out of
every tier so that a budget overrun degrades the product instead of breaking
it — but their cost still counts towards the water level. 在途任务同理放行：
熔断只拦网关**新请求**，worker 侧不做二次检查。

手动干预（运营/演练，redis-cli 可直接操作）::

    SET llm:budget:override 4   # 强制 tier 4（加严演练）
    SET llm:budget:override 2   # 水位已到 tier 4 时手动放宽回 tier 2
    DEL  llm:budget:override    # 解除，回落成本水位判定

键 ``llm:budget:override`` 的值为 "1"~"4"（档位序号），优先级高于成本水位；
读不到或值非法则忽略。程序化写入走 :meth:`DailyBudgetGuard.set_override`，
TTL 到当日 UTC+8 24 点 —— override 最多活一个预算日，第二天必然回落水位。
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta, timezone
from decimal import Decimal
from typing import Literal, Protocol, runtime_checkable

from prometheus_client import Counter, Gauge
from pydantic import BaseModel, ConfigDict, Field, field_validator

from packages.llm.errors import DailyBudgetExceededError, SceneDegradedError
from packages.shared.logging import get_logger
from packages.shared.observability import REGISTRY

logger = get_logger(__name__)

_DAY_SECONDS = 60 * 60 * 24

#: 运营手动干预档位的 Redis 键（值为 "1"~"4"；TTL 到当日 UTC+8 24 点）。
OVERRIDE_KEY = "llm:budget:override"

#: blocked 名单中的通配符：拒绝所有非 critical 场景（仅 block_all 档配置 allow 名单）。
BLOCK_ALL: Literal["*"] = "*"

#: 内置默认四档，与 ``Settings.llm_budget_tiers`` 的默认值保持一致（架构 §13.3）。
_DEFAULT_TIERS: list[dict[str, object]] = [
    {"threshold": 0.6, "name": "warn", "blocked": []},
    {"threshold": 0.8, "name": "degrade_secondary", "blocked": ["hd_export", "report_regen"]},
    {"threshold": 0.9, "name": "degrade_generation", "blocked": ["hairstyle", "tryon"]},
    {"threshold": 1.0, "name": "block_all", "blocked": BLOCK_ALL, "allow": ["agent_chat"]},
]

# 档位是全局状态（整个平台只有一个水位），挂到共享 REGISTRY 上随 /metrics 暴露。
LLM_BUDGET_TIER = Gauge(
    "llm_budget_tier",
    "当前生效的日预算降级档位（0=未触发，1~4 递进；手动 override 优先）",
    registry=REGISTRY,
)
LLM_BUDGET_BLOCKED = Counter(
    "llm_budget_blocked_total",
    "被预算降级档位拒绝的网关请求数",
    labelnames=("scene",),
    registry=REGISTRY,
)


class BudgetTier(BaseModel):
    """一档降级配置（``LLM_BUDGET_TIERS_JSON`` 的元素，架构 §13.3）。

    ``blocked`` 里的 ``"*"`` 表示该档拒绝全部非 critical 场景（配置允许写成
    裸字符串，这里统一归一成列表）；``allow`` 仅 block_all 档有意义。
    """

    model_config = ConfigDict(frozen=True)

    threshold: float
    name: str
    blocked: list[str] = Field(default_factory=list)
    allow: list[str] = Field(default_factory=list)

    @field_validator("blocked", "allow", mode="before")
    @classmethod
    def _wrap_bare_string(cls, value: object) -> object:
        # Settings 默认值里 "blocked": "*" 是裸字符串写法，包成列表让类型恒为 list[str]。
        if isinstance(value, str):
            return [value]
        return value

    @property
    def blocked_set(self) -> frozenset[str]:
        return frozenset(self.blocked)

    @property
    def allow_set(self) -> frozenset[str]:
        return frozenset(self.allow)

    def blocks_all(self) -> bool:
        return BLOCK_ALL in self.blocked


@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: ...

    async def set(self, key: str, value: str, ttl_seconds: int) -> None: ...

    async def delete(self, key: str) -> None: ...


class InMemoryCostCounter:
    """Process-local counter. Correct only for a single process — tests only."""

    def __init__(self) -> None:
        self._values: dict[str, Decimal] = {}
        self._strings: dict[str, str] = {}

    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:
        if key in self._strings:
            # 非法字符串在这里抛 InvalidOperation，与 Redis 实现的行为一致，
            # 调用方（override 读取）按「值非法则忽略」处理。
            return Decimal(self._strings[key])
        return self._values.get(key, Decimal("0"))

    async def set(self, key: str, value: str, ttl_seconds: int) -> None:
        self._strings[key] = value

    async def delete(self, key: str) -> None:
        self._strings.pop(key, None)
        self._values.pop(key, None)


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")

    async def set(self, key: str, value: str, ttl_seconds: int) -> None:
        await self._redis.set(key, value, ex=ttl_seconds)  # type: ignore[attr-defined]

    async def delete(self, key: str) -> None:
        await self._redis.delete(key)  # type: ignore[attr-defined]


class DailyBudgetGuard:
    """Platform-wide daily AI budget with four-tier graded degradation (§13.3)."""

    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
        # 档位配置（LLM_BUDGET_TIERS_JSON）。None = 内置默认四档。解析失败在装配期
        # 就炸出来：一份写坏的降级配置不应该等到平台烧超预算时才被发现。
        self._tiers = [BudgetTier.model_validate(t) for t in (tiers if tiers else _DEFAULT_TIERS)]
        self._tiers.sort(key=lambda t: t.threshold)
        self._thresholds = [Decimal(str(t.threshold)) for t in self._tiers]
        # 上升沿告警的实例内游标：避免每个请求都刷一遍 budget_tier_raised。
        self._last_tier = 0

    @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())

    # -- tiered degradation ------------------------------------------------

    async def current_tier(self) -> int:
        """当前生效档位（0=未触发，1~4 递进）。手动 override 优先于成本水位。"""
        if not self._enabled:
            return 0
        override = await self._override_tier()
        if override is not None:
            return override
        return self._tier_for_spend(await self.spent_today())

    async def tier_name(self) -> str:
        """当前档位名（warn/degrade_secondary/degrade_generation/block_all）。"""
        tier = await self.current_tier()
        return self._tiers[tier - 1].name if tier else ""

    async def blocked_scenes(self) -> frozenset[str] | Literal["*"]:
        """当前档位**累计**拒绝的场景名单；``"*"`` 表示拒绝全部非 critical。"""
        return self._blocked_scenes_at(await self.current_tier())

    async def ensure_scene_allowed(self, *, scene: str, critical: bool = False) -> None:
        """Raise when the current tier rejects ``scene``.

        - critical 场景全档豁免（但照常计入成本水位）；
        - tier 2/3 且 scene 命中累计 blocked 名单 → :class:`SceneDegradedError`；
        - tier 4（blocked="*"）且 scene 不在 allow 名单 →
          :class:`DailyBudgetExceededError`——保持 100% 档的既有语义与 60003
          错误码不变（现有测试与告警文案都依赖它）。
        """
        if not self._enabled:
            return
        override = await self._override_tier()
        spent = await self.spent_today()
        tier = override if override is not None else self._tier_for_spend(spent)
        self._observe_tier(tier, spent=spent)
        if critical or tier == 0:
            return

        blocked = self._blocked_scenes_at(tier)
        if blocked == BLOCK_ALL:
            if scene not in self._tiers[tier - 1].allow_set:
                LLM_BUDGET_BLOCKED.labels(scene).inc()
                raise DailyBudgetExceededError(str(spent), str(self._budget), scene)
            return
        if scene in blocked:
            LLM_BUDGET_BLOCKED.labels(scene).inc()
            raise SceneDegradedError(scene, self._tiers[tier - 1].name)

    async def ensure_within_budget(self, *, scene: str, critical: bool = False) -> None:
        """兼容入口：与 :meth:`ensure_scene_allowed` 等价（存量调用方在用）。"""
        await self.ensure_scene_allowed(scene=scene, critical=critical)

    # -- manual override ---------------------------------------------------

    async def set_override(self, tier: int) -> None:
        """运营/演练：强制指定当前档位，TTL 到当日 UTC+8 24 点。"""
        if not 1 <= tier <= len(self._tiers):
            msg = f"override tier must be in 1..{len(self._tiers)}, got {tier}"
            raise ValueError(msg)
        await self._counter.set(OVERRIDE_KEY, str(tier), self._seconds_until_day_end())

    async def clear_override(self) -> None:
        """解除 override，回落成本水位判定。"""
        await self._counter.delete(OVERRIDE_KEY)

    # -- internals ----------------------------------------------------------

    async def _override_tier(self) -> int | None:
        """Read the manual override; unreadable/invalid values are ignored."""
        try:
            value = await self._counter.get(OVERRIDE_KEY)
            tier = int(value)
        except Exception:
            # 键不存在时计数器返回 0（→ 下方范围校验拒绝）；这里兜底的是
            # 值非法（如 "abc"）或存储故障：override 是运营工具，绝不能
            # 因为它把请求路径打挂。
            return None
        if 1 <= tier <= len(self._tiers):
            return tier
        return None

    def _tier_for_spend(self, spent: Decimal) -> int:
        """水位判定：取「最后一个被跨过阈值」的档位（阈值已升序）。"""
        tier = 0
        for index, threshold in enumerate(self._thresholds, start=1):
            if spent >= threshold * self._budget:
                tier = index
        return tier

    def _blocked_scenes_at(self, tier: int) -> frozenset[str] | Literal["*"]:
        """档位越高拒绝越多：取 1..tier 各档 blocked 的并集（「追加拒绝」）。"""
        scenes: set[str] = set()
        for spec in self._tiers[:tier]:
            if spec.blocks_all():
                return BLOCK_ALL
            scenes.update(spec.blocked_set)
        return frozenset(scenes)

    def _observe_tier(self, tier: int, *, spent: Decimal) -> None:
        LLM_BUDGET_TIER.set(tier)
        if tier <= self._last_tier:
            return
        self._last_tier = tier
        if tier == 0:
            return
        # 只在上升沿告警一次：水位在 90% 徘徊时每个请求都打日志等于没有告警。
        logger.warning(
            "budget_tier_raised",
            tier=tier,
            tier_name=self._tiers[tier - 1].name,
            spent=str(spent),
            budget=str(self._budget),
        )

    def _seconds_until_day_end(self) -> int:
        """到当日 UTC+8 24 点的秒数（预算日切换即失效）。"""
        local = self._now().astimezone(self._tz)
        day_end = local.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
        return max(int((day_end - local).total_seconds()), 1)

    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)