obol_core/
lib.rs

1//! obol-core: parse agent transcripts and estimate token cost.
2
3pub mod cost;
4pub mod error;
5pub mod model;
6pub mod pricing;
7pub mod transcript;
8
9pub use error::ObolError;
10pub use model::{Approximation, CostEstimate, MessageUsage, ModelCost, Provider, TokenBuckets};
11pub use transcript::Dialect;
12
13use std::path::{Path, PathBuf};
14
15/// Where to read the transcript bytes from. The dialect hint is a separate
16/// argument to `estimate_cost`, so it applies equally to a path or to bytes.
17pub enum Source<'a> {
18    Path(&'a Path),
19    Bytes(&'a [u8]),
20}
21
22/// Report from a pricing refresh.
23#[derive(Debug, serde::Serialize)]
24pub struct RefreshReport {
25    pub models: usize,
26    pub as_of: String,
27    pub written_to: PathBuf,
28}
29
30/// Estimate the cost of a transcript. `dialect` is an optional hint; when `None`,
31/// the dialect is detected from the content. Loads the active price snapshot from
32/// disk (errors with `PricingTablesMissing` if absent).
33pub fn estimate_cost(source: Source, dialect: Option<Dialect>) -> Result<CostEstimate, ObolError> {
34    let store = pricing::PriceStore::load(&pricing::current_path())?;
35    let bytes: Vec<u8> = match source {
36        Source::Path(p) => std::fs::read(p)?,
37        Source::Bytes(b) => b.to_vec(),
38    };
39    let dialect = match dialect {
40        Some(d) => d,
41        None => transcript::detect(&bytes)?,
42    };
43    let usages = transcript::parse(&bytes, dialect)?;
44    Ok(cost::estimate(&usages, &store))
45}
46
47/// Fetch the LiteLLM sheet and write it as the active snapshot. `as_of` is the
48/// caller's date string (the library has no clock).
49pub fn refresh_pricing_tables(as_of: &str) -> Result<RefreshReport, ObolError> {
50    let mut store = pricing::refresh::fetch_litellm(as_of)?; // {litellm: …}
51    let openrouter = pricing::refresh::fetch_openrouter()?;
52    store
53        .namespaces
54        .insert("openrouter".to_string(), openrouter);
55    let models: usize = store.namespaces.values().map(|m| m.len()).sum();
56    let dir = pricing::pricing_dir();
57    store.save(&dir.join(format!("prices-{as_of}.json")))?;
58    let current = pricing::current_path();
59    store.save(&current)?;
60    Ok(RefreshReport {
61        models,
62        as_of: as_of.to_string(),
63        written_to: current,
64    })
65}
66
67#[cfg(test)]
68mod api_tests {
69    use super::*;
70
71    #[test]
72    fn estimate_cost_on_bytes_with_missing_tables_errors() {
73        std::env::set_var("OBOL_PRICING_DIR", "/nonexistent/obol-xyz");
74        let data = include_bytes!("../tests/fixtures/claude-mini.jsonl");
75        let r = estimate_cost(Source::Bytes(data), Some(Dialect::Claude));
76        assert!(matches!(r, Err(ObolError::PricingTablesMissing(_))));
77        std::env::remove_var("OBOL_PRICING_DIR");
78    }
79
80    #[test]
81    fn estimate_cost_end_to_end_with_seeded_store() {
82        let dir = std::env::temp_dir().join(format!("obol-api-{}", std::process::id()));
83        std::env::set_var("OBOL_PRICING_DIR", &dir);
84        // seed the store from the sample sheet
85        let store = pricing::refresh::normalize_litellm(
86            include_bytes!("../tests/fixtures/litellm-sample.json"),
87            "2026-06-04",
88        )
89        .unwrap();
90        store.save(&pricing::current_path()).unwrap();
91
92        let data = include_bytes!("../tests/fixtures/claude-mini.jsonl");
93        let est = estimate_cost(Source::Bytes(data), Some(Dialect::Claude)).unwrap();
94        assert!(est.total_usd > 0.0);
95        assert_eq!(est.pricing_as_of, "2026-06-04");
96
97        std::fs::remove_dir_all(&dir).ok();
98        std::env::remove_var("OBOL_PRICING_DIR");
99    }
100
101    #[test]
102    fn estimate_cost_from_path_with_autodetect() {
103        let dir = std::env::temp_dir().join(format!("obol-path-{}", std::process::id()));
104        std::env::set_var("OBOL_PRICING_DIR", &dir);
105        let store = pricing::refresh::normalize_litellm(
106            include_bytes!("../tests/fixtures/litellm-sample.json"),
107            "2026-06-04",
108        )
109        .unwrap();
110        store.save(&pricing::current_path()).unwrap();
111
112        // Write the Claude fixture to a real file and price it via Source::Path
113        // with NO dialect hint — exercises both the path input and auto-detect.
114        let transcript = dir.join("session.jsonl");
115        std::fs::write(
116            &transcript,
117            include_bytes!("../tests/fixtures/claude-mini.jsonl"),
118        )
119        .unwrap();
120        let est = estimate_cost(Source::Path(&transcript), None).unwrap();
121        assert!(est.total_usd > 0.0);
122
123        std::fs::remove_dir_all(&dir).ok();
124        std::env::remove_var("OBOL_PRICING_DIR");
125    }
126
127    #[test]
128    fn refresh_report_serializes() {
129        let r = RefreshReport {
130            models: 7,
131            as_of: "2026-06-05".into(),
132            written_to: "/x/current.json".into(),
133        };
134        let v = serde_json::to_value(&r).unwrap();
135        assert_eq!(v["models"], 7);
136        assert_eq!(v["as_of"], "2026-06-05");
137        assert_eq!(v["written_to"], "/x/current.json");
138    }
139}
obol_core/lib.rs

obol_core/
lib.rs
