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}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn provider_serializes_as_lowercase_string() {
104        assert_eq!(
105            serde_json::to_value(Provider::OpenRouter).unwrap(),
106            serde_json::json!("openrouter")
107        );
108        assert_eq!(
109            serde_json::to_value(Provider::Other("bedrock".into())).unwrap(),
110            serde_json::json!("bedrock")
111        );
112        assert_eq!(
113            serde_json::to_value(Provider::OpenAI).unwrap(),
114            serde_json::json!("openai")
115        );
116    }
117
118    #[test]
119    fn cost_estimate_serializes_with_expected_fields() {
120        let est = CostEstimate {
121            total_usd: 1.5,
122            per_model: vec![],
123            tokens: TokenBuckets::default(),
124            unpriced_models: vec![],
125            approximations: vec![Approximation::AssumedStandardTier],
126            pricing_as_of: "2026-06-04".into(),
127            pricing_source: PricingSource::Bundled,
128        };
129        let v: serde_json::Value = serde_json::to_value(&est).unwrap();
130        assert_eq!(v["total_usd"], 1.5);
131        assert_eq!(v["pricing_as_of"], "2026-06-04");
132        assert_eq!(v["approximations"][0]["kind"], "AssumedStandardTier");
133        assert_eq!(v["pricing_source"], "bundled");
134    }
135
136    #[test]
137    fn pricing_source_serializes_lowercase() {
138        assert_eq!(
139            serde_json::to_value(PricingSource::Bundled).unwrap(),
140            serde_json::json!("bundled")
141        );
142        assert_eq!(
143            serde_json::to_value(PricingSource::Local).unwrap(),
144            serde_json::json!("local")
145        );
146    }
147}