a_agent/
pricing.rs

1use std::collections::BTreeMap;
2use std::time::Duration;
3
4use anyhow::Result;
5use serde::{Deserialize, Serialize};
6
7use crate::config::Rates;
8use crate::model::Usage;
9
10pub const SOURCE_URL: &str = "https://models.dev/api.json";
11pub const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);
12
13/// The catalog endpoint, overridable with `A_PRICING_URL` for mirrors, proxies,
14/// and tests that must not reach a third-party service.
15pub fn source_url() -> String {
16    std::env::var("A_PRICING_URL").unwrap_or_else(|_| SOURCE_URL.to_owned())
17}
18
19/// What is known about a model's prices. Ambiguity is reported rather than
20/// resolved by guessing: mirrors of the same model id charge different rates, so
21/// picking one would produce a plausible but wrong number.
22#[derive(Debug, Clone, PartialEq)]
23pub enum Resolution {
24    Known {
25        schedule: Schedule,
26        source: String,
27    },
28    Ambiguous {
29        model: String,
30        providers: Vec<String>,
31    },
32    Unknown(String),
33}
34
35#[derive(Deserialize)]
36struct Provider {
37    #[serde(default)]
38    models: BTreeMap<String, Model>,
39}
40
41#[derive(Deserialize)]
42struct Model {
43    #[serde(default)]
44    cost: Option<Cost>,
45}
46
47#[derive(Deserialize)]
48struct Cost {
49    #[serde(default)]
50    input: f64,
51    #[serde(default)]
52    output: f64,
53    #[serde(default)]
54    cache_read: f64,
55    #[serde(default)]
56    cache_write: f64,
57    /// Rates that replace the base ones once a request is large enough.
58    #[serde(default)]
59    tiers: Vec<Tier>,
60}
61
62#[derive(Deserialize)]
63struct Tier {
64    #[serde(default)]
65    input: f64,
66    #[serde(default)]
67    output: f64,
68    #[serde(default)]
69    cache_read: f64,
70    #[serde(default)]
71    cache_write: f64,
72    tier: TierBound,
73}
74
75#[derive(Deserialize)]
76struct TierBound {
77    #[serde(rename = "type")]
78    kind: String,
79    #[serde(default)]
80    size: u64,
81}
82
83impl From<&Cost> for Rates {
84    fn from(cost: &Cost) -> Self {
85        Self {
86            input: cost.input,
87            output: cost.output,
88            cache_read: cost.cache_read,
89            cache_write: cost.cache_write,
90        }
91    }
92}
93
94impl From<&Tier> for Rates {
95    fn from(tier: &Tier) -> Self {
96        Self {
97            input: tier.input,
98            output: tier.output,
99            cache_read: tier.cache_read,
100            cache_write: tier.cache_write,
101        }
102    }
103}
104
105/// Base rates plus any context-size tiers, so a request can be priced by how
106/// large it actually was.
107#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
108#[serde(default)]
109pub struct Schedule {
110    pub base: Rates,
111    /// Ascending by threshold; the highest threshold at or below the request's
112    /// context size wins.
113    pub tiers: Vec<(u64, Rates)>,
114}
115
116impl Schedule {
117    pub fn flat(base: Rates) -> Self {
118        Self {
119            base,
120            tiers: Vec::new(),
121        }
122    }
123
124    fn from_cost(cost: &Cost) -> Self {
125        let mut tiers = cost
126            .tiers
127            .iter()
128            .filter(|tier| tier.tier.kind == "context" && tier.tier.size > 0)
129            .map(|tier| (tier.tier.size, Rates::from(tier)))
130            .collect::<Vec<_>>();
131        tiers.sort_by_key(|(size, _)| *size);
132        Self {
133            base: Rates::from(cost),
134            tiers,
135        }
136    }
137
138    /// The rates that apply to a request whose context was `context_tokens`.
139    pub fn rates_for(&self, context_tokens: u64) -> Rates {
140        self.tiers
141            .iter()
142            .rev()
143            .find(|(size, _)| context_tokens >= *size)
144            .map_or(self.base, |(_, rates)| *rates)
145    }
146
147    pub fn is_tiered(&self) -> bool {
148        !self.tiers.is_empty()
149    }
150}
151
152/// Looks up `model` in a models.dev catalog. `key` is an explicit
153/// `provider/model` selector; without one the model id must be unique.
154pub fn resolve_from_catalog(catalog: &str, model: &str, key: Option<&str>) -> Resolution {
155    let providers = match serde_json::from_str::<BTreeMap<String, Provider>>(catalog) {
156        Ok(providers) => providers,
157        Err(error) => return Resolution::Unknown(format!("catalog is not valid JSON: {error}")),
158    };
159    if let Some(key) = key {
160        let Some((provider, model)) = key.split_once('/') else {
161            return Resolution::Unknown(format!(
162                "pricing key {key:?} must look like provider/model"
163            ));
164        };
165        return match providers
166            .get(provider)
167            .and_then(|entry| entry.models.get(model))
168        {
169            Some(entry) => match &entry.cost {
170                Some(cost) => Resolution::Known {
171                    schedule: Schedule::from_cost(cost),
172                    source: key.to_owned(),
173                },
174                None => Resolution::Unknown(format!("{key} lists no prices")),
175            },
176            None => Resolution::Unknown(format!("{key} was not found on models.dev")),
177        };
178    }
179    let matches = providers
180        .iter()
181        .filter_map(|(name, provider)| {
182            provider
183                .models
184                .get(model)
185                .map(|entry| (name.clone(), entry.cost.as_ref().map(Schedule::from_cost)))
186        })
187        .collect::<Vec<_>>();
188    match matches.len() {
189        0 => Resolution::Unknown(format!("{model} was not found on models.dev")),
190        1 => {
191            let (provider, schedule) = matches.into_iter().next().expect("one match");
192            match schedule {
193                Some(schedule) => Resolution::Known {
194                    schedule,
195                    source: format!("{provider}/{model}"),
196                },
197                None => Resolution::Unknown(format!("{provider}/{model} lists no prices")),
198            }
199        }
200        _ => Resolution::Ambiguous {
201            model: model.to_owned(),
202            providers: matches.into_iter().map(|(name, _)| name).collect(),
203        },
204    }
205}
206
207pub async fn fetch_catalog(url: &str, timeout: Duration) -> Result<String> {
208    let client = reqwest::Client::builder().timeout(timeout).build()?;
209    Ok(client
210        .get(url)
211        .send()
212        .await?
213        .error_for_status()?
214        .text()
215        .await?)
216}
217
218/// Cost in USD of one request.
219pub fn request_cost(usage: Usage, rates: Rates) -> f64 {
220    let scale = |tokens: Option<u64>, rate: f64| tokens.unwrap_or(0) as f64 * rate;
221    (scale(usage.input_tokens, rates.input)
222        + scale(usage.output_tokens, rates.output)
223        + scale(usage.cached_tokens, rates.cache_read)
224        + scale(usage.cache_write_tokens, rates.cache_write))
225        / 1_000_000.0
226}
227
228/// Cost of a whole session. Each request is priced on its own, because a tiered
229/// schedule charges by how large that individual request was.
230pub fn session_cost(requests: &[Usage], schedule: &Schedule) -> f64 {
231    requests
232        .iter()
233        .map(|usage| request_cost(*usage, schedule.rates_for(request_context(*usage))))
234        .sum()
235}
236
237/// Tokens the provider had to read for a request, which is what a context tier
238/// is measured against.
239fn request_context(usage: Usage) -> u64 {
240    usage.input_tokens.unwrap_or(0)
241        + usage.cached_tokens.unwrap_or(0)
242        + usage.cache_write_tokens.unwrap_or(0)
243}
244
245/// The candidate most likely to be the model's own vendor, used only to make the
246/// suggested `pricing` line useful. A provider whose id prefixes the model id,
247/// such as `deepseek` for `deepseek-v4-flash`, beats an alphabetical first pick.
248pub fn likely_provider<'a>(model: &str, providers: &'a [String]) -> Option<&'a String> {
249    providers
250        .iter()
251        .find(|provider| model.starts_with(provider.as_str()))
252        .or_else(|| providers.first())
253}
254
255pub fn format_cost(cost: f64) -> String {
256    // `f64`'s Sum identity is -0.0, so an empty or all-zero session would
257    // otherwise print "$-0.00".
258    let cost = if cost == 0.0 { 0.0 } else { cost };
259    if cost > 0.0 && cost < 0.01 {
260        format!("${cost:.4}")
261    } else {
262        format!("${cost:.2}")
263    }
264}
265
266/// Whether the provider reported any tokens for this session. Without that,
267/// there is nothing to price, and a dollar figure would claim a measurement that
268/// was never made.
269pub fn has_measured_usage(requests: &[Usage]) -> bool {
270    requests.iter().any(|usage| {
271        [
272            usage.input_tokens,
273            usage.output_tokens,
274            usage.cached_tokens,
275            usage.cache_write_tokens,
276        ]
277        .iter()
278        .any(|tokens| tokens.unwrap_or(0) > 0)
279    })
280}
a_agent/pricing.rs

a_agent/
pricing.rs
