meida_backend/packages/llm/cost.py

"""Cost attribution (架构文档 §13, FR-BILL-08, NFR-COST-01).

**每一次付费调用都必须落 ``ob_ai_call``，并可归因到用户与功能。**
V1 不收费让这张表的地位更高，不是更低 —— 它是 V2 定价的唯一数据来源，
不允许「先上线后补埋点」。

Rates are configuration, never constants at a call site: providers re-price
without warning and V1's whole cost model depends on these numbers being
correct at the time of the call.
"""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass, replace
from datetime import UTC, datetime, time, timedelta, timezone
from decimal import ROUND_HALF_UP, Decimal
from typing import Protocol, runtime_checkable

from packages.llm.types import Usage

_MILLION = Decimal("1000000")
_QUANT = Decimal("0.000001")

BEIJING = timezone(timedelta(hours=8))

# DeepSeek 自 2026-08-17 00:00（北京时间）起改为峰谷定价，高峰单价是空闲时段的 2 倍。
# 官方口径的高峰时段：北京时间 09:00-12:00 与 14:00-18:00。
PEAK_WINDOWS: tuple[tuple[time, time], ...] = (
    (time(9, 0), time(12, 0)),
    (time(14, 0), time(18, 0)),
)


@dataclass(frozen=True, slots=True)
class ModelRate:
    """Cents per 1M tokens. Cached input is billed at a lower rate by DeepSeek."""

    input_cents_per_mtok: Decimal
    cached_input_cents_per_mtok: Decimal
    output_cents_per_mtok: Decimal

    @classmethod
    def from_mapping(cls, data: Mapping[str, Decimal | str | int | float]) -> ModelRate:
        return cls(
            input_cents_per_mtok=Decimal(str(data["input_cents_per_mtok"])),
            cached_input_cents_per_mtok=Decimal(str(data["cached_input_cents_per_mtok"])),
            output_cents_per_mtok=Decimal(str(data["output_cents_per_mtok"])),
        )


class RateTable:
    """model → :class:`ModelRate`, with an explicit fallback so an unknown model
    is still costed (a zero-cost blind spot is worse than an approximate one)."""

    def __init__(self, rates: dict[str, ModelRate], default: ModelRate) -> None:
        self._rates = dict(rates)
        self._default = default

    def for_model(self, model: str) -> ModelRate:
        return self._rates.get(model, self._default)

    def known_models(self) -> tuple[str, ...]:
        return tuple(sorted(self._rates))


def is_peak(at: datetime | None = None) -> bool:
    """调用时刻是否落在供应商的高峰时段（按北京时间判定）。"""
    moment = (at or datetime.now(UTC)).astimezone(BEIJING).timetz()
    return any(start <= moment.replace(tzinfo=None) < end for start, end in PEAK_WINDOWS)


class PeakAwareRateTable:
    """按调用时刻在「高峰 / 空闲」两套单价之间选择。

    为什么不能只存一套单价：高峰价是空闲价的 2 倍，而报告生成这类任务大多发生在
    白天高峰时段。用单一均价记账会让 ``ob_ai_call`` 的成本与账单对不上，
    而这张表是 V2 定价的唯一依据 —— 记错等于把定价建在错的数字上。

    时段判断放在算成本的这一刻，不是事后按当前价重算：跨时段重算会把历史成本改掉。
    """

    def __init__(
        self, off_peak: RateTable, peak: RateTable | None = None, *, enabled: bool = True
    ) -> None:
        self._off_peak = off_peak
        self._peak = peak or off_peak
        self._enabled = enabled

    def for_model(self, model: str, at: datetime | None = None) -> ModelRate:
        table = self._peak if self._enabled and is_peak(at) else self._off_peak
        return table.for_model(model)

    def known_models(self) -> tuple[str, ...]:
        return self._off_peak.known_models()


def compute_cost_cents(usage: Usage, rate: ModelRate) -> Decimal:
    """Cost of one call, in cents, with cache-hit tokens billed separately."""
    billable_input = Decimal(usage.billable_input) * rate.input_cents_per_mtok
    cached_input = Decimal(usage.tokens_cached) * rate.cached_input_cents_per_mtok
    output = Decimal(usage.tokens_out) * rate.output_cents_per_mtok
    total = (billable_input + cached_input + output) / _MILLION
    return total.quantize(_QUANT, rounding=ROUND_HALF_UP)


@runtime_checkable
class RateResolver(Protocol):
    """网关只需要「给我这个模型此刻的单价」，不关心背后是一套还是两套价目表。"""

    def for_model(self, model: str) -> ModelRate: ...


@dataclass(frozen=True, slots=True)
class AICallRecord:
    """One row of ``ob_ai_call``. Written for successes *and* failures."""

    call_id: int
    scene: str
    provider: str
    model: str
    success: bool
    usage: Usage
    latency_ms: int
    cost_cents: Decimal
    user_id: int | None = None
    prompt_id: str | None = None
    prompt_version: str | None = None
    error_code: str | None = None
    error_message: str | None = None
    ref_type: str | None = None
    ref_id: str | None = None
    created_at: datetime | None = None

    def with_defaults(self) -> AICallRecord:
        if self.created_at is not None:
            return self
        # 本类是 slots dataclass，没有 __dict__，只能用 replace 复制
        return replace(self, created_at=datetime.now(UTC))


@runtime_checkable
class AICallRecorder(Protocol):
    """Persistence port. The gateway never touches the ORM directly."""

    async def record(self, record: AICallRecord) -> None: ...


class NullAICallRecorder:
    """Explicitly opts out of cost attribution. Only valid in unit tests."""

    async def record(self, record: AICallRecord) -> None:
        return None


class InMemoryAICallRecorder:
    """Test double that keeps every record for assertions."""

    def __init__(self) -> None:
        self.records: list[AICallRecord] = []

    async def record(self, record: AICallRecord) -> None:
        self.records.append(record)


class CompositeAICallRecorder:
    """Fan out to several recorders (e.g. PostgreSQL + Prometheus).

    A failure in one recorder must not lose the others, and must never fail the
    user's request — cost bookkeeping is mandatory but not user-facing.
    """

    def __init__(self, *recorders: AICallRecorder) -> None:
        self._recorders = recorders

    async def record(self, record: AICallRecord) -> None:
        for recorder in self._recorders:
            await recorder.record(record)