Skip to main content

obol_core/
model.rs

1//! Public result types + the internal per-message usage record.
2
3use serde::Serialize;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum Provider {
7    Anthropic,
8    OpenAI,
9    OpenRouter,
10    Other(String),
11}
12
13impl Provider {
14    pub fn label(&self) -> &str {
15        match self {
16            Provider::Anthropic => "anthropic",
17            Provider::OpenAI => "openai",
18            Provider::OpenRouter => "openrouter",
19            Provider::Other(s) => s.as_str(),
20        }
21    }
22}
23
24impl serde::Serialize for Provider {
25    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
26        s.serialize_str(self.label())
27    }
28}
29
30#[derive(Debug, Clone, Default, PartialEq, Serialize)]
31pub struct TokenBuckets {
32    pub input: u64, // uncached input
33    pub output: u64,
34    pub cache_read: u64,
35    pub cache_write: u64, // 5m + 1h combined, for the summary
36}
37
38#[derive(Debug, Clone, PartialEq, Serialize)]
39pub struct ModelCost {
40    pub model: String,
41    pub provider: Provider,
42    pub tokens: TokenBuckets,
43    pub subtotal_usd: f64,
44}
45
46/// Which snapshot priced this estimate: the one compiled into the library, or one
47/// read from disk (`refresh`ed).
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum PricingSource {
50    Bundled,
51    Local,
52}
53
54impl serde::Serialize for PricingSource {
55    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
56        s.serialize_str(match self {
57            PricingSource::Bundled => "bundled",
58            PricingSource::Local => "local",
59        })
60    }
61}
62
63#[derive(Debug, Clone, PartialEq, Serialize)]
64#[serde(tag = "kind", content = "detail")]
65pub enum Approximation {
66    UnpricedModel(String),
67    AssumedStandardTier,
68    UnknownModelForTurn,
69}
70
71#[derive(Debug, Clone, PartialEq, Serialize)]
72pub struct CostEstimate {
73    pub total_usd: f64,
74    pub per_model: Vec<ModelCost>,
75    pub tokens: TokenBuckets,
76    pub unpriced_models: Vec<String>,
77    pub approximations: Vec<Approximation>,
78    pub pricing_as_of: String,
79    pub pricing_source: PricingSource,
80}
81
82/// One billable API call extracted from a transcript. Produced by the dialect
83/// parsers, consumed by the cost engine.
84#[derive(Debug, Clone, PartialEq)]
85pub struct MessageUsage {
86    pub model: String, // verbatim; empty string == unknown (e.g. Codex cleared turn_context)
87    pub provider: Provider,
88    pub namespace: String, // "litellm" in v1
89    pub input_uncached: u64,
90    pub cache_read: u64,
91    pub cache_write_5m: u64,
92    pub cache_write_1h: u64,
93    pub output: u64,
94    pub request_input_tokens: u64, // full billed input for THIS call, for tier selection
95    pub service_tier: Option<String>,
96    /// Provider-reported all-in cost for this call, in USD, when the transcript
97    /// carries one (e.g. Pi's `usage.cost.total`). `Some` is ground truth and is
98    /// preferred over list-price math; `None` means price from the tables.
99    pub native_cost_usd: Option<f64>,
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn provider_serializes_as_lowercase_string() {
108        assert_eq!(
109            serde_json::to_value(Provider::OpenRouter).unwrap(),
110            serde_json::json!("openrouter")
111        );
112        assert_eq!(
113            serde_json::to_value(Provider::Other("bedrock".into())).unwrap(),
114            serde_json::json!("bedrock")
115        );
116        assert_eq!(
117            serde_json::to_value(Provider::OpenAI).unwrap(),
118            serde_json::json!("openai")
119        );
120    }
121
122    #[test]
123    fn cost_estimate_serializes_with_expected_fields() {
124        let est = CostEstimate {
125            total_usd: 1.5,
126            per_model: vec![],
127            tokens: TokenBuckets::default(),
128            unpriced_models: vec![],
129            approximations: vec![Approximation::AssumedStandardTier],
130            pricing_as_of: "2026-06-04".into(),
131            pricing_source: PricingSource::Bundled,
132        };
133        let v: serde_json::Value = serde_json::to_value(&est).unwrap();
134        assert_eq!(v["total_usd"], 1.5);
135        assert_eq!(v["pricing_as_of"], "2026-06-04");
136        assert_eq!(v["approximations"][0]["kind"], "AssumedStandardTier");
137        assert_eq!(v["pricing_source"], "bundled");
138    }
139
140    #[test]
141    fn pricing_source_serializes_lowercase() {
142        assert_eq!(
143            serde_json::to_value(PricingSource::Bundled).unwrap(),
144            serde_json::json!("bundled")
145        );
146        assert_eq!(
147            serde_json::to_value(PricingSource::Local).unwrap(),
148            serde_json::json!("local")
149        );
150    }
151}