1use jiff::Timestamp;
8use std::collections::BTreeMap;
9
10use serde::{Deserialize, Serialize};
11
12use super::{CostSource, RunCost, TokenUsage};
13use crate::agent::AgentName;
14use crate::flight::{ItineraryId, RunId};
15
16#[derive(Debug, Clone, Copy, Default, PartialEq, Deserialize, Serialize)]
22#[serde(deny_unknown_fields)]
23pub struct ModelRates {
24 #[serde(default)]
26 pub input_usd: f64,
27 #[serde(default)]
29 pub output_usd: f64,
30 #[serde(default)]
32 pub cache_read_usd: f64,
33 #[serde(default)]
35 pub cache_write_usd: f64,
36}
37
38impl ModelRates {
39 #[must_use]
44 pub fn cost_of(&self, usage: TokenUsage) -> Option<f64> {
45 let rates = [
46 (self.input_usd, usage.input),
47 (self.output_usd, usage.output),
48 (self.cache_read_usd, usage.cache_read),
49 (self.cache_write_usd, usage.cache_write),
50 ];
51
52 let mut total = 0.0;
53 for (rate, tokens) in rates {
54 if !rate.is_finite() || rate < 0.0 {
55 return None;
56 }
57 #[allow(clippy::cast_precision_loss)]
58 let tokens = tokens as f64;
59 total += rate * tokens / 1_000_000.0;
60 }
61
62 total.is_finite().then_some(total)
63 }
64}
65
66#[derive(Debug, Clone, Default, Deserialize, Serialize)]
72#[serde(transparent)]
73pub struct RateCard {
74 models: BTreeMap<String, ModelRates>,
75}
76
77impl RateCard {
78 #[must_use]
80 pub fn new() -> Self {
81 Self::default()
82 }
83
84 #[must_use]
86 pub fn with(mut self, model: impl Into<String>, rates: ModelRates) -> Self {
87 self.models.insert(model.into(), rates);
88 self
89 }
90
91 #[must_use]
93 pub fn rates_for(&self, model: &str) -> Option<&ModelRates> {
94 self.models.get(model)
95 }
96
97 #[must_use]
99 pub fn is_empty(&self) -> bool {
100 self.models.is_empty()
101 }
102
103 #[must_use]
108 pub fn estimate(&self, model: &str, usage: TokenUsage) -> Option<f64> {
109 self.rates_for(model)?.cost_of(usage)
110 }
111
112 #[must_use]
117 pub fn price(
118 &self,
119 run: RunId,
120 itinerary: ItineraryId,
121 agent: AgentName,
122 model: Option<String>,
123 usage: TokenUsage,
124 ) -> RunCost {
125 let estimate = model
126 .as_deref()
127 .filter(|_| !usage.is_empty())
128 .and_then(|model| self.estimate(model, usage));
129
130 let (usd, source) = match estimate {
131 Some(usd) => (usd, CostSource::RateCard),
132 None => (0.0, CostSource::Unreported),
133 };
134
135 RunCost {
136 run,
137 itinerary,
138 agent,
139 pipeline: None,
140 model,
141 usage,
142 usd,
143 source,
144 at: Timestamp::now(),
145 }
146 }
147}
148
149#[cfg(test)]
150mod tests {
151 use super::*;
152
153 fn opus() -> ModelRates {
155 ModelRates {
156 input_usd: 5.0,
157 output_usd: 25.0,
158 cache_read_usd: 0.5,
159 cache_write_usd: 6.25,
160 }
161 }
162
163 fn card() -> RateCard {
164 RateCard::new().with("claude-opus-5", opus())
165 }
166
167 fn usage() -> TokenUsage {
168 TokenUsage {
169 input: 1_000_000,
170 output: 1_000_000,
171 cache_read: 1_000_000,
172 cache_write: 1_000_000,
173 }
174 }
175
176 fn priced(card: &RateCard, model: Option<&str>, usage: TokenUsage) -> RunCost {
177 card.price(
178 RunId::generate(),
179 ItineraryId::generate(),
180 "analyst".into(),
181 model.map(ToOwned::to_owned),
182 usage,
183 )
184 }
185
186 #[test]
187 fn a_million_of_each_costs_the_sum_of_the_rates() {
188 let cost = card().estimate("claude-opus-5", usage()).expect("priced");
189
190 assert!((cost - 36.75).abs() < 1e-9, "got {cost}");
191 }
192
193 #[test]
194 fn cached_tokens_are_priced_apart_from_fresh_input() {
195 let fresh = card()
197 .estimate(
198 "claude-opus-5",
199 TokenUsage {
200 input: 1_000_000,
201 ..TokenUsage::default()
202 },
203 )
204 .expect("priced");
205 let cached = card()
206 .estimate(
207 "claude-opus-5",
208 TokenUsage {
209 cache_read: 1_000_000,
210 ..TokenUsage::default()
211 },
212 )
213 .expect("priced");
214
215 assert!((fresh - 5.0).abs() < 1e-9);
216 assert!((cached - 0.5).abs() < 1e-9);
217 assert!(cached < fresh);
218 }
219
220 #[test]
221 fn an_estimate_is_labelled_as_an_estimate() {
222 let cost = priced(&card(), Some("claude-opus-5"), usage());
223
224 assert_eq!(cost.source, CostSource::RateCard);
225 assert!(
226 !cost.source.is_measured(),
227 "an estimate must never count as measured"
228 );
229 }
230
231 #[test]
232 fn an_unknown_model_is_unreported_rather_than_free() {
233 let cost = priced(&card(), Some("some-new-model"), usage());
235
236 assert_eq!(cost.source, CostSource::Unreported);
237 assert!((cost.usd - 0.0).abs() < f64::EPSILON);
238 }
239
240 #[test]
241 fn a_run_with_no_model_cannot_be_priced() {
242 assert_eq!(
243 priced(&card(), None, usage()).source,
244 CostSource::Unreported
245 );
246 }
247
248 #[test]
249 fn a_run_with_no_tokens_cannot_be_priced() {
250 assert_eq!(
251 priced(&card(), Some("claude-opus-5"), TokenUsage::default()).source,
252 CostSource::Unreported
253 );
254 }
255
256 #[test]
257 fn a_nonsense_rate_is_refused_rather_than_propagated() {
258 for bad in [f64::NAN, f64::INFINITY, -1.0] {
259 let card = RateCard::new().with(
260 "broken",
261 ModelRates {
262 output_usd: bad,
263 ..ModelRates::default()
264 },
265 );
266
267 assert_eq!(
268 card.estimate("broken", usage()),
269 None,
270 "{bad} must not produce a price"
271 );
272 }
273 }
274
275 #[test]
276 fn rates_parse_from_configuration() {
277 let card: RateCard = toml::from_str(
278 r"
279 [claude-opus-5]
280 input_usd = 5.0
281 output_usd = 25.0
282 cache_read_usd = 0.5
283 cache_write_usd = 6.25
284 ",
285 )
286 .expect("a rate card parses");
287
288 assert!(!card.is_empty());
289 assert_eq!(card.rates_for("claude-opus-5"), Some(&opus()));
290 }
291
292 #[test]
293 fn an_unknown_rate_field_is_rejected() {
294 let error = toml::from_str::<RateCard>(
296 r"
297 [claude-opus-5]
298 output_used = 25.0
299 ",
300 );
301
302 assert!(error.is_err());
303 }
304}