1use std::collections::BTreeMap;
28
29use serde::{Deserialize, Serialize};
30
31use crate::provider::Usage;
32
33pub type Micros = u128;
36
37pub const MICROS: Micros = 1_000_000;
39
40pub const TOKENS_PER_QUOTE: Micros = 1_000_000;
43
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
46#[serde(deny_unknown_fields, rename_all = "kebab-case")]
47pub struct ModelPrice {
48 pub model: String,
54 pub input: String,
56 pub output: String,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
61 pub cache_read: Option<String>,
62 pub currency: String,
64}
65
66#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
68#[serde(transparent)]
69pub struct Pricing {
70 models: Vec<ModelPrice>,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq)]
75pub enum Charge {
76 Priced { micros: Micros, currency: String },
78 Unpriced,
80 WrongCurrency { priced_in: String },
84}
85
86impl Pricing {
87 pub fn new(models: Vec<ModelPrice>) -> Pricing {
88 Pricing { models }
89 }
90
91 pub fn is_empty(&self) -> bool {
92 self.models.is_empty()
93 }
94
95 pub fn models(&self) -> impl Iterator<Item = &str> {
97 self.models.iter().map(|price| price.model.as_str())
98 }
99
100 pub fn charge(&self, model: &str, usage: Usage, budget_currency: &str) -> Charge {
102 let Some(price) = self.models.iter().find(|price| price.model == model) else {
103 return Charge::Unpriced;
104 };
105 if !price.currency.eq_ignore_ascii_case(budget_currency) {
106 return Charge::WrongCurrency {
107 priced_in: price.currency.clone(),
108 };
109 }
110
111 let uncached = usage.input_tokens.saturating_sub(usage.cache_read_tokens);
115 let cache_rate = price.cache_read.as_deref().unwrap_or(&price.input);
116
117 let mut micros: Micros = 0;
118 for (tokens, rate) in [
119 (uncached, price.input.as_str()),
120 (usage.cache_read_tokens, cache_rate),
121 (usage.output_tokens, price.output.as_str()),
122 ] {
123 let Some(per_quote) = parse_micros(rate) else {
124 return Charge::Unpriced;
125 };
126 micros += Micros::from(tokens) * per_quote / TOKENS_PER_QUOTE;
127 }
128
129 Charge::Priced {
130 micros,
131 currency: price.currency.clone(),
132 }
133 }
134}
135
136pub fn parse_micros(amount: &str) -> Option<Micros> {
142 let amount = amount.trim();
143 if amount.is_empty() || amount.starts_with('-') {
144 return None;
145 }
146 let (whole, fraction) = match amount.split_once('.') {
147 Some((whole, fraction)) => (whole, fraction),
148 None => (amount, ""),
149 };
150 if whole.is_empty() && fraction.is_empty() {
151 return None;
152 }
153 if !whole.chars().all(|ch| ch.is_ascii_digit())
154 || !fraction.chars().all(|ch| ch.is_ascii_digit())
155 || fraction.len() > 6
156 {
157 return None;
158 }
159
160 let whole: Micros = if whole.is_empty() {
161 0
162 } else {
163 whole.parse().ok()?
164 };
165 let mut padded = fraction.to_string();
166 while padded.len() < 6 {
167 padded.push('0');
168 }
169 let fraction: Micros = if padded.is_empty() {
170 0
171 } else {
172 padded.parse().ok()?
173 };
174 whole.checked_mul(MICROS)?.checked_add(fraction)
175}
176
177pub fn render_micros(micros: Micros) -> String {
182 let whole = micros / MICROS;
183 let fraction = micros % MICROS;
184 if fraction == 0 {
185 return whole.to_string();
186 }
187 let rendered = format!("{whole}.{fraction:06}");
188 rendered.trim_end_matches('0').to_string()
189}
190
191#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
196#[serde(rename_all = "camelCase", default)]
197pub struct Spend {
198 micros: Micros,
199 currency: Option<String>,
200 unpriced: BTreeMap<String, String>,
203}
204
205impl Spend {
206 pub fn add(&mut self, model: &str, charge: Charge) {
208 match charge {
209 Charge::Priced { micros, currency } => {
210 self.micros += micros;
211 self.currency = Some(currency);
212 }
213 Charge::Unpriced => {
214 self.unpriced
215 .insert(model.to_string(), "no price is configured".to_string());
216 }
217 Charge::WrongCurrency { priced_in } => {
218 self.unpriced.insert(
219 model.to_string(),
220 format!("priced in {priced_in}, which the budget is not"),
221 );
222 }
223 }
224 }
225
226 pub fn micros(&self) -> Micros {
227 self.micros
228 }
229
230 pub fn is_complete(&self) -> bool {
236 self.unpriced.is_empty()
237 }
238
239 pub fn unpriced(&self) -> impl Iterator<Item = (&str, &str)> {
241 self.unpriced
242 .iter()
243 .map(|(model, reason)| (model.as_str(), reason.as_str()))
244 }
245
246 pub fn rendered(&self) -> Option<String> {
248 let currency = self.currency.as_ref()?;
249 Some(format!(
250 "{} {}",
251 render_micros(self.micros),
252 currency.to_ascii_uppercase()
253 ))
254 }
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260
261 fn usage(input: u64, output: u64, cached: u64) -> Usage {
262 Usage {
263 input_tokens: input,
264 output_tokens: output,
265 cache_read_tokens: cached,
266 }
267 }
268
269 fn opus() -> Pricing {
270 Pricing::new(vec![ModelPrice {
271 model: "claude-opus-5".into(),
272 input: "3".into(),
273 output: "15".into(),
274 cache_read: Some("0.3".into()),
275 currency: "usd".into(),
276 }])
277 }
278
279 #[test]
280 fn a_decimal_string_becomes_exact_millionths() {
281 assert_eq!(parse_micros("3"), Some(3_000_000));
282 assert_eq!(parse_micros("0.3"), Some(300_000));
283 assert_eq!(parse_micros("3.5"), Some(3_500_000));
284 assert_eq!(parse_micros("0.000001"), Some(1));
285 assert_eq!(parse_micros(".5"), Some(500_000));
286 }
287
288 #[test]
289 fn an_amount_this_format_cannot_hold_is_refused_rather_than_rounded() {
290 assert_eq!(parse_micros("0.0000001"), None);
293 assert_eq!(parse_micros("-1"), None);
294 assert_eq!(parse_micros("free"), None);
295 assert_eq!(parse_micros(""), None);
296 assert_eq!(parse_micros("1.2.3"), None);
297 }
298
299 #[test]
300 fn rendering_round_trips_through_the_ir_encoding() {
301 for amount in ["0", "5", "0.25", "0.0042", "1234.567891"] {
302 let micros = parse_micros(amount).expect("a valid amount");
303 assert_eq!(render_micros(micros), amount);
304 }
305 }
306
307 #[test]
308 fn a_call_is_charged_at_the_quoted_rate() {
309 let charge = opus().charge("claude-opus-5", usage(1000, 500, 0), "usd");
311 assert_eq!(
312 charge,
313 Charge::Priced {
314 micros: 3_000 + 7_500,
315 currency: "usd".into()
316 }
317 );
318 assert_eq!(render_micros(10_500), "0.0105");
319 }
320
321 #[test]
322 fn cached_input_is_charged_at_its_own_rate_when_one_is_given() {
323 let charge = opus().charge("claude-opus-5", usage(1000, 0, 800), "usd");
325 assert_eq!(
326 charge,
327 Charge::Priced {
328 micros: 600 + 240,
329 currency: "usd".into()
330 }
331 );
332 }
333
334 #[test]
335 fn cached_input_falls_back_to_the_input_rate() {
336 let pricing = Pricing::new(vec![ModelPrice {
339 model: "m".into(),
340 input: "3".into(),
341 output: "15".into(),
342 cache_read: None,
343 currency: "usd".into(),
344 }]);
345 let charge = pricing.charge("m", usage(1000, 0, 1000), "usd");
346 assert_eq!(
347 charge,
348 Charge::Priced {
349 micros: 3_000,
350 currency: "usd".into()
351 }
352 );
353 }
354
355 #[test]
356 fn an_unknown_model_is_unpriced_rather_than_free() {
357 assert_eq!(
358 opus().charge("claude-opus-5-mini", usage(1000, 500, 0), "usd"),
359 Charge::Unpriced,
360 "a prefix rule would price a different model at this one's rate"
361 );
362 assert_eq!(
363 Pricing::default().charge("anything", usage(1, 1, 0), "usd"),
364 Charge::Unpriced
365 );
366 }
367
368 #[test]
369 fn a_price_in_another_currency_does_not_get_converted() {
370 let charge = opus().charge("claude-opus-5", usage(1000, 0, 0), "eur");
371 assert_eq!(
372 charge,
373 Charge::WrongCurrency {
374 priced_in: "usd".into()
375 },
376 "converting needs a rate, which is a second time-dependent input"
377 );
378 }
379
380 #[test]
381 fn a_spend_that_missed_a_call_is_not_a_total() {
382 let mut spend = Spend::default();
383 spend.add(
384 "claude-opus-5",
385 opus().charge("claude-opus-5", usage(1000, 500, 0), "usd"),
386 );
387 assert!(spend.is_complete());
388 assert_eq!(spend.rendered().as_deref(), Some("0.0105 USD"));
389
390 spend.add("mystery", Charge::Unpriced);
391 assert!(
392 !spend.is_complete(),
393 "a total that missed a call is not a total"
394 );
395 let unpriced: Vec<&str> = spend.unpriced().map(|(model, _)| model).collect();
396 assert_eq!(unpriced, vec!["mystery"]);
397
398 spend.add("mystery", Charge::Unpriced);
400 assert_eq!(spend.unpriced().count(), 1);
401 }
402}