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#[derive(Debug, Clone, PartialEq, Serialize)]
47#[serde(tag = "kind", content = "detail")]
48pub enum Approximation {
49    UnpricedModel(String),
50    AssumedStandardTier,
51    UnknownModelForTurn,
52}
53
54#[derive(Debug, Clone, PartialEq, Serialize)]
55pub struct CostEstimate {
56    pub total_usd: f64,
57    pub per_model: Vec<ModelCost>,
58    pub tokens: TokenBuckets,
59    pub unpriced_models: Vec<String>,
60    pub approximations: Vec<Approximation>,
61    pub pricing_as_of: String,
62}
63
64/// One billable API call extracted from a transcript. Produced by the dialect
65/// parsers, consumed by the cost engine.
66#[derive(Debug, Clone, PartialEq)]
67pub struct MessageUsage {
68    pub model: String, // verbatim; empty string == unknown (e.g. Codex cleared turn_context)
69    pub provider: Provider,
70    pub namespace: String, // "litellm" in v1
71    pub input_uncached: u64,
72    pub cache_read: u64,
73    pub cache_write_5m: u64,
74    pub cache_write_1h: u64,
75    pub output: u64,
76    pub request_input_tokens: u64, // full billed input for THIS call, for tier selection
77    pub service_tier: Option<String>,
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn provider_serializes_as_lowercase_string() {
86        assert_eq!(
87            serde_json::to_value(Provider::OpenRouter).unwrap(),
88            serde_json::json!("openrouter")
89        );
90        assert_eq!(
91            serde_json::to_value(Provider::Other("bedrock".into())).unwrap(),
92            serde_json::json!("bedrock")
93        );
94        assert_eq!(
95            serde_json::to_value(Provider::OpenAI).unwrap(),
96            serde_json::json!("openai")
97        );
98    }
99
100    #[test]
101    fn cost_estimate_serializes_with_expected_fields() {
102        let est = CostEstimate {
103            total_usd: 1.5,
104            per_model: vec![],
105            tokens: TokenBuckets::default(),
106            unpriced_models: vec![],
107            approximations: vec![Approximation::AssumedStandardTier],
108            pricing_as_of: "2026-06-04".into(),
109        };
110        let v: serde_json::Value = serde_json::to_value(&est).unwrap();
111        assert_eq!(v["total_usd"], 1.5);
112        assert_eq!(v["pricing_as_of"], "2026-06-04");
113        assert_eq!(v["approximations"][0]["kind"], "AssumedStandardTier");
114    }
115}