Cache Savings Calculator — Python source

"""cache_savings — uncached vs prompt-cached LLM cost comparison.

Language: Python (3.9+, standard library only)
Source:   CosmoDev polyglot showcase port of the Cache Savings Calculator
          tool, ported from src/lib/cacheSavings.ts (the canonical
          TypeScript implementation).
Tool:     https://dev.cosmolabs.org/tools/cache-savings-calculator
License:  display source — part of CosmoDev's polyglot tool pages.

Design goals:
  - Pure + deterministic; never raises.
  - Functionally equivalent to the TS reference: same inputs -> same outputs.
  - Self-contained: stdlib only (no pip packages).

The TS original takes a full ``AiModel`` record but reads only its four
pricing rates, so this port narrows the parameter to exactly those fields
(snake_cased per Python convention: ``inputPerM`` -> ``input_per_m`` etc.).
Any ``None`` rate makes every output ``None`` — the caller renders an
explanatory empty state instead of partial math. All rates are per-1M-token
USD, mirroring the cost conventions of llmCost.ts.

Reference vectors (fixture model 10 / 50 / 1 / 12.5 — see
cacheSavings.test.ts, the lock-step contract every port mirrors):
  10k in / 1k out / 5 hits -> uncached 0.75, cached 0.425, savings 0.325,
  43.333...% saved, break-even 13 hits. At 1 hit caching LOSES 0.035 (an
  honest negative saving). hits < 1 counts as 1. Zero tokens -> zero costs
  with 0%. cache_read 0 -> break_even_hits None (write premium never repaid).
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

__all__ = ["ModelRates", "CacheInput", "CacheMath", "cache_math"]


@dataclass(frozen=True)
class ModelRates:
    """The four per-1M-token USD pricing rates ``cache_math`` reads from the
    TS ``AiModel`` record."""

    input_per_m: Optional[float]
    """Uncached prompt (input) rate, USD per 1M tokens."""

    output_per_m: Optional[float]
    """Completion (output) rate, USD per 1M tokens."""

    cache_read_per_m: Optional[float]
    """Cached prompt read rate, USD per 1M tokens."""

    cache_write_per_m: Optional[float]
    """Cache write premium rate, USD per 1M tokens."""


@dataclass(frozen=True)
class CacheInput:
    """Request shape (TS ``CacheInput``)."""

    prompt_tokens: float
    """Prompt (input) tokens per request."""

    output_tokens: float
    """Completion (output) tokens per request."""

    hits: float
    """Requests reusing the cached prompt. Values < 1 are treated as 1."""


@dataclass(frozen=True)
class CacheMath:
    """Result shape. Any missing rate nulls every field."""

    uncached: Optional[float]
    """hits × (prompt·in$/M + output·out$/M) / 1e6."""

    cached: Optional[float]
    """(prompt·write$/M + hits × (prompt·read$/M + output·out$/M)) / 1e6 —
    one cache write, ``hits`` cache reads, output billed every request."""

    savings: Optional[float]
    """uncached − cached (negative when caching costs more)."""

    savings_pct: Optional[float]
    """savings / uncached × 100; 0 when uncached is 0."""

    break_even_hits: Optional[int]
    """ceil(write$/M / read$/M) when read$/M > 0 — cache hits needed for
    cumulative READ spend to equal ONE write premium; None otherwise."""


def cache_math(model: ModelRates, inp: CacheInput) -> CacheMath:
    """Compare uncached vs prompt-cached cost for one model. Never raises;
    an unpriced rate yields the all-``None`` result so callers can render an
    explanatory empty state instead of partial math."""
    if (
        model.input_per_m is None
        or model.output_per_m is None
        or model.cache_read_per_m is None
        or model.cache_write_per_m is None
    ):
        return CacheMath(None, None, None, None, None)

    hits = max(1, inp.hits)
    in_t = inp.prompt_tokens
    out_t = inp.output_tokens

    # One cache write, `hits` cache reads; output tokens are billed on every request.
    uncached = (hits * (in_t * model.input_per_m + out_t * model.output_per_m)) / 1_000_000
    cached = (
        in_t * model.cache_write_per_m
        + hits * (in_t * model.cache_read_per_m + out_t * model.output_per_m)
    ) / 1_000_000
    savings = uncached - cached
    savings_pct = 0 if uncached == 0 else (savings / uncached) * 100
    break_even_hits = (
        math.ceil(model.cache_write_per_m / model.cache_read_per_m)
        if model.cache_read_per_m > 0
        else None
    )

    return CacheMath(uncached, cached, savings, savings_pct, break_even_hits)