obol_core/pricing/
store.rs

1use super::ModelPrice;
2use crate::error::ObolError;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::path::{Path, PathBuf};
6
7/// All price tables, keyed by namespace ("litellm") then verbatim model string.
8#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
9pub struct PriceStore {
10    pub as_of: String,
11    pub namespaces: HashMap<String, HashMap<String, ModelPrice>>,
12}
13
14impl PriceStore {
15    pub fn lookup(&self, namespace: &str, model: &str) -> Option<&ModelPrice> {
16        self.namespaces.get(namespace)?.get(model)
17    }
18
19    pub fn from_json(bytes: &[u8]) -> Result<Self, ObolError> {
20        Ok(serde_json::from_slice(bytes)?)
21    }
22
23    pub fn load(path: &Path) -> Result<Self, ObolError> {
24        if !path.exists() {
25            return Err(ObolError::PricingTablesMissing(path.to_path_buf()));
26        }
27        Self::from_json(&std::fs::read(path)?)
28    }
29
30    pub fn save(&self, path: &Path) -> Result<(), ObolError> {
31        if let Some(parent) = path.parent() {
32            std::fs::create_dir_all(parent)?;
33        }
34        std::fs::write(path, serde_json::to_vec_pretty(self)?)?;
35        Ok(())
36    }
37}
38
39/// Directory holding price snapshots: $OBOL_PRICING_DIR, else $XDG_DATA_HOME/obol,
40/// else $HOME/.local/share/obol.
41pub fn pricing_dir() -> PathBuf {
42    if let Ok(d) = std::env::var("OBOL_PRICING_DIR") {
43        return PathBuf::from(d);
44    }
45    let base = std::env::var("XDG_DATA_HOME")
46        .map(PathBuf::from)
47        .unwrap_or_else(|_| {
48            let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
49            PathBuf::from(home).join(".local").join("share")
50        });
51    base.join("obol")
52}
53
54/// The active snapshot the library reads.
55pub fn current_path() -> PathBuf {
56    pricing_dir().join("current.json")
57}
58
59/// The price snapshot compiled into the library — the out-of-the-box fallback used
60/// when no on-disk snapshot is newer (see `lib::estimate_cost`).
61pub fn embedded() -> Result<PriceStore, ObolError> {
62    const BYTES: &[u8] = include_bytes!("../../prices/bundled.json");
63    PriceStore::from_json(BYTES)
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    fn store() -> PriceStore {
71        let mut litellm = HashMap::new();
72        litellm.insert(
73            "claude-opus-4-8".to_string(),
74            ModelPrice {
75                input: 5.0,
76                output: 25.0,
77                cache_read: 0.5,
78                cache_write: 6.25,
79                cache_write_1h: Some(10.0),
80                tier_boundary: None,
81                input_above: None,
82                output_above: None,
83                cache_read_above: None,
84                cache_write_above: None,
85            },
86        );
87        let mut namespaces = HashMap::new();
88        namespaces.insert("litellm".to_string(), litellm);
89        PriceStore {
90            as_of: "2026-06-04".into(),
91            namespaces,
92        }
93    }
94
95    #[test]
96    fn lookup_finds_model_in_namespace() {
97        let s = store();
98        assert!(s.lookup("litellm", "claude-opus-4-8").is_some());
99        assert!(s.lookup("litellm", "nonsense").is_none());
100        assert!(s.lookup("openrouter", "claude-opus-4-8").is_none());
101    }
102
103    #[test]
104    fn save_then_load_roundtrips() {
105        let dir = std::env::temp_dir().join(format!("obol-test-{}", std::process::id()));
106        let path = dir.join("current.json");
107        store().save(&path).unwrap();
108        let loaded = PriceStore::load(&path).unwrap();
109        assert_eq!(loaded, store());
110        std::fs::remove_dir_all(&dir).ok();
111    }
112
113    #[test]
114    fn load_missing_is_pricing_tables_missing() {
115        let path = PathBuf::from("/nonexistent/obol/current.json");
116        match PriceStore::load(&path) {
117            Err(ObolError::PricingTablesMissing(_)) => {}
118            other => panic!("expected PricingTablesMissing, got {other:?}"),
119        }
120    }
121
122    #[test]
123    fn pricing_dir_honors_env_override() {
124        std::env::set_var("OBOL_PRICING_DIR", "/tmp/obol-x");
125        assert_eq!(pricing_dir(), PathBuf::from("/tmp/obol-x"));
126        std::env::remove_var("OBOL_PRICING_DIR");
127    }
128
129    #[test]
130    fn embedded_snapshot_loads_and_has_models() {
131        let s = embedded().expect("embedded snapshot parses");
132        assert!(!s.as_of.is_empty(), "embedded snapshot must carry an as_of");
133        assert!(
134            s.lookup("litellm", "claude-opus-4-8").is_some(),
135            "embedded snapshot should price a known model"
136        );
137    }
138
139    #[test]
140    fn embedded_snapshot_bundles_litellm_and_openrouter() {
141        let s = embedded().expect("embedded snapshot parses");
142        assert!(
143            s.namespaces.contains_key("litellm"),
144            "embedded snapshot must carry the litellm namespace"
145        );
146        let or = s
147            .namespaces
148            .get("openrouter")
149            .expect("embedded snapshot must carry the openrouter namespace");
150        assert!(!or.is_empty(), "openrouter namespace should be non-empty");
151        // OpenRouter keys are `<vendor>/<model>` — a run billed through OpenRouter
152        // prices from this namespace out of the box.
153        assert!(
154            or.keys().any(|k| k.contains('/')),
155            "openrouter keys are vendor/model form"
156        );
157    }
158}
obol_core/pricing/store.rs

obol_core/pricing/
store.rs
