use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CompletionTokensDetails {
pub reasoning_tokens: usize,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Usage {
pub prompt_tokens: usize,
pub completion_tokens: usize,
pub total_tokens: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_tokens_details: Option<CompletionTokensDetails>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_per_second: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub predicted_per_second: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_eval_duration_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub generation_duration_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub time_to_first_token_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cached_tokens: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub acceptance_length: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub draft_tokens: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub accepted_draft_tokens: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub draft_accept_rate_per_position: Option<Vec<f64>>,
}
impl Usage {
pub fn new(prompt_tokens: usize, completion_tokens: usize) -> Self {
Usage {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens + completion_tokens,
completion_tokens_details: None,
prompt_per_second: None,
predicted_per_second: None,
prompt_eval_duration_ms: None,
generation_duration_ms: None,
time_to_first_token_ms: None,
cached_tokens: None,
acceptance_length: None,
draft_tokens: None,
accepted_draft_tokens: None,
draft_accept_rate_per_position: None,
}
}
pub fn with_timings(mut self, prompt_secs: f64, predicted_secs: f64) -> Self {
self.prompt_eval_duration_ms = Some(prompt_secs * 1000.0);
self.generation_duration_ms = Some(predicted_secs * 1000.0);
if prompt_secs > 0.0 && self.prompt_tokens > 0 {
self.prompt_per_second = Some(self.prompt_tokens as f64 / prompt_secs);
}
if predicted_secs > 0.0 && self.completion_tokens > 0 {
self.predicted_per_second = Some(self.completion_tokens as f64 / predicted_secs);
}
self
}
pub fn with_ttft(mut self, secs: f64) -> Self {
if self.completion_tokens > 0 {
self.time_to_first_token_ms = Some(secs * 1000.0);
}
self
}
pub fn with_reasoning_tokens(mut self, reasoning: usize) -> Self {
self.completion_tokens_details = Some(CompletionTokensDetails {
reasoning_tokens: reasoning,
});
self
}
pub fn with_cached_tokens(mut self, cached: usize) -> Self {
self.cached_tokens = Some(cached);
self
}
pub fn with_speculation(
mut self,
verification_steps: usize,
accepted: usize,
drafted: usize,
per_position: Vec<f64>,
) -> Self {
if verification_steps > 0 {
self.acceptance_length =
Some(self.completion_tokens as f64 / verification_steps as f64);
}
self.accepted_draft_tokens = Some(accepted);
self.draft_tokens = Some(drafted);
self.draft_accept_rate_per_position = Some(per_position);
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn totals_are_the_sum_of_the_two_phases() {
let usage = Usage::new(7, 3);
assert_eq!(usage.total_tokens, 10);
}
#[test]
fn phase_durations_stay_separate() {
let usage = Usage::new(100, 10).with_timings(1.0, 1.0);
assert_eq!(usage.prompt_per_second, Some(100.0));
assert_eq!(usage.predicted_per_second, Some(10.0));
assert_eq!(usage.prompt_eval_duration_ms, Some(1000.0));
assert_eq!(usage.generation_duration_ms, Some(1000.0));
}
#[test]
fn zero_length_phases_do_not_become_infinite_rates() {
let usage = Usage::new(5, 5).with_timings(0.0, 0.0);
assert_eq!(usage.prompt_per_second, None);
assert_eq!(usage.predicted_per_second, None);
assert_eq!(usage.prompt_eval_duration_ms, Some(0.0));
}
#[test]
fn ttft_is_unset_when_nothing_was_generated() {
let usage = Usage::new(5, 0).with_ttft(0.25);
assert_eq!(usage.time_to_first_token_ms, None);
assert_eq!(
Usage::new(5, 1).with_ttft(0.25).time_to_first_token_ms,
Some(250.0)
);
}
#[test]
fn untimed_usage_serializes_to_the_plain_openai_shape() {
let json = serde_json::to_string(&Usage::new(2, 3)).unwrap();
assert_eq!(
json,
"{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}"
);
}
#[test]
fn a_non_speculative_request_reports_no_acceptance_length_at_all() {
let plain = Usage::new(10, 5);
assert_eq!(plain.acceptance_length, None);
assert_eq!(plain.draft_tokens, None);
let json = serde_json::to_value(&plain).unwrap();
assert!(json.get("acceptance_length").is_none());
assert!(json.get("draft_accept_rate_per_position").is_none());
}
#[test]
fn acceptance_length_is_completion_tokens_per_verification_step() {
let usage = Usage::new(20, 12).with_speculation(5, 7, 10, vec![0.9, 0.6, 0.2]);
assert_eq!(usage.acceptance_length, Some(2.4));
assert_eq!(usage.accepted_draft_tokens, Some(7));
assert_eq!(usage.draft_tokens, Some(10));
assert_eq!(
usage.draft_accept_rate_per_position,
Some(vec![0.9, 0.6, 0.2])
);
}
#[test]
fn speculation_that_verified_nothing_reports_no_length_rather_than_infinity() {
let usage = Usage::new(20, 0).with_speculation(0, 0, 0, Vec::new());
assert_eq!(usage.acceptance_length, None);
assert_eq!(usage.draft_tokens, Some(0));
}
#[test]
fn a_prefix_cache_miss_is_distinguishable_from_no_prefix_cache() {
assert_eq!(Usage::new(2, 3).cached_tokens, None);
assert_eq!(
Usage::new(2, 3).with_cached_tokens(0).cached_tokens,
Some(0)
);
}
}