llm/usage/
tokens.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use std::fmt;
4use std::ops::{Add, AddAssign};
5
6/// A count of tokens. Sums saturate instead of wrapping.
7#[repr(transparent)]
8#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, JsonSchema)]
9#[serde(transparent)]
10#[schemars(transparent)]
11pub struct Tokens(u64);
12
13impl Tokens {
14    pub const ZERO: Self = Self(0);
15
16    pub const fn new(value: u64) -> Self {
17        Self(value)
18    }
19
20    pub const fn get(self) -> u64 {
21        self.0
22    }
23
24    pub const fn is_zero(self) -> bool {
25        self.0 == 0
26    }
27
28    pub const fn saturating_add(self, rhs: Self) -> Self {
29        Self(self.0.saturating_add(rhs.0))
30    }
31
32    pub const fn saturating_sub(self, rhs: Self) -> Self {
33        Self(self.0.saturating_sub(rhs.0))
34    }
35}
36
37impl Add for Tokens {
38    type Output = Self;
39    fn add(self, rhs: Self) -> Self::Output {
40        self.saturating_add(rhs)
41    }
42}
43
44impl AddAssign for Tokens {
45    fn add_assign(&mut self, rhs: Self) {
46        *self = self.saturating_add(rhs);
47    }
48}
49
50impl From<u32> for Tokens {
51    fn from(value: u32) -> Self {
52        Self(value.into())
53    }
54}
55
56impl From<Tokens> for u64 {
57    fn from(value: Tokens) -> Self {
58        value.0
59    }
60}
61
62impl TryFrom<Tokens> for i64 {
63    type Error = std::num::TryFromIntError;
64    fn try_from(value: Tokens) -> Result<Self, Self::Error> {
65        value.0.try_into()
66    }
67}
68
69impl From<Tokens> for f64 {
70    #[allow(clippy::cast_precision_loss)]
71    fn from(value: Tokens) -> Self {
72        value.0 as f64
73    }
74}
75
76impl fmt::Display for Tokens {
77    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
78        self.0.fmt(formatter)
79    }
80}
81
82/// Token counts for a single LLM call or an aggregate of calls. Providers fill
83/// in only the dimensions they report; the rest stay `None`.
84#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
85pub struct TokenUsage {
86    /// All input tokens, including cache reads and cache writes.
87    pub input_tokens: Tokens,
88    pub output_tokens: Tokens,
89    #[serde(default)]
90    pub cache_read_tokens: Option<Tokens>,
91    #[serde(default)]
92    pub cache_creation_tokens: Option<Tokens>,
93    #[serde(default)]
94    pub input_audio_tokens: Option<Tokens>,
95    #[serde(default)]
96    pub input_video_tokens: Option<Tokens>,
97    #[serde(default)]
98    pub reasoning_tokens: Option<Tokens>,
99    #[serde(default)]
100    pub output_audio_tokens: Option<Tokens>,
101    #[serde(default)]
102    pub accepted_prediction_tokens: Option<Tokens>,
103    #[serde(default)]
104    pub rejected_prediction_tokens: Option<Tokens>,
105}
106
107impl TokenUsage {
108    pub fn new(input_tokens: u64, output_tokens: u64) -> Self {
109        Self { input_tokens: Tokens::new(input_tokens), output_tokens: Tokens::new(output_tokens), ..Self::default() }
110    }
111
112    pub fn is_zero(self) -> bool {
113        let reported = [
114            self.cache_read_tokens,
115            self.cache_creation_tokens,
116            self.input_audio_tokens,
117            self.input_video_tokens,
118            self.reasoning_tokens,
119            self.output_audio_tokens,
120            self.accepted_prediction_tokens,
121            self.rejected_prediction_tokens,
122        ];
123        self.total_tokens().is_zero() && reported.iter().all(|dimension| dimension.unwrap_or_default().is_zero())
124    }
125
126    pub fn total_tokens(self) -> Tokens {
127        self.input_tokens + self.output_tokens
128    }
129
130    pub(crate) fn from_uncached_input(
131        input: u32,
132        output: u32,
133        cache_read: Option<u32>,
134        cache_write: Option<u32>,
135    ) -> Self {
136        let cache_read = cache_read.map(Tokens::from);
137        let cache_creation = cache_write.map(Tokens::from);
138        Self {
139            input_tokens: Tokens::from(input) + cache_read.unwrap_or_default() + cache_creation.unwrap_or_default(),
140            output_tokens: output.into(),
141            cache_read_tokens: cache_read,
142            cache_creation_tokens: cache_creation,
143            ..Self::default()
144        }
145    }
146}
147
148impl Add for TokenUsage {
149    type Output = Self;
150    fn add(self, rhs: Self) -> Self::Output {
151        Self {
152            input_tokens: self.input_tokens + rhs.input_tokens,
153            output_tokens: self.output_tokens + rhs.output_tokens,
154            cache_read_tokens: add_reported(self.cache_read_tokens, rhs.cache_read_tokens),
155            cache_creation_tokens: add_reported(self.cache_creation_tokens, rhs.cache_creation_tokens),
156            input_audio_tokens: add_reported(self.input_audio_tokens, rhs.input_audio_tokens),
157            input_video_tokens: add_reported(self.input_video_tokens, rhs.input_video_tokens),
158            reasoning_tokens: add_reported(self.reasoning_tokens, rhs.reasoning_tokens),
159            output_audio_tokens: add_reported(self.output_audio_tokens, rhs.output_audio_tokens),
160            accepted_prediction_tokens: add_reported(self.accepted_prediction_tokens, rhs.accepted_prediction_tokens),
161            rejected_prediction_tokens: add_reported(self.rejected_prediction_tokens, rhs.rejected_prediction_tokens),
162        }
163    }
164}
165
166impl AddAssign for TokenUsage {
167    fn add_assign(&mut self, rhs: Self) {
168        *self = *self + rhs;
169    }
170}
171
172fn add_reported(lhs: Option<Tokens>, rhs: Option<Tokens>) -> Option<Tokens> {
173    match (lhs, rhs) {
174        (None, None) => None,
175        _ => Some(lhs.unwrap_or_default() + rhs.unwrap_or_default()),
176    }
177}
llm/usage/tokens.rs

llm/usage/
tokens.rs
