1use serde::{Deserialize, Serialize};
22
23#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
24pub struct Usage {
25 pub prompt_tokens: usize,
26 pub completion_tokens: usize,
27 pub total_tokens: usize,
28 #[serde(skip_serializing_if = "Option::is_none")]
30 pub prompt_per_second: Option<f64>,
31 #[serde(skip_serializing_if = "Option::is_none")]
33 pub predicted_per_second: Option<f64>,
34 #[serde(skip_serializing_if = "Option::is_none")]
36 pub prompt_eval_duration_ms: Option<f64>,
37 #[serde(skip_serializing_if = "Option::is_none")]
41 pub generation_duration_ms: Option<f64>,
42 #[serde(skip_serializing_if = "Option::is_none")]
47 pub time_to_first_token_ms: Option<f64>,
48 #[serde(skip_serializing_if = "Option::is_none")]
53 pub cached_tokens: Option<usize>,
54}
55
56impl Usage {
57 pub fn new(prompt_tokens: usize, completion_tokens: usize) -> Self {
58 Usage {
59 prompt_tokens,
60 completion_tokens,
61 total_tokens: prompt_tokens + completion_tokens,
62 prompt_per_second: None,
63 predicted_per_second: None,
64 prompt_eval_duration_ms: None,
65 generation_duration_ms: None,
66 time_to_first_token_ms: None,
67 cached_tokens: None,
68 }
69 }
70
71 pub fn with_timings(mut self, prompt_secs: f64, predicted_secs: f64) -> Self {
75 self.prompt_eval_duration_ms = Some(prompt_secs * 1000.0);
76 self.generation_duration_ms = Some(predicted_secs * 1000.0);
77 if prompt_secs > 0.0 && self.prompt_tokens > 0 {
78 self.prompt_per_second = Some(self.prompt_tokens as f64 / prompt_secs);
79 }
80 if predicted_secs > 0.0 && self.completion_tokens > 0 {
81 self.predicted_per_second = Some(self.completion_tokens as f64 / predicted_secs);
82 }
83 self
84 }
85
86 pub fn with_ttft(mut self, secs: f64) -> Self {
89 if self.completion_tokens > 0 {
90 self.time_to_first_token_ms = Some(secs * 1000.0);
91 }
92 self
93 }
94
95 pub fn with_cached_tokens(mut self, cached: usize) -> Self {
98 self.cached_tokens = Some(cached);
99 self
100 }
101}
102
103#[cfg(test)]
104mod tests {
105 use super::*;
106
107 #[test]
108 fn totals_are_the_sum_of_the_two_phases() {
109 let usage = Usage::new(7, 3);
110 assert_eq!(usage.total_tokens, 10);
111 }
112
113 #[test]
114 fn phase_durations_stay_separate() {
115 let usage = Usage::new(100, 10).with_timings(1.0, 1.0);
118 assert_eq!(usage.prompt_per_second, Some(100.0));
119 assert_eq!(usage.predicted_per_second, Some(10.0));
120 assert_eq!(usage.prompt_eval_duration_ms, Some(1000.0));
121 assert_eq!(usage.generation_duration_ms, Some(1000.0));
122 }
123
124 #[test]
125 fn zero_length_phases_do_not_become_infinite_rates() {
126 let usage = Usage::new(5, 5).with_timings(0.0, 0.0);
127 assert_eq!(usage.prompt_per_second, None);
128 assert_eq!(usage.predicted_per_second, None);
129 assert_eq!(usage.prompt_eval_duration_ms, Some(0.0));
130 }
131
132 #[test]
133 fn ttft_is_unset_when_nothing_was_generated() {
134 let usage = Usage::new(5, 0).with_ttft(0.25);
135 assert_eq!(usage.time_to_first_token_ms, None);
136 assert_eq!(
137 Usage::new(5, 1).with_ttft(0.25).time_to_first_token_ms,
138 Some(250.0)
139 );
140 }
141
142 #[test]
143 fn untimed_usage_serializes_to_the_plain_openai_shape() {
144 let json = serde_json::to_string(&Usage::new(2, 3)).unwrap();
146 assert_eq!(
147 json,
148 "{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}"
149 );
150 }
151
152 #[test]
153 fn a_prefix_cache_miss_is_distinguishable_from_no_prefix_cache() {
154 assert_eq!(Usage::new(2, 3).cached_tokens, None);
155 assert_eq!(
156 Usage::new(2, 3).with_cached_tokens(0).cached_tokens,
157 Some(0)
158 );
159 }
160}