Skip to main content

crabllm_core/
usage.rs

1use serde::{Deserialize, Serialize};
2use std::collections::BTreeMap;
3
4/// Canonical token usage. Axes are disjoint — billing computes cost as a sum
5/// over (axis, rate) pairs without any subtraction or clamping.
6///
7/// Input axes
8/// - `input_tokens`: uncached new input (the portion the model has to read fresh)
9/// - `cache_read_tokens`: served from prompt cache (cheaper than base input)
10/// - `cache_write_tokens`: written to prompt cache this turn (often *more*
11///   expensive than base input — Anthropic charges ~1.25× for cache writes)
12///
13/// Output axes
14/// - `output_tokens`: regular completion tokens, excluding reasoning
15/// - `reasoning_tokens`: thinking/reasoning tokens (may have a separate rate)
16///
17/// Side-channel
18/// - `server_tool_calls`: per-call billing for upstream-side tools like web
19///   search. Keyed by tool name (`"web_search"`, etc.).
20#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
21#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
22pub struct Usage {
23    pub input_tokens: u32,
24    #[serde(default)]
25    pub cache_read_tokens: u32,
26    #[serde(default)]
27    pub cache_write_tokens: u32,
28    pub output_tokens: u32,
29    #[serde(default)]
30    pub reasoning_tokens: u32,
31    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
32    pub server_tool_calls: BTreeMap<String, u32>,
33}
34
35impl Usage {
36    /// Total prompt size (cached + uncached + writes). Equivalent to OpenAI's
37    /// `prompt_tokens`.
38    pub fn prompt_tokens(&self) -> u32 {
39        self.input_tokens + self.cache_read_tokens + self.cache_write_tokens
40    }
41
42    /// Total completion size including reasoning. Equivalent to OpenAI's
43    /// `completion_tokens`.
44    pub fn completion_tokens(&self) -> u32 {
45        self.output_tokens + self.reasoning_tokens
46    }
47
48    /// Prompt plus completion — equivalent to OpenAI's `total_tokens`.
49    pub fn total_tokens(&self) -> u32 {
50        self.prompt_tokens() + self.completion_tokens()
51    }
52}
53
54mod peek {
55    use crate::types::{AnthropicUsage, GeminiUsage, OpenAiUsage};
56    use serde::Deserialize;
57
58    #[derive(Deserialize)]
59    pub struct OpenAi {
60        pub usage: Option<OpenAiUsage>,
61    }
62
63    #[derive(Deserialize)]
64    pub struct Anthropic {
65        pub usage: Option<AnthropicUsage>,
66    }
67
68    #[derive(Deserialize)]
69    #[serde(rename_all = "camelCase")]
70    pub struct Gemini {
71        pub usage_metadata: Option<GeminiUsage>,
72    }
73}
74
75impl From<&[u8]> for Usage {
76    fn from(raw: &[u8]) -> Self {
77        if let Ok(peek::OpenAi { usage: Some(u) }) = crate::json::from_slice(raw)
78            && (u.prompt_tokens > 0 || u.completion_tokens > 0)
79        {
80            return Usage::from(&u);
81        }
82
83        if let Ok(peek::Anthropic { usage: Some(u) }) = crate::json::from_slice(raw)
84            && (u.input_tokens > 0 || u.output_tokens > 0)
85        {
86            return Usage::from(&u);
87        }
88
89        if let Ok(peek::Gemini {
90            usage_metadata: Some(u),
91        }) = crate::json::from_slice(raw)
92            && u.total_token_count > 0
93        {
94            return Usage::from(&u);
95        }
96
97        Usage::default()
98    }
99}