Skip to main content

gateway_core/
catalog.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
6pub struct Usage {
7    pub input_tokens: u64,
8    pub output_tokens: u64,
9    #[serde(default)]
10    pub reasoning_tokens: u64,
11    pub cache_read_tokens: u64,
12    pub cache_write_tokens: u64,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub struct ModelPrice {
17    pub input_microdollars_per_million: u64,
18    pub output_microdollars_per_million: u64,
19    #[serde(default)]
20    pub reasoning_microdollars_per_million: Option<u64>,
21    #[serde(default)]
22    pub cache_read_microdollars_per_million: Option<u64>,
23    #[serde(default)]
24    pub cache_write_microdollars_per_million: Option<u64>,
25}
26
27impl ModelPrice {
28    /// Cost of a usage record in integer micro-dollars. Reasoning tokens are a
29    /// subset of output tokens and are priced separately (falling back to the
30    /// output rate); cache reads/writes fall back to the input rate. Saturates
31    /// rather than overflowing.
32    pub fn cost_microdollars(&self, usage: Usage) -> u64 {
33        let billed_output_tokens = usage.output_tokens.saturating_sub(usage.reasoning_tokens);
34        let cost = component(usage.input_tokens, self.input_microdollars_per_million)
35            .saturating_add(component(
36                billed_output_tokens,
37                self.output_microdollars_per_million,
38            ))
39            .saturating_add(component(
40                usage.reasoning_tokens,
41                self.reasoning_microdollars_per_million
42                    .unwrap_or(self.output_microdollars_per_million),
43            ))
44            .saturating_add(component(
45                usage.cache_read_tokens,
46                self.cache_read_microdollars_per_million
47                    .unwrap_or(self.input_microdollars_per_million),
48            ))
49            .saturating_add(component(
50                usage.cache_write_tokens,
51                self.cache_write_microdollars_per_million
52                    .unwrap_or(self.input_microdollars_per_million),
53            ));
54        u64::try_from(cost).unwrap_or(u64::MAX)
55    }
56}
57
58#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
59pub struct CatalogModel {
60    pub id: String,
61    pub price: ModelPrice,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub struct ModelCatalog {
66    pub version: u64,
67    pub models: Vec<CatalogModel>,
68}
69
70#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
71pub struct UsageReceipt {
72    pub catalog_version: u64,
73    pub model: String,
74    pub price: ModelPrice,
75    pub usage: Usage,
76    pub cost_microdollars: u64,
77}
78
79impl UsageReceipt {
80    pub fn normalized_cost_microdollars(&self) -> u64 {
81        self.cost_microdollars
82    }
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
86pub enum CatalogError {
87    #[error("duplicate model '{0}' in catalog")]
88    DuplicateModel(String),
89    #[error("model '{0}' is not priced in catalog")]
90    UnknownModel(String),
91    #[error("usage cost overflow")]
92    CostOverflow,
93}
94
95impl ModelCatalog {
96    pub fn validate(&self) -> Result<(), CatalogError> {
97        let mut seen = HashMap::new();
98        for model in &self.models {
99            if seen.insert(&model.id, ()).is_some() {
100                return Err(CatalogError::DuplicateModel(model.id.clone()));
101            }
102        }
103        Ok(())
104    }
105
106    pub fn receipt(&self, model: &str, usage: Usage) -> Result<UsageReceipt, CatalogError> {
107        let price = self
108            .models
109            .iter()
110            .find(|entry| entry.id == model)
111            .map(|entry| entry.price)
112            .ok_or_else(|| CatalogError::UnknownModel(model.to_owned()))?;
113        let cost_microdollars = price.cost_microdollars(usage);
114        Ok(UsageReceipt {
115            catalog_version: self.version,
116            model: model.to_owned(),
117            price,
118            usage,
119            cost_microdollars,
120        })
121    }
122}
123
124fn component(tokens: u64, microdollars_per_million: u64) -> u128 {
125    u128::from(tokens).saturating_mul(u128::from(microdollars_per_million)) / 1_000_000
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    fn catalog(version: u64, output_price: u64) -> ModelCatalog {
133        ModelCatalog {
134            version,
135            models: vec![CatalogModel {
136                id: "openai/model".into(),
137                price: ModelPrice {
138                    input_microdollars_per_million: 1_000_000,
139                    output_microdollars_per_million: output_price,
140                    reasoning_microdollars_per_million: None,
141                    cache_read_microdollars_per_million: Some(100_000),
142                    cache_write_microdollars_per_million: Some(1_250_000),
143                },
144            }],
145        }
146    }
147
148    #[test]
149    fn receipt_pins_catalog_version_and_normalizes_cost() {
150        let usage = Usage {
151            input_tokens: 1_000,
152            output_tokens: 500,
153            reasoning_tokens: 0,
154            cache_read_tokens: 1_000,
155            cache_write_tokens: 200,
156        };
157        let receipt = catalog(7, 2_000_000)
158            .receipt("openai/model", usage)
159            .unwrap();
160        assert_eq!(receipt.catalog_version, 7);
161        assert_eq!(receipt.price, catalog(7, 2_000_000).models[0].price);
162        assert_eq!(receipt.cost_microdollars, 2_350);
163        assert_eq!(receipt.normalized_cost_microdollars(), 2_350);
164        assert_eq!(
165            catalog(8, 4_000_000)
166                .receipt("openai/model", usage)
167                .unwrap()
168                .cost_microdollars,
169            3_350
170        );
171    }
172
173    #[test]
174    fn optional_reasoning_price_replaces_output_price_for_reasoning_subset() {
175        let mut catalog = catalog(9, 2_000_000);
176        catalog.models[0].price.reasoning_microdollars_per_million = Some(4_000_000);
177        let receipt = catalog
178            .receipt(
179                "openai/model",
180                Usage {
181                    input_tokens: 0,
182                    output_tokens: 1_000,
183                    reasoning_tokens: 250,
184                    cache_read_tokens: 0,
185                    cache_write_tokens: 0,
186                },
187            )
188            .unwrap();
189        assert_eq!(receipt.cost_microdollars, 2_500);
190        assert_eq!(receipt.catalog_version, 9);
191        assert_eq!(
192            receipt.price.reasoning_microdollars_per_million,
193            Some(4_000_000)
194        );
195    }
196}