klieo_runlog/pricing/

mod.rs

//! Pricing tables and per-(provider, model) USD rate lookup.
//!
//! Used by [`crate::projector::project_with_price_table`] to compute a
//! [`crate::types::Cost`] across all `Episode::LlmCall` steps in a run.
//!
//! ## Rate sources (2026-Q1 published list prices, USD)
//!
//! - Anthropic — <https://www.anthropic.com/pricing>
//! - OpenAI    — <https://openai.com/api/pricing/>
//! - Google    — <https://ai.google.dev/gemini-api/docs/pricing>
//! - Ollama    — local execution, zero list price
//!
//! These figures are list-price snapshots; revalidate against the upstream
//! pricing pages on the URLs above before billing or commercial routing
//! decisions. A periodic refresh job for live pricing is a deferred follow-up
//! — until then, callers override entries via [`PriceTable::set`].
//!
//! ## Model-id matching strategy
//!
//! Provider naming conventions vary widely: Anthropic ships date-suffixed
//! aliases (`claude-3-5-sonnet-20241022`), Ollama uses colon tags
//! (`qwen2.5:14b`), and others ship versioned -N family heads
//! (`claude-sonnet-4-6`). The price table keys family aliases with a literal
//! `-x` suffix (e.g. `claude-sonnet-4-x`), and [`PriceTable::lookup`]
//! resolves a `(provider, model)` query through this precedence chain:
//!
//! 1. **Exact match**: `(provider, model)` verbatim.
//! 2. **Date-stripped**: if `model` ends in an ISO date `-YYYY-MM-DD` or a
//!    compact date `-YYYYMMDD`, drop the date and re-resolve.
//! 3. **Colon-stripped**: if `model` contains `:` (Ollama tag form), try the
//!    head before the first `:` and re-resolve.
//! 4. **Trailing-digits rewrite**: if `model` ends in `-<digits>`, rewrite to
//!    `-x` (e.g. `claude-sonnet-4-6` → `claude-sonnet-4-x`).
//! 5. **Provider wildcard**: try `(provider, "*")` so the built-in Ollama `*`
//!    rate hits real Ollama model strings without callers passing `"*"`.
//! 6. **Miss**: return `None`.
//!
//! The trailing-digits rewrite is conservative: it only fires when the
//! trailing segment is purely digits, so `gpt-4o-mini` (alpha tail) does not
//! collapse to `gpt-4o`. Callers wanting non-default mappings populate the
//! table explicitly via [`PriceTable::set`].

mod model_alias;
mod tables;

use crate::types::Cost;
use model_alias::{family_alias, strip_trailing_date};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use thiserror::Error;

/// Token-count denominator for the `*_per_1k` rate convention: rates are
/// quoted per 1,000 tokens, so cost = (tokens / 1000) * rate.
const TOKENS_PER_RATE_UNIT: f64 = 1_000.0;

/// Validation failure for [`Rates`] construction.
///
/// Returned by [`Rates::try_new`] and [`PriceTable::set`] when a rate field
/// is non-finite (`NaN`, `+inf`, `-inf`) or negative. Negative or non-finite
/// rates would silently corrupt downstream USD totals — billing-grade input
/// validation rejects them at the boundary.
#[derive(Debug, Clone, Error, PartialEq)]
#[non_exhaustive]
pub enum RatesError {
    /// One of the rate fields was `NaN` or infinite.
    #[error("rate must be finite: prompt={prompt_usd_per_1k}, completion={completion_usd_per_1k}")]
    NotFinite {
        /// Offending prompt rate.
        prompt_usd_per_1k: f64,
        /// Offending completion rate.
        completion_usd_per_1k: f64,
    },
    /// One of the rate fields was negative.
    #[error(
        "rate must be non-negative: prompt={prompt_usd_per_1k}, completion={completion_usd_per_1k}"
    )]
    Negative {
        /// Offending prompt rate.
        prompt_usd_per_1k: f64,
        /// Offending completion rate.
        completion_usd_per_1k: f64,
    },
    /// A cache-tier rate was non-finite or negative.
    #[error(
        "cache rate must be finite and non-negative: \
         read={cache_read_usd_per_1k}, creation={cache_creation_usd_per_1k}"
    )]
    CacheRateInvalid {
        /// Offending cache-read rate.
        cache_read_usd_per_1k: f64,
        /// Offending cache-creation rate.
        cache_creation_usd_per_1k: f64,
    },
}

/// Wire-format mirror of [`Rates`] used to gate `Deserialize` through
/// [`Rates::try_new`].
///
/// JSON / serde deserialisation routes through `Rates::try_from(RatesUnchecked)`
/// so any out-of-range or non-finite input fails closed with a serde error
/// rather than landing as a poisoned [`Rates`] that would silently corrupt
/// signed audit-evidence totals.
#[derive(Debug, Clone, Copy, Deserialize)]
pub(crate) struct RatesUnchecked {
    pub(crate) prompt_usd_per_1k: f64,
    pub(crate) completion_usd_per_1k: f64,
    #[serde(default)]
    pub(crate) cache_read_usd_per_1k: f64,
    #[serde(default)]
    pub(crate) cache_creation_usd_per_1k: f64,
}

impl TryFrom<RatesUnchecked> for Rates {
    type Error = RatesError;

    fn try_from(value: RatesUnchecked) -> Result<Self, Self::Error> {
        let base = Rates::try_new(value.prompt_usd_per_1k, value.completion_usd_per_1k)?;
        let (read, creation) = (value.cache_read_usd_per_1k, value.cache_creation_usd_per_1k);
        if !read.is_finite() || !creation.is_finite() || read < 0.0 || creation < 0.0 {
            return Err(RatesError::CacheRateInvalid {
                cache_read_usd_per_1k: read,
                cache_creation_usd_per_1k: creation,
            });
        }
        Ok(base.with_cache_rates(read, creation))
    }
}

/// Per-1k-token USD rates for a single (provider, model) tier.
///
/// # Construction
///
/// Construct via [`Rates::new`] (const-context, debug-validated) or
/// [`Rates::try_new`] (runtime-validated, returns [`RatesError`]). Direct
/// struct-literal init is restricted to in-crate code paths via `pub(crate)`
/// fields so external callers cannot bypass the finite + non-negative
/// contract.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(try_from = "RatesUnchecked")]
pub struct Rates {
    /// USD cost per 1,000 prompt tokens.
    pub(crate) prompt_usd_per_1k: f64,
    /// USD cost per 1,000 completion tokens.
    pub(crate) completion_usd_per_1k: f64,
    /// USD cost per 1,000 cache-read tokens (defaults to 0 — providers without
    /// prompt caching, or tables that predate cache-aware pricing).
    pub(crate) cache_read_usd_per_1k: f64,
    /// USD cost per 1,000 cache-creation tokens (defaults to 0).
    pub(crate) cache_creation_usd_per_1k: f64,
}

impl Rates {
    /// Construct a rate pair without runtime validation.
    ///
    /// # Contract
    ///
    /// Both inputs MUST be finite (`!is_nan() && !is_infinite()`) and
    /// non-negative. Violations panic in debug builds via `debug_assert!`
    /// and are silently accepted in release builds for backward compat with
    /// the const-context callers in the built-in rate table.
    ///
    /// Untrusted callers should prefer [`Rates::try_new`], which returns a
    /// [`RatesError`] instead of relying on the debug-only assertion.
    pub const fn new(prompt_usd_per_1k: f64, completion_usd_per_1k: f64) -> Self {
        debug_assert!(
            prompt_usd_per_1k.is_finite() && completion_usd_per_1k.is_finite(),
            "Rates::new: both rates must be finite"
        );
        debug_assert!(
            prompt_usd_per_1k >= 0.0 && completion_usd_per_1k >= 0.0,
            "Rates::new: both rates must be non-negative"
        );
        Self {
            prompt_usd_per_1k,
            completion_usd_per_1k,
            cache_read_usd_per_1k: 0.0,
            cache_creation_usd_per_1k: 0.0,
        }
    }

    /// Attach cache-tier rates (read = served-from-cache, creation =
    /// written-to-cache). Both default to 0; set them where the provider bills
    /// cached prompt tokens at a distinct tier (e.g. Anthropic). Debug-validated
    /// finite + non-negative, like [`Rates::new`].
    #[must_use]
    pub const fn with_cache_rates(
        mut self,
        cache_read_usd_per_1k: f64,
        cache_creation_usd_per_1k: f64,
    ) -> Self {
        debug_assert!(
            cache_read_usd_per_1k.is_finite() && cache_creation_usd_per_1k.is_finite(),
            "Rates::with_cache_rates: both rates must be finite"
        );
        debug_assert!(
            cache_read_usd_per_1k >= 0.0 && cache_creation_usd_per_1k >= 0.0,
            "Rates::with_cache_rates: both rates must be non-negative"
        );
        self.cache_read_usd_per_1k = cache_read_usd_per_1k;
        self.cache_creation_usd_per_1k = cache_creation_usd_per_1k;
        self
    }

    /// Construct a rate pair with runtime validation.
    ///
    /// Rejects `NaN`, `±inf`, and negative inputs with [`RatesError`]. This
    /// is the constructor untrusted callers (config loaders, dynamic rate
    /// updates) should use.
    pub fn try_new(prompt_usd_per_1k: f64, completion_usd_per_1k: f64) -> Result<Self, RatesError> {
        if !prompt_usd_per_1k.is_finite() || !completion_usd_per_1k.is_finite() {
            return Err(RatesError::NotFinite {
                prompt_usd_per_1k,
                completion_usd_per_1k,
            });
        }
        if prompt_usd_per_1k < 0.0 || completion_usd_per_1k < 0.0 {
            return Err(RatesError::Negative {
                prompt_usd_per_1k,
                completion_usd_per_1k,
            });
        }
        Ok(Self {
            prompt_usd_per_1k,
            completion_usd_per_1k,
            cache_read_usd_per_1k: 0.0,
            cache_creation_usd_per_1k: 0.0,
        })
    }

    /// Compute the [`Cost`] for prompt + completion tokens at these rates (no
    /// cache tokens). Equivalent to [`Rates::cost_for_cached`] with both cache
    /// counts 0; `Cost::cache_usd` is 0.
    pub fn cost_for(&self, prompt_tokens: u32, completion_tokens: u32) -> Cost {
        self.cost_for_cached(prompt_tokens, completion_tokens, 0, 0)
    }

    /// Compute the [`Cost`] across all four token tiers — prompt, completion,
    /// cache-read, cache-creation — each at its own rate. `total_usd` is their
    /// sum. Token counts are `u32`; conversion to `f64` is lossless across the
    /// full `u32` range.
    pub fn cost_for_cached(
        &self,
        prompt_tokens: u32,
        completion_tokens: u32,
        cache_read_tokens: u32,
        cache_creation_tokens: u32,
    ) -> Cost {
        // Defense-in-depth: the constructors guarantee finite + non-negative,
        // and pub(crate) fields block external literal-init bypass — but if an
        // in-crate edit ever lands an invalid Rates, surface it loudly in dev.
        debug_assert!(
            self.prompt_usd_per_1k.is_finite() && self.prompt_usd_per_1k >= 0.0,
            "Rates::cost_for: prompt rate must be finite and non-negative"
        );
        debug_assert!(
            self.completion_usd_per_1k.is_finite() && self.completion_usd_per_1k >= 0.0,
            "Rates::cost_for: completion rate must be finite and non-negative"
        );
        let per_1k = |tokens: u32, rate: f64| (f64::from(tokens) / TOKENS_PER_RATE_UNIT) * rate;
        let prompt_usd = per_1k(prompt_tokens, self.prompt_usd_per_1k);
        let completion_usd = per_1k(completion_tokens, self.completion_usd_per_1k);
        let cache_usd = per_1k(cache_read_tokens, self.cache_read_usd_per_1k)
            + per_1k(cache_creation_tokens, self.cache_creation_usd_per_1k);
        Cost {
            cache_usd,
            prompt_usd,
            completion_usd,
            total_usd: prompt_usd + completion_usd + cache_usd,
        }
    }

    /// USD cost per 1,000 prompt tokens.
    pub fn prompt_usd_per_1k(&self) -> f64 {
        self.prompt_usd_per_1k
    }

    /// USD cost per 1,000 completion tokens.
    pub fn completion_usd_per_1k(&self) -> f64 {
        self.completion_usd_per_1k
    }

    /// USD cost per 1,000 cache-read tokens.
    pub fn cache_read_usd_per_1k(&self) -> f64 {
        self.cache_read_usd_per_1k
    }

    /// USD cost per 1,000 cache-creation tokens.
    pub fn cache_creation_usd_per_1k(&self) -> f64 {
        self.cache_creation_usd_per_1k
    }
}

/// `(provider, model)` keyed table of [`Rates`].
///
/// See [module docs](self) for the lookup-fallback rule.
#[derive(Debug, Clone, Default)]
pub struct PriceTable {
    rates: HashMap<(String, String), Rates>,
}

impl PriceTable {
    /// Construct an empty price table.
    pub fn new() -> Self {
        Self {
            rates: HashMap::new(),
        }
    }

    /// Construct a price table pre-populated with built-in 2026-Q1 list-price
    /// rates for the providers shipped with klieo.
    ///
    /// See the [module-level rate-source documentation](self) for the
    /// upstream pricing URLs.
    pub fn with_default_rates() -> Self {
        tables::with_default_rates()
    }

    /// Insert or overwrite the rates for a `(provider, model)` pair.
    ///
    /// Validates `rates` via the same finite + non-negative contract enforced
    /// by [`Rates::try_new`]; returns [`RatesError`] without mutating the
    /// table on failure.
    pub fn set(
        &mut self,
        provider: impl Into<String>,
        model: impl Into<String>,
        rates: Rates,
    ) -> Result<(), RatesError> {
        // Re-validate at the boundary: the const Rates::new debug-asserts but
        // is silent in release builds, and untrusted callers may have
        // bypassed that path entirely.
        let _ = Rates::try_new(rates.prompt_usd_per_1k, rates.completion_usd_per_1k)?;
        self.rates.insert((provider.into(), model.into()), rates);
        Ok(())
    }

    /// Look up the rates for a `(provider, model)` pair.
    ///
    /// Resolves through the precedence chain documented at the
    /// [module level](self): exact match → date-stripped → colon-stripped →
    /// trailing-digits rewrite → `(provider, "*")` wildcard.
    pub fn lookup(&self, provider: &str, model: &str) -> Option<Rates> {
        let provider = normalize_provider(provider);
        // 1. exact
        if let Some(r) = self.lookup_exact(provider, model) {
            return Some(r);
        }
        // 2. date-stripped — strip the date suffix and try the head, then the
        //    head with `-x` appended (matching the family-alias convention
        //    used by the built-in Anthropic tiers).
        if let Some(stripped) = strip_trailing_date(model) {
            if let Some(r) = self.lookup_exact(provider, &stripped) {
                return Some(r);
            }
            let family = format!("{stripped}-x");
            if let Some(r) = self.lookup_exact(provider, &family) {
                return Some(r);
            }
            // Recurse so a stripped form like `qwen2.5:14b-2024-11-20` (date
            // first, then colon tag) still falls through the rest of the
            // chain.
            if let Some(r) = self.lookup(provider, &stripped) {
                return Some(r);
            }
        }
        // 3. colon-stripped (Ollama tag)
        if let Some((head, _tag)) = model.split_once(':') {
            if let Some(r) = self.lookup_exact(provider, head) {
                return Some(r);
            }
            // Also try the family-alias rewrite on the head form.
            if let Some(family) = family_alias(head) {
                if let Some(r) = self.lookup_exact(provider, &family) {
                    return Some(r);
                }
            }
        }
        // 4. trailing-digits rewrite (-N → -x)
        if let Some(family) = family_alias(model) {
            if let Some(r) = self.lookup_exact(provider, &family) {
                return Some(r);
            }
        }
        // 5. provider wildcard
        self.lookup_exact(provider, "*")
    }

    fn lookup_exact(&self, provider: &str, model: &str) -> Option<Rates> {
        self.rates
            .get(&(provider.to_string(), model.to_string()))
            .copied()
    }
}

/// Canonicalize a client-reported provider name to its price-table key. LLM
/// clients name themselves by wire identity (`gemini`), but the built-in table
/// keys Google models under `google`.
fn normalize_provider(provider: &str) -> &str {
    match provider {
        "gemini" => "google",
        other => other,
    }
}

#[cfg(test)]
mod tests;