forge_agent/env/token_usage/
pricing.rs

1//! LiteLLM-sourced per-model pricing: the parser plus the proxy-aware
2//! fetch that populates the runtime cache.
3//!
4//! Pricing is sourced at runtime from LiteLLM's raw file and cached in
5//! redb (`forge_workspace::store::pricing`); there is no bundled copy.
6
7use std::collections::HashMap;
8use std::time::Duration;
9
10use forge_primitives::token_usage::ModelPricing;
11
12/// LiteLLM's community-maintained price table (prices Claude and the
13/// OpenAI/Codex ids in one file).
14const LITELLM_URL: &str =
15    "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json";
16
17/// Bound on the pricing fetch; the file is ~1.6 MB.
18const FETCH_TIMEOUT: Duration = Duration::from_secs(20);
19
20/// Per-model USD pricing keyed by the exact `message.model` id.
21pub struct PricingTable {
22    map: HashMap<String, ModelPricing>,
23}
24
25impl PricingTable {
26    /// Parse a LiteLLM `model_prices_and_context_window.json` body.
27    /// Entries missing either core per-token cost are skipped; a
28    /// malformed body yields an empty table (logged), never a panic.
29    pub fn from_litellm_json(json: &str) -> Self {
30        let value: serde_json::Value = match serde_json::from_str(json) {
31            Ok(value) => value,
32            Err(error) => {
33                tracing::warn!(
34                    target: "forge_agent::env::token_usage::pricing",
35                    %error,
36                    "litellm pricing json did not parse; using an empty table",
37                );
38                return Self { map: HashMap::new() };
39            }
40        };
41        let Some(entries) = value.as_object() else {
42            return Self { map: HashMap::new() };
43        };
44        let mut map = HashMap::with_capacity(entries.len());
45        for (model_id, entry) in entries {
46            // LiteLLM ships a `sample_spec` documentation entry with
47            // placeholder cost fields; it is not a real model id.
48            if model_id == "sample_spec" {
49                continue;
50            }
51            let field = |name: &str| entry.get(name).and_then(serde_json::Value::as_f64);
52            let (Some(input), Some(output)) =
53                (field("input_cost_per_token"), field("output_cost_per_token"))
54            else {
55                continue;
56            };
57            map.insert(
58                model_id.clone(),
59                ModelPricing {
60                    input,
61                    output,
62                    cache_write_5m: field("cache_creation_input_token_cost").unwrap_or(0.0),
63                    cache_write_1h: field("cache_creation_input_token_cost_above_1hr")
64                        .unwrap_or(0.0),
65                    cache_read: field("cache_read_input_token_cost").unwrap_or(0.0),
66                },
67            );
68        }
69        Self { map }
70    }
71
72    /// Exact-match price for `model_id`; `None` when the id is absent
73    /// (an unknown/new model, or `<synthetic>`).
74    pub fn price(&self, model_id: &str) -> Option<&ModelPricing> {
75        self.map.get(model_id)
76    }
77
78    /// Number of priced models; the refresh path rejects an empty parse
79    /// so a bad fetch never replaces a good cache.
80    pub fn is_empty(&self) -> bool {
81        self.map.is_empty()
82    }
83}
84
85/// Fetch the raw LiteLLM pricing file. Honours an ambient `HTTPS_PROXY`
86/// first, then retries once directly so a dead proxy doesn't block the
87/// fetch. `None` on any network failure (the caller keeps the last-good
88/// cache).
89pub async fn fetch_litellm() -> Option<String> {
90    if let Some(body) = attempt_fetch(false).await {
91        return Some(body);
92    }
93    attempt_fetch(true).await
94}
95
96async fn attempt_fetch(direct: bool) -> Option<String> {
97    let mut builder =
98        crate::http_trust::with_extra_roots(reqwest::Client::builder().timeout(FETCH_TIMEOUT));
99    if direct {
100        builder = builder.no_proxy();
101    }
102    let client = match builder.build() {
103        Ok(client) => client,
104        Err(error) => {
105            tracing::warn!(
106                target: "forge_agent::env::token_usage::pricing",
107                %error,
108                direct,
109                "pricing http client build failed",
110            );
111            return None;
112        }
113    };
114    let response = match client.get(LITELLM_URL).send().await {
115        Ok(response) => response,
116        Err(error) => {
117            tracing::warn!(
118                target: "forge_agent::env::token_usage::pricing",
119                %error,
120                direct,
121                "litellm pricing fetch failed",
122            );
123            return None;
124        }
125    };
126    if !response.status().is_success() {
127        tracing::warn!(
128            target: "forge_agent::env::token_usage::pricing",
129            status = %response.status(),
130            direct,
131            "litellm pricing fetch returned non-success",
132        );
133        return None;
134    }
135    match response.text().await {
136        Ok(body) => Some(body),
137        Err(error) => {
138            tracing::warn!(
139                target: "forge_agent::env::token_usage::pricing",
140                %error,
141                direct,
142                "reading the pricing response body failed",
143            );
144            None
145        }
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    /// Bit-exact float equality: both sides parse the same decimal, so
154    /// the values are identical, not merely close.
155    fn exact(actual: f64, expected: f64) {
156        assert_eq!(actual.to_bits(), expected.to_bits(), "expected {expected}, got {actual}");
157    }
158
159    const SAMPLE: &str = r#"{
160        "sample_spec": {"input_cost_per_token": 1.0, "output_cost_per_token": 1.0},
161        "claude-opus-4-8": {
162            "input_cost_per_token": 5e-06,
163            "output_cost_per_token": 2.5e-05,
164            "cache_creation_input_token_cost": 6.25e-06,
165            "cache_creation_input_token_cost_above_1hr": 1e-05,
166            "cache_read_input_token_cost": 5e-07
167        },
168        "gpt-5-codex": {
169            "input_cost_per_token": 1.25e-06,
170            "output_cost_per_token": 1e-05,
171            "cache_read_input_token_cost": 1.25e-07
172        },
173        "text-embedding-3-small": {"max_tokens": 8191}
174    }"#;
175
176    #[test]
177    fn maps_litellm_fields_onto_model_pricing() {
178        let table = PricingTable::from_litellm_json(SAMPLE);
179        let opus = table.price("claude-opus-4-8").expect("opus priced");
180        exact(opus.input, 5e-6);
181        exact(opus.output, 2.5e-5);
182        exact(opus.cache_write_5m, 6.25e-6);
183        exact(opus.cache_write_1h, 1e-5);
184        exact(opus.cache_read, 5e-7);
185    }
186
187    #[test]
188    fn gpt_entry_without_cache_creation_zeroes_cache_write() {
189        let table = PricingTable::from_litellm_json(SAMPLE);
190        let codex = table.price("gpt-5-codex").expect("codex priced");
191        exact(codex.input, 1.25e-6);
192        exact(codex.cache_read, 1.25e-7);
193        exact(codex.cache_write_5m, 0.0);
194        exact(codex.cache_write_1h, 0.0);
195    }
196
197    #[test]
198    fn skips_doc_entry_and_entries_without_core_costs() {
199        let table = PricingTable::from_litellm_json(SAMPLE);
200        assert!(table.price("sample_spec").is_none(), "the doc entry is skipped");
201        assert!(table.price("text-embedding-3-small").is_none(), "no core cost -> skipped");
202        assert!(table.price("no-such-model").is_none());
203        assert!(table.price("<synthetic>").is_none());
204    }
205
206    #[test]
207    fn malformed_json_is_an_empty_table() {
208        assert!(PricingTable::from_litellm_json("not json").is_empty());
209        assert!(PricingTable::from_litellm_json("{}").is_empty());
210    }
211}
forge_agent/env/token_usage/pricing.rs

forge_agent/env/token_usage/
pricing.rs
