Token 计费、吞吐与成本意识：先做预算再谈规模

"""按输入输出Token数量计算费用，并在请求前预留预算。"""

from dataclasses import dataclass
from decimal import Decimal


MILLION = Decimal("1000000")


def _token_count(name: str, value: int) -> int:
    if type(value) is not int:
        raise TypeError(f"{name} must be an integer")
    if value < 0:
        raise ValueError(f"{name} must be non-negative")
    return value


def _money(name: str, value: Decimal, *, allow_zero: bool = True) -> Decimal:
    if not isinstance(value, Decimal):
        raise TypeError(f"{name} must be Decimal")
    if not value.is_finite():
        raise ValueError(f"{name} must be finite")
    if value < 0 or (not allow_zero and value == 0):
        raise ValueError(f"{name} must be positive")
    return value


@dataclass(frozen=True)
class Pricing:
    currency: str
    input_cache_hit_per_million: Decimal
    input_cache_miss_per_million: Decimal
    output_per_million: Decimal

    def __post_init__(self) -> None:
        if not isinstance(self.currency, str) or not self.currency.strip():
            raise ValueError("currency is required")
        _money("input_cache_hit_per_million", self.input_cache_hit_per_million)
        _money("input_cache_miss_per_million", self.input_cache_miss_per_million)
        _money("output_per_million", self.output_per_million)


@dataclass(frozen=True)
class Usage:
    cache_hit_input: int
    cache_miss_input: int
    output: int

    def __post_init__(self) -> None:
        _token_count("cache_hit_input", self.cache_hit_input)
        _token_count("cache_miss_input", self.cache_miss_input)
        _token_count("output", self.output)


def actual_cost(usage: Usage, pricing: Pricing) -> Decimal:
    """按供应方报告的实际用量计算这次请求费用。"""
    return (
        Decimal(usage.cache_hit_input) * pricing.input_cache_hit_per_million
        + Decimal(usage.cache_miss_input) * pricing.input_cache_miss_per_million
        + Decimal(usage.output) * pricing.output_per_million
    ) / MILLION


def reserve_cost(max_input_tokens: int, max_output_tokens: int, pricing: Pricing) -> Decimal:
    """按输入全部未命中缓存、输出达到上限来保守预留。"""
    input_count = _token_count("max_input_tokens", max_input_tokens)
    output_count = _token_count("max_output_tokens", max_output_tokens)
    return (
        Decimal(input_count) * pricing.input_cache_miss_per_million
        + Decimal(output_count) * pricing.output_per_million
    ) / MILLION


class BudgetLedger:
    """用内存演示预留；生产环境需要持久化且原子更新的存储。"""

    def __init__(self, limit: Decimal, currency: str) -> None:
        self.limit = _money("limit", limit, allow_zero=False)
        if not isinstance(currency, str) or not currency.strip():
            raise ValueError("currency is required")
        self.currency = currency
        self.spent = Decimal("0")
        self._requests: dict[str, dict[str, Decimal | str]] = {}

    @property
    def reserved(self) -> Decimal:
        return sum(
            (record["amount"] for record in self._requests.values()
             if record["status"] in {"reserved", "unknown"}),
            start=Decimal("0"),
        )

    @property
    def available(self) -> Decimal:
        return self.limit - self.spent - self.reserved

    def reserve(self, request_id: str, amount: Decimal) -> Decimal:
        if not isinstance(request_id, str) or not request_id.strip():
            raise ValueError("request_id is required")
        amount = _money("amount", amount, allow_zero=False)
        existing = self._requests.get(request_id)
        if existing:
            if existing["status"] == "unknown":
                raise ValueError("reconcile required before another reservation")
            if existing["amount"] != amount:
                raise ValueError("request_id already has a different amount")
            return existing["amount"]
        if amount > self.available:
            raise ValueError("budget is insufficient")
        self._requests[request_id] = {"status": "reserved", "amount": amount}
        return amount

    def settle(self, request_id: str, actual: Decimal) -> Decimal:
        actual = _money("actual", actual)
        record = self._requests.get(request_id)
        if not record:
            raise ValueError("reservation not found")
        if record["status"] == "settled":
            if record["actual"] != actual:
                raise ValueError("request_id already settled with a different amount")
            return record["actual"]
        if record["status"] == "unknown":
            raise ValueError("reconcile required before settlement")
        if actual > record["amount"]:
            raise ValueError("actual cost exceeds reservation")
        record["status"] = "settled"
        record["actual"] = actual
        self.spent += actual
        return actual

    def mark_unknown(self, request_id: str) -> None:
        record = self._requests.get(request_id)
        if not record:
            raise ValueError("reservation not found")
        if record["status"] != "reserved":
            raise ValueError("only a reserved request can become unknown")
        record["status"] = "unknown"

    def reconcile(self, request_id: str, confirmed_actual: Decimal) -> Decimal:
        """独立确认实际金额后，结清这条结果未知的请求。"""
        actual = _money("confirmed_actual", confirmed_actual)
        record = self._requests.get(request_id)
        if not record:
            raise ValueError("reservation not found")
        if record["status"] == "settled":
            if record["actual"] != actual:
                raise ValueError("request_id already settled with a different amount")
            return record["actual"]
        if record["status"] != "unknown":
            raise ValueError("only an unknown request can be reconciled")
        if actual > record["amount"]:
            raise ValueError("actual cost exceeds reservation")
        record["status"] = "settled"
        record["actual"] = actual
        self.spent += actual
        return actual


if __name__ == "__main__":
    usage = Usage(cache_hit_input=200_000, cache_miss_input=300_000, output=50_000)
    estimate = reserve_cost(500_000, 100_000, PRICING := Pricing(
        "CNY", Decimal("0.20"), Decimal("2.00"), Decimal("8.00")
    ))
    observed = actual_cost(usage, PRICING)
    ledger = BudgetLedger(limit=Decimal("100.00"), currency="CNY")
    ledger.reserve("demo-001", estimate)
    ledger.settle("demo-001", observed)
    print(f"reserved={estimate} {ledger.currency}")
    print(f"actual={observed} {ledger.currency}")
    print(f"available={ledger.available} {ledger.currency}")

# 核对Token 计费、吞吐与成本意识的正常结果和错误分支；断言失败时会显示对应的测试。
import unittest
from decimal import Decimal

from cost_guard import BudgetLedger, Pricing, Usage, actual_cost, reserve_cost


PRICING = Pricing(
    currency="CNY",
    input_cache_hit_per_million=Decimal("0.20"),
    input_cache_miss_per_million=Decimal("2.00"),
    output_per_million=Decimal("8.00"),
)


class CostFormulaTests(unittest.TestCase):
    def test_cost_separates_hit_miss_and_output_tokens(self):
        usage = Usage(cache_hit_input=200_000, cache_miss_input=300_000, output=50_000)
        self.assertEqual(actual_cost(usage, PRICING), Decimal("1.04"))

    def test_reservation_uses_cache_miss_and_maximum_output(self):
        self.assertEqual(reserve_cost(500_000, 100_000, PRICING), Decimal("1.80"))

    def test_boolean_negative_and_inconsistent_counts_are_rejected(self):
        invalid = [
            lambda: Usage(cache_hit_input=True, cache_miss_input=0, output=0),
            lambda: Usage(cache_hit_input=-1, cache_miss_input=0, output=0),
            lambda: reserve_cost(10, -1, PRICING),
            lambda: Pricing("CNY", Decimal("NaN"), Decimal("1"), Decimal("1")),
        ]
        for operation in invalid:
            with self.subTest(operation=operation), self.assertRaises((TypeError, ValueError)):
                operation()


class BudgetLedgerTests(unittest.TestCase):
    def test_unknown_reconciles_once_from_confirmed_amount(self):
        ledger = BudgetLedger(Decimal("100"), "CNY")
        ledger.reserve("req-1", Decimal("1.80"))
        ledger.mark_unknown("req-1")
        ledger.reconcile("req-1", Decimal("1.04"))
        ledger.reconcile("req-1", Decimal("1.04"))
        self.assertEqual(ledger.spent, Decimal("1.04"))
        self.assertEqual(ledger.available, Decimal("98.96"))
        with self.assertRaises(ValueError):
            ledger.reconcile("req-1", Decimal("1.05"))

    def test_reconcile_rejects_unconfirmed_state_and_over_reservation(self):
        ledger = BudgetLedger(Decimal("100"), "CNY")
        ledger.reserve("req-1", Decimal("1.80"))
        with self.assertRaises(ValueError):
            ledger.reconcile("req-1", Decimal("1.04"))
        ledger.mark_unknown("req-1")
        with self.assertRaises(ValueError):
            ledger.reconcile("req-1", Decimal("2"))
        self.assertEqual(ledger.available, Decimal("98.20"))

    def test_duplicate_reservation_is_idempotent(self):
        ledger = BudgetLedger(limit=Decimal("100.00"), currency="CNY")
        first = ledger.reserve("req-1", Decimal("1.80"))
        second = ledger.reserve("req-1", Decimal("1.80"))
        self.assertEqual(first, second)
        self.assertEqual(ledger.reserved, Decimal("1.80"))

    def test_over_budget_is_rejected_before_external_call(self):
        ledger = BudgetLedger(limit=Decimal("1.00"), currency="CNY")
        with self.assertRaisesRegex(ValueError, "budget"):
            ledger.reserve("req-1", Decimal("1.80"))
        self.assertEqual(ledger.reserved, Decimal("0"))

    def test_settlement_charges_actual_and_releases_unused_reservation(self):
        ledger = BudgetLedger(limit=Decimal("100.00"), currency="CNY")
        ledger.reserve("req-1", Decimal("1.80"))
        ledger.settle("req-1", Decimal("1.04"))
        self.assertEqual(ledger.spent, Decimal("1.04"))
        self.assertEqual(ledger.reserved, Decimal("0.00"))
        self.assertEqual(ledger.available, Decimal("98.96"))

    def test_unknown_result_keeps_reservation(self):
        ledger = BudgetLedger(limit=Decimal("100.00"), currency="CNY")
        ledger.reserve("req-1", Decimal("1.80"))
        ledger.mark_unknown("req-1")
        self.assertEqual(ledger.reserved, Decimal("1.80"))
        self.assertEqual(ledger.available, Decimal("98.20"))
        with self.assertRaisesRegex(ValueError, "reconcile"):
            ledger.reserve("req-1", Decimal("1.80"))

    def test_actual_cost_above_reservation_requires_reconciliation(self):
        ledger = BudgetLedger(limit=Decimal("100.00"), currency="CNY")
        ledger.reserve("req-1", Decimal("1.80"))
        with self.assertRaisesRegex(ValueError, "reservation"):
            ledger.settle("req-1", Decimal("2.10"))
        self.assertEqual(ledger.spent, Decimal("0"))
        self.assertEqual(ledger.reserved, Decimal("1.80"))

    def test_duplicate_settlement_does_not_charge_twice(self):
        ledger = BudgetLedger(limit=Decimal("100.00"), currency="CNY")
        ledger.reserve("req-1", Decimal("1.80"))
        first = ledger.settle("req-1", Decimal("1.04"))
        second = ledger.settle("req-1", Decimal("1.04"))
        self.assertEqual(first, second)
        self.assertEqual(ledger.spent, Decimal("1.04"))


if __name__ == "__main__":
    unittest.main()