use std::collections::HashSet;
use std::path::PathBuf;
#[cfg(test)]
use std::sync::Mutex;
use std::sync::{Arc, LazyLock, RwLock};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use super::response::{AnnotatedLlmResponse, CostEstimate, CostSource, Usage};
const PRICING_CATALOG_VERSION: u32 = 1;
static ACTIVE_PRICING_RESOLVER: LazyLock<RwLock<Arc<PricingResolver>>> =
LazyLock::new(|| RwLock::new(Arc::new(PricingResolver::default())));
#[cfg(test)]
pub(crate) fn pricing_test_mutex() -> &'static Mutex<()> {
static PRICING_TEST_MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
&PRICING_TEST_MUTEX
}
#[derive(Debug, Error)]
pub enum PricingCatalogError {
#[error("invalid model pricing catalog JSON: {0}")]
Json(#[from] serde_json::Error),
#[error("duplicate model pricing alias '{model}'")]
DuplicateModelAlias {
model: String,
},
#[error("unsupported model pricing catalog version {version}")]
UnsupportedVersion {
version: u32,
},
#[error("model pricing entry {entry_index} has empty {field}")]
EmptyField {
entry_index: usize,
field: String,
},
#[error("model pricing entry {entry_index} has invalid {field}: {value}")]
InvalidRate {
entry_index: usize,
field: String,
value: f64,
},
#[error("could not read model pricing catalog file '{}': {source}", path.display())]
FileRead {
path: PathBuf,
source: std::io::Error,
},
#[error("model pricing resolver lock poisoned: {0}")]
LockPoisoned(String),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PricingCatalog {
pub version: u32,
pub entries: Vec<ModelPricing>,
}
impl PricingCatalog {
pub fn from_json_str(catalog_json: &str) -> Result<Self, PricingCatalogError> {
let catalog: Self = serde_json::from_str(catalog_json)?;
catalog.validate()?;
Ok(catalog)
}
#[must_use]
pub fn pricing_for_model(&self, model: &str) -> Option<ModelPricing> {
self.pricing_for(None, model)
}
#[must_use]
pub fn pricing_for(&self, provider: Option<&str>, model: &str) -> Option<ModelPricing> {
let model_keys = normalized_model_lookup_keys(provider, model);
if model_keys.is_empty() {
return None;
}
model_keys.iter().find_map(|model_key| {
self.entries
.iter()
.find(|entry| entry.matches_model(model_key))
.cloned()
})
}
fn validate(&self) -> Result<(), PricingCatalogError> {
if self.version != PRICING_CATALOG_VERSION {
return Err(PricingCatalogError::UnsupportedVersion {
version: self.version,
});
}
let mut seen = HashSet::new();
for (entry_index, entry) in self.entries.iter().enumerate() {
entry.validate(entry_index)?;
for model_key in entry.provider_model_keys() {
if !seen.insert(model_key.clone()) {
return Err(PricingCatalogError::DuplicateModelAlias { model: model_key });
}
}
}
Ok(())
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct PricingConfig {
#[serde(default)]
pub sources: Vec<PricingSourceConfig>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum PricingSourceConfig {
Inline {
catalog: PricingCatalog,
},
File {
path: PathBuf,
},
}
pub trait PricingSource: Send + Sync {
fn source_name(&self) -> &str;
fn load_catalog(&self) -> Result<Option<PricingCatalog>, PricingCatalogError>;
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct PricingResolver {
catalogs: Vec<PricingCatalog>,
}
impl PricingResolver {
#[must_use]
pub fn from_catalogs(catalogs: Vec<PricingCatalog>) -> Self {
Self { catalogs }
}
pub fn from_config(config: &PricingConfig) -> Result<Self, PricingCatalogError> {
let mut catalogs = Vec::new();
for source in &config.sources {
match source {
PricingSourceConfig::Inline { catalog } => {
catalog.validate()?;
catalogs.push(catalog.clone());
}
PricingSourceConfig::File { path } => {
let raw = std::fs::read_to_string(path).map_err(|source| {
PricingCatalogError::FileRead {
path: path.clone(),
source,
}
})?;
catalogs.push(PricingCatalog::from_json_str(&raw)?);
}
}
}
Ok(Self { catalogs })
}
pub fn from_sources(sources: Vec<Box<dyn PricingSource>>) -> Result<Self, PricingCatalogError> {
let mut catalogs = Vec::new();
for source in sources {
if let Some(catalog) = source.load_catalog()? {
catalog.validate()?;
catalogs.push(catalog);
}
}
Ok(Self { catalogs })
}
#[must_use]
pub fn pricing_for_model(&self, model: &str) -> Option<ModelPricing> {
self.pricing_for(None, model)
}
#[must_use]
pub fn pricing_for(&self, provider: Option<&str>, model: &str) -> Option<ModelPricing> {
self.catalogs
.iter()
.find_map(|catalog| catalog.pricing_for(provider, model))
}
#[must_use]
pub fn estimate_cost(&self, model: &str, usage: &Usage) -> Option<CostEstimate> {
self.estimate_cost_for_provider(None, model, usage)
}
#[must_use]
pub fn estimate_cost_for_provider(
&self,
provider: Option<&str>,
model: &str,
usage: &Usage,
) -> Option<CostEstimate> {
self.pricing_for(provider, model)
.and_then(|pricing| pricing.estimate_cost(usage))
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ModelPricing {
pub provider: String,
pub model_id: String,
#[serde(default)]
pub aliases: Vec<String>,
#[serde(default = "default_pricing_currency")]
pub currency: String,
#[serde(default)]
pub unit: PricingUnit,
#[serde(skip_serializing_if = "Option::is_none")]
pub rates: Option<TokenPricingRates>,
#[serde(skip_serializing_if = "Option::is_none")]
pub rate_schedule: Option<TokenRateSchedule>,
pub prompt_cache: PromptCachePricing,
pub pricing_as_of: String,
pub pricing_source: String,
}
impl ModelPricing {
#[must_use]
pub fn estimate_cost(&self, usage: &Usage) -> Option<CostEstimate> {
if self.unit != PricingUnit::PerToken {
return None;
}
let prompt_tokens = usage.prompt_tokens.unwrap_or(0);
let completion_tokens = usage.completion_tokens.unwrap_or(0);
let cache_read_tokens = usage.cache_read_tokens.unwrap_or(0);
let cache_write_tokens = usage.cache_write_tokens.unwrap_or(0);
let rates = self.rates_for_usage(usage)?;
if prompt_tokens == 0
&& completion_tokens == 0
&& cache_read_tokens == 0
&& cache_write_tokens == 0
{
return None;
}
let billable_prompt_tokens =
if self.prompt_cache.read_accounting == CacheReadAccounting::IncludedInPromptTokens {
prompt_tokens.saturating_sub(cache_read_tokens)
} else {
prompt_tokens
};
let input_cost = cost_component_if_nonzero(billable_prompt_tokens, rates.input_per_million);
let output_cost = cost_component_if_nonzero(completion_tokens, rates.output_per_million);
let cache_read_cost = rates
.cache_read_per_million
.and_then(|price| cost_component_if_nonzero(cache_read_tokens, price));
let cache_write_cost = rates
.cache_write_per_million
.and_then(|price| cost_component_if_nonzero(cache_write_tokens, price));
let total: f64 = [input_cost, output_cost, cache_read_cost, cache_write_cost]
.into_iter()
.flatten()
.sum();
Some(CostEstimate {
total: Some(round_cost_amount(total)),
currency: self.currency.clone(),
input: input_cost,
output: output_cost,
cache_read: cache_read_cost,
cache_write: cache_write_cost,
source: CostSource::ModelPricing,
pricing_provider: Some(self.provider.clone()),
pricing_model: Some(self.model_id.clone()),
pricing_as_of: Some(self.pricing_as_of.clone()),
pricing_source: Some(self.pricing_source.clone()),
})
}
fn rates_for_usage(&self, usage: &Usage) -> Option<TokenPricingRates> {
if let Some(schedule) = &self.rate_schedule {
return schedule.rates_for_usage(usage);
}
self.rates
}
fn matches_model(&self, lookup: &ModelLookupKey) -> bool {
if let Some(provider) = lookup.provider.as_deref()
&& normalized_provider_name(&self.provider) != provider
{
return false;
}
self.model_keys().any(|key| key == lookup.model)
}
fn model_keys(&self) -> impl Iterator<Item = String> + '_ {
std::iter::once(&self.model_id)
.chain(self.aliases.iter())
.map(|model| normalized_model_name(model))
.filter(|model| !model.is_empty())
}
fn provider_model_keys(&self) -> impl Iterator<Item = String> + '_ {
let provider = normalized_provider_name(&self.provider);
self.model_keys()
.map(move |model| format!("{provider}/{model}"))
}
fn validate(&self, entry_index: usize) -> Result<(), PricingCatalogError> {
validate_nonempty(entry_index, "provider", &self.provider)?;
validate_nonempty(entry_index, "model_id", &self.model_id)?;
validate_nonempty(entry_index, "currency", &self.currency)?;
validate_nonempty(entry_index, "pricing_as_of", &self.pricing_as_of)?;
validate_nonempty(entry_index, "pricing_source", &self.pricing_source)?;
if self.unit == PricingUnit::PerToken
&& self.rates.is_none()
&& self.rate_schedule.is_none()
{
return Err(PricingCatalogError::EmptyField {
entry_index,
field: "rates or rate_schedule".to_string(),
});
}
if let Some(rates) = &self.rates {
rates.validate(entry_index, "rates")?;
}
if let Some(schedule) = &self.rate_schedule {
schedule.validate(entry_index)?;
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum PricingUnit {
#[default]
PerToken,
PerRequest,
PerSecond,
GpuHour,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct TokenPricingRates {
pub input_per_million: f64,
pub output_per_million: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_per_million: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_write_per_million: Option<f64>,
}
impl TokenPricingRates {
fn validate(&self, entry_index: usize, field_prefix: &str) -> Result<(), PricingCatalogError> {
validate_rate(
entry_index,
format!("{field_prefix}.input_per_million"),
self.input_per_million,
)?;
validate_rate(
entry_index,
format!("{field_prefix}.output_per_million"),
self.output_per_million,
)?;
if let Some(value) = self.cache_read_per_million {
validate_rate(
entry_index,
format!("{field_prefix}.cache_read_per_million"),
value,
)?;
}
if let Some(value) = self.cache_write_per_million {
validate_rate(
entry_index,
format!("{field_prefix}.cache_write_per_million"),
value,
)?;
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TokenRateSchedule {
PromptTokenThreshold {
#[serde(default)]
applies_to: RateScheduleApplication,
tiers: Vec<TokenRateTier>,
},
}
impl TokenRateSchedule {
fn rates_for_usage(&self, usage: &Usage) -> Option<TokenPricingRates> {
match self {
Self::PromptTokenThreshold { applies_to, tiers } => {
if *applies_to != RateScheduleApplication::FullRequest {
return None;
}
let prompt_tokens = usage.prompt_tokens?;
tiers
.iter()
.find(|tier| tier.matches_prompt_tokens(prompt_tokens))
.map(|tier| tier.rates)
}
}
}
fn validate(&self, entry_index: usize) -> Result<(), PricingCatalogError> {
match self {
Self::PromptTokenThreshold { tiers, .. } if tiers.is_empty() => {
Err(PricingCatalogError::EmptyField {
entry_index,
field: "rate_schedule.tiers".to_string(),
})
}
Self::PromptTokenThreshold { tiers, .. } => {
for (tier_index, tier) in tiers.iter().enumerate() {
tier.validate(entry_index, tier_index)?;
}
Ok(())
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum RateScheduleApplication {
#[default]
FullRequest,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct TokenRateTier {
#[serde(skip_serializing_if = "Option::is_none")]
pub min_prompt_tokens: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_prompt_tokens: Option<u64>,
pub rates: TokenPricingRates,
}
impl TokenRateTier {
fn matches_prompt_tokens(&self, prompt_tokens: u64) -> bool {
self.min_prompt_tokens
.is_none_or(|min| prompt_tokens >= min)
&& self
.max_prompt_tokens
.is_none_or(|max| prompt_tokens <= max)
}
fn validate(&self, entry_index: usize, tier_index: usize) -> Result<(), PricingCatalogError> {
if let (Some(min), Some(max)) = (self.min_prompt_tokens, self.max_prompt_tokens)
&& min > max
{
return Err(PricingCatalogError::InvalidRate {
entry_index,
field: "rate_schedule.tiers.prompt_tokens".to_string(),
value: min as f64,
});
}
self.rates.validate(
entry_index,
&format!("rate_schedule.tiers[{tier_index}].rates"),
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptCachePricing {
pub read_accounting: CacheReadAccounting,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheReadAccounting {
IncludedInPromptTokens,
Separate,
}
#[must_use]
pub fn pricing_for_model(model: &str) -> Option<ModelPricing> {
active_pricing_resolver().pricing_for_model(model)
}
#[must_use]
pub fn pricing_for_provider(provider: Option<&str>, model: &str) -> Option<ModelPricing> {
active_pricing_resolver().pricing_for(provider, model)
}
#[must_use]
pub fn estimate_cost(model: &str, usage: &Usage) -> Option<CostEstimate> {
active_pricing_resolver().estimate_cost(model, usage)
}
#[must_use]
pub fn estimate_cost_for_provider(
provider: Option<&str>,
model: &str,
usage: &Usage,
) -> Option<CostEstimate> {
active_pricing_resolver().estimate_cost_for_provider(provider, model, usage)
}
#[must_use]
pub fn estimate_cost_with_catalog(
catalog: &PricingCatalog,
model: &str,
usage: &Usage,
) -> Option<CostEstimate> {
catalog
.pricing_for_model(model)
.and_then(|pricing| pricing.estimate_cost(usage))
}
#[must_use]
pub fn estimate_cost_with_provider(
catalog: &PricingCatalog,
provider: Option<&str>,
model: &str,
usage: &Usage,
) -> Option<CostEstimate> {
catalog
.pricing_for(provider, model)
.and_then(|pricing| pricing.estimate_cost(usage))
}
#[must_use]
pub fn active_pricing_resolver() -> Arc<PricingResolver> {
ACTIVE_PRICING_RESOLVER
.read()
.map(|resolver| Arc::clone(&resolver))
.unwrap_or_else(|_| Arc::new(PricingResolver::default()))
}
pub fn set_active_pricing_resolver(resolver: PricingResolver) -> Result<(), PricingCatalogError> {
let mut guard = ACTIVE_PRICING_RESOLVER
.write()
.map_err(|err| PricingCatalogError::LockPoisoned(err.to_string()))?;
*guard = Arc::new(resolver);
Ok(())
}
pub fn reset_active_pricing_resolver() -> Result<(), PricingCatalogError> {
set_active_pricing_resolver(PricingResolver::default())
}
pub fn attach_estimated_cost(response: &mut AnnotatedLlmResponse) {
attach_estimated_cost_for_provider(response, None);
}
pub fn attach_estimated_cost_for_provider(
response: &mut AnnotatedLlmResponse,
provider: Option<&str>,
) {
if response
.usage
.as_ref()
.and_then(|usage| usage.cost.as_ref())
.is_some()
{
return;
}
let Some(model) = response.model.clone() else {
return;
};
let Some(usage) = response.usage.as_mut() else {
return;
};
usage.cost = estimate_cost_for_provider(provider, &model, usage);
}
fn validate_nonempty(
entry_index: usize,
field: &'static str,
value: &str,
) -> Result<(), PricingCatalogError> {
if value.trim().is_empty() {
return Err(PricingCatalogError::EmptyField {
entry_index,
field: field.to_string(),
});
}
Ok(())
}
fn validate_rate(
entry_index: usize,
field: impl Into<String>,
value: f64,
) -> Result<(), PricingCatalogError> {
if !value.is_finite() || value < 0.0 {
return Err(PricingCatalogError::InvalidRate {
entry_index,
field: field.into(),
value,
});
}
Ok(())
}
fn default_pricing_currency() -> String {
"USD".into()
}
fn normalized_model_name(model: &str) -> String {
model.trim().to_ascii_lowercase()
}
fn normalized_provider_name(provider: &str) -> String {
provider.trim().trim_matches('/').to_ascii_lowercase()
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ModelLookupKey {
provider: Option<String>,
model: String,
}
fn normalized_model_lookup_keys(provider: Option<&str>, model: &str) -> Vec<ModelLookupKey> {
let normalized = normalized_model_name(model);
if normalized.is_empty() {
return vec![];
}
let parts: Vec<&str> = normalized
.split('/')
.map(str::trim)
.filter(|part| !part.is_empty())
.collect();
let mut keys = Vec::with_capacity(parts.len() + 3);
let explicit_provider = provider
.map(normalized_provider_name)
.filter(|provider| !provider.is_empty());
let terminal_model = parts
.last()
.copied()
.unwrap_or(normalized.as_str())
.to_string();
if let Some(provider) = explicit_provider {
push_lookup_key(&mut keys, Some(provider.clone()), normalized.clone());
push_lookup_key(&mut keys, Some(provider), terminal_model.clone());
} else if parts.len() > 1 {
push_lookup_key(
&mut keys,
Some(parts[..parts.len() - 1].join("/")),
terminal_model,
);
}
for start in 0..parts.len() {
let key = parts[start..].join("/");
push_lookup_key(&mut keys, None, key);
}
keys
}
fn push_lookup_key(keys: &mut Vec<ModelLookupKey>, provider: Option<String>, model: String) {
let key = ModelLookupKey { provider, model };
if !key.model.is_empty() && !keys.contains(&key) {
keys.push(key);
}
}
#[must_use]
pub fn infer_model_provider(default_provider: &str, model: Option<&str>) -> Option<String> {
let normalized_default = normalized_provider_name(default_provider);
if let Some(model) = model {
let normalized = normalized_model_name(model);
let parts: Vec<&str> = normalized
.split('/')
.map(str::trim)
.filter(|part| !part.is_empty())
.collect();
if parts.len() > 1 {
return Some(parts[..parts.len() - 1].join("/"));
}
}
(!normalized_default.is_empty()).then_some(normalized_default)
}
fn cost_component(tokens: u64, price_per_million: f64) -> f64 {
tokens as f64 * price_per_million / 1_000_000.0
}
fn cost_component_if_nonzero(tokens: u64, price_per_million: f64) -> Option<f64> {
(tokens > 0).then(|| round_cost_amount(cost_component(tokens, price_per_million)))
}
fn round_cost_amount(cost: f64) -> f64 {
const SCALE: f64 = 1_000_000_000_000.0;
(cost * SCALE).round() / SCALE
}