use std::collections::BTreeMap;
use serde_json::Value;
use crate::budgets::BudgetSpend;
pub const MICRO_PER_USD: u64 = 1_000_000;
pub const MICRO_PER_CENT: u64 = 10_000;
const TOKENS_PER_QUOTE: u64 = 1_000_000;
const MICRO_PER_USD_F: f64 = 1_000_000.0;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Price {
pub input: u64,
pub output: u64,
pub cache_read: u64,
pub cache_write: u64,
}
impl Price {
pub fn cost(&self, spend: BudgetSpend) -> u64 {
rate(spend.uncached_prompt_tokens(), self.input)
.saturating_add(rate(spend.output_tokens, self.output))
.saturating_add(rate(spend.cache_read_tokens, self.cache_read))
.saturating_add(rate(spend.cache_write_tokens, self.cache_write))
}
}
fn rate(tokens: u64, per_quote: u64) -> u64 {
tokens.saturating_mul(per_quote) / TOKENS_PER_QUOTE
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Prices(BTreeMap<String, Price>);
impl Prices {
pub fn from_json(value: &Value) -> Self {
let Some(rows) = value.as_object() else {
return Self::default();
};
Self(
rows.iter()
.filter_map(|(model, row)| Some((model.clone(), price(row)?)))
.collect(),
)
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn of(&self, model: Option<&str>) -> Option<Price> {
self.0.get(model?).copied()
}
}
fn price(value: &Value) -> Option<Price> {
let row = value.as_object()?;
Some(Price {
input: micros(row.get("input")),
output: micros(row.get("output")),
cache_read: micros(row.get("cache_read")),
cache_write: micros(row.get("cache_write")),
})
}
pub(super) fn quoted(value: Option<&Value>) -> Option<u64> {
let usd = value.and_then(Value::as_f64).filter(|usd| *usd >= 0.0)?;
Some((usd * MICRO_PER_USD_F).round() as u64)
}
fn micros(value: Option<&Value>) -> u64 {
quoted(value).unwrap_or(0)
}