use serde::{Deserialize, Serialize};
use crate::catalog::{CatalogOffering, CatalogSource};
use crate::models_dev::ModelsDevCost;
use crate::route::PricingSku;
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Currency {
#[default]
Usd,
Cny,
Other(String),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "source", rename_all = "snake_case")]
pub enum PricingProvenance {
ModelsDevBundled,
ProviderLive,
ProviderDocs,
UserOverride,
Unknown,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct TokenUsage {
pub input: u64,
pub output: u64,
pub cache_read: u64,
pub cache_write: u64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct OfferingPricing {
pub provider: String,
pub wire_model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub canonical_model: Option<String>,
pub currency: Currency,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_per_million: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_per_million: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_per_million: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_write_per_million: Option<f64>,
pub provenance: PricingProvenance,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub effective_at: Option<u64>,
}
impl OfferingPricing {
#[must_use]
pub fn from_catalog_offering(offering: &CatalogOffering) -> Option<Self> {
let cost = offering.cost.as_ref()?;
if cost.input.is_none()
&& cost.output.is_none()
&& cost.cache_read.is_none()
&& cost.cache_write.is_none()
{
return None;
}
Some(Self {
provider: offering.provider.clone(),
wire_model_id: offering.wire_model_id.clone(),
canonical_model: offering.canonical_model.clone(),
currency: Currency::Usd,
input_per_million: cost.input,
output_per_million: cost.output,
cache_read_per_million: cost.cache_read,
cache_write_per_million: cost.cache_write,
provenance: provenance_from_source(&offering.source),
effective_at: effective_at_from_source(&offering.source),
})
}
#[must_use]
pub fn has_any_price(&self) -> bool {
self.input_per_million.is_some()
|| self.output_per_million.is_some()
|| self.cache_read_per_million.is_some()
|| self.cache_write_per_million.is_some()
}
#[must_use]
pub fn is_stale(&self, now_unix: u64, max_age_secs: u64) -> bool {
match self.effective_at {
Some(t) => now_unix.saturating_sub(t) >= max_age_secs,
None => false,
}
}
#[must_use]
pub fn estimate_cost(&self, usage: &TokenUsage) -> Option<f64> {
let mut total = 0.0_f64;
for (tokens, price) in [
(usage.input, self.input_per_million),
(usage.output, self.output_per_million),
(usage.cache_read, self.cache_read_per_million),
(usage.cache_write, self.cache_write_per_million),
] {
if tokens > 0 {
let price = price?;
total += (tokens as f64 / 1_000_000.0) * price;
}
}
Some(total)
}
#[must_use]
pub fn to_route_sku(&self) -> PricingSku {
if self.input_per_million.is_none() && self.output_per_million.is_none() {
return PricingSku::UnknownOrStale;
}
PricingSku::Token {
input_per_mtok: self.input_per_million,
output_per_mtok: self.output_per_million,
}
}
}
#[must_use]
pub fn route_pricing_sku(offering: &CatalogOffering) -> PricingSku {
OfferingPricing::from_catalog_offering(offering)
.map_or(PricingSku::UnknownOrStale, |pricing| pricing.to_route_sku())
}
#[must_use]
pub(crate) fn route_pricing_sku_from_cost(cost: Option<&ModelsDevCost>) -> PricingSku {
let Some(cost) = cost else {
return PricingSku::UnknownOrStale;
};
if cost.input.is_none() && cost.output.is_none() {
return PricingSku::UnknownOrStale;
}
PricingSku::Token {
input_per_mtok: cost.input,
output_per_mtok: cost.output,
}
}
fn provenance_from_source(source: &CatalogSource) -> PricingProvenance {
match source {
CatalogSource::Bundled => PricingProvenance::ModelsDevBundled,
CatalogSource::Live { .. } => PricingProvenance::ProviderLive,
CatalogSource::UserOverride => PricingProvenance::UserOverride,
}
}
fn effective_at_from_source(source: &CatalogSource) -> Option<u64> {
match source {
CatalogSource::Live { fetched_at, .. } => Some(*fetched_at),
CatalogSource::Bundled | CatalogSource::UserOverride => None,
}
}
#[cfg(test)]
mod tests;