use crate::generator::CompletionParameters;
use crate::message::{ContentPart, Message, MessageContent};
use crate::provider::bpe;
use crate::generator::ModelRates;
use crate::provider::wire::TokenPrice;
pub const SAFETY_MULTIPLIER: f64 = 1.6;
const TOKENS_PER_MESSAGE: u32 = 4;
const TOKENS_PER_STILL_IMAGE: u32 = 1_600;
const TOKENS_PER_VIDEO_FRAME: u32 = 258;
const _: () = assert!(
TOKENS_PER_STILL_IMAGE > TOKENS_PER_VIDEO_FRAME,
"a still of unknown size must cost more than a resampled video frame"
);
const VIDEO_FRAMES_PER_SECOND: f64 = 1.0;
const AUDIO_TOKENS_PER_SECOND: f64 = 32.0;
const DEFAULT_MEDIA_SECONDS: f64 = 60.0;
fn media_seconds(declared: Option<f64>) -> f64 {
match declared {
Some(secs) if secs.is_finite() && secs >= 0.0 => secs,
_ => DEFAULT_MEDIA_SECONDS,
}
}
fn audio_tokens_for(secs: f64) -> u64 {
(secs * AUDIO_TOKENS_PER_SECOND).ceil() as u64
}
fn video_frame_tokens_for(secs: f64) -> u64 {
let frames = (secs * VIDEO_FRAMES_PER_SECOND).ceil() as u64;
frames.saturating_mul(u64::from(TOKENS_PER_VIDEO_FRAME))
}
pub fn estimate_prompt_tokens(messages: &[Message]) -> PromptEstimate {
let mut text_tokens = 0u64;
let mut image_tokens = 0u64;
let mut audio_tokens = 0u64;
for message in messages {
text_tokens += u64::from(TOKENS_PER_MESSAGE);
match &message.content {
MessageContent::Text(t) => text_tokens += bpe::count_tokens(t) as u64,
MessageContent::Parts(parts) => {
for part in parts {
match part {
ContentPart::Text { text } => text_tokens += bpe::count_tokens(text) as u64,
ContentPart::Image { .. } => {
image_tokens += u64::from(TOKENS_PER_STILL_IMAGE)
}
ContentPart::Video { video_url } => {
let secs = media_seconds(video_url.duration_secs);
image_tokens = image_tokens.saturating_add(video_frame_tokens_for(secs));
audio_tokens = audio_tokens.saturating_add(audio_tokens_for(secs));
}
ContentPart::Audio { input_audio } => {
audio_tokens = audio_tokens
.saturating_add(audio_tokens_for(media_seconds(input_audio.duration_secs)))
}
}
}
}
}
}
PromptEstimate {
text_tokens: (text_tokens as f64 * SAFETY_MULTIPLIER).ceil() as u64,
image_tokens,
audio_tokens,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct PromptEstimate {
pub text_tokens: u64,
pub image_tokens: u64,
pub audio_tokens: u64,
}
impl PromptEstimate {
pub fn total_tokens(&self) -> u64 {
self.text_tokens
.saturating_add(self.image_tokens)
.saturating_add(self.audio_tokens)
}
fn input_cost_usd(self, context_length: u32, price: &TokenPrice) -> f64 {
let mut buckets: [(u64, f64); 3] = [
(self.audio_tokens, price.audio_rate()),
(self.image_tokens, price.image_rate()),
(self.text_tokens, price.input_per_mtok),
];
buckets.sort_by(|a, b| b.1.total_cmp(&a.1));
let mut room = u64::from(context_length);
let mut cost = 0.0;
for (tokens, rate) in buckets {
let billed = tokens.min(room);
cost += billed as f64 * rate;
room -= billed;
}
cost
}
}
fn max_output_tokens(params: &CompletionParameters, rates: &ModelRates) -> u64 {
let ceiling = rates.max_completion_tokens.unwrap_or(rates.context_length);
let visible = params.max_tokens.unwrap_or(ceiling).min(ceiling);
let thinking = match ¶ms.reasoning {
None => 0,
Some(r) if r.effort.as_deref() == Some("none") => 0,
Some(r) => r.max_tokens.unwrap_or(ceiling).min(ceiling),
};
u64::from(visible) + u64::from(thinking)
}
pub fn estimate_cost_usd(
messages: &[Message],
params: &CompletionParameters,
rates: &ModelRates,
) -> f64 {
let price = &rates.price;
let input = estimate_prompt_tokens(messages).input_cost_usd(rates.context_length, price);
let output = max_output_tokens(params, rates) as f64 * price.output_per_mtok;
(input + output) / 1_000_000.0
}
#[cfg(test)]
mod tests {
use super::*;
use crate::generator::ReasoningConfig;
use crate::message::AudioInput;
use crate::provider::wire::{TokenPrice, AUDIO_RATE_FALLBACK_MULTIPLE};
fn user(text: &str) -> Message {
Message::user(text)
}
fn parts(parts: Vec<ContentPart>) -> Message {
Message::user(MessageContent::Parts(parts))
}
fn rates(max_completion: Option<u32>, context: u32) -> ModelRates {
ModelRates {
price: TokenPrice::new(3.0, 15.0),
max_completion_tokens: max_completion,
context_length: context,
}
}
#[test]
fn the_estimate_is_an_upper_bound_on_real_claude_token_counts() {
let corpus: &[(&str, u64)] = &[
("The quick brown fox jumps over the lazy dog. Tokenization is not standardized across model families, which makes cross-provider estimation inherently approximate.", 39),
("fn main(){let x:Vec<u32>=(0..10).filter(|n|n%2==0).map(|n|n*n).collect();println!(\"{:?}\",x);}", 55),
("{\"a\": 1, \"bb\": [1, 2, 3], \"ccc\": {\"d\": true, \"e\": null}}", 39),
("这是一段中文文本,用于测试分词器的行为差异。", 30),
];
for (text, truth) in corpus {
let est = estimate_prompt_tokens(&[user(text)]).text_tokens;
assert!(est >= *truth, "estimate {est} under-counts {truth} real tokens for {text:?}");
}
}
#[test]
fn an_explicit_max_tokens_bounds_the_output_but_never_exceeds_the_model_ceiling() {
let r = rates(Some(8_000), 200_000);
let p = CompletionParameters { max_tokens: Some(500), ..Default::default() };
assert_eq!(max_output_tokens(&p, &r), 500);
let p = CompletionParameters { max_tokens: Some(999_999), ..Default::default() };
assert_eq!(max_output_tokens(&p, &r), 8_000);
}
#[test]
fn the_default_max_tokens_is_what_gets_priced() {
let r = rates(Some(64_000), 200_000);
assert_eq!(max_output_tokens(&CompletionParameters::default(), &r), 4_096);
}
#[test]
fn with_no_max_tokens_and_no_completion_cap_the_context_window_is_the_ceiling() {
let r = rates(None, 32_768);
let p = CompletionParameters { max_tokens: None, ..Default::default() };
assert_eq!(max_output_tokens(&p, &r), 32_768);
let capped = rates(Some(8_000), 32_768);
assert_eq!(max_output_tokens(&p, &capped), 8_000);
}
#[test]
fn reasoning_tokens_are_added_on_top_of_the_visible_output_budget() {
let r = rates(Some(8_000), 200_000);
let base = CompletionParameters { max_tokens: Some(1_000), ..Default::default() };
assert_eq!(max_output_tokens(&base, &r), 1_000);
let thinking = CompletionParameters {
max_tokens: Some(1_000),
reasoning: Some(ReasoningConfig { effort: None, max_tokens: Some(4_000), exclude: None }),
..Default::default()
};
assert_eq!(max_output_tokens(&thinking, &r), 5_000);
let effort = CompletionParameters {
max_tokens: Some(1_000),
reasoning: Some(ReasoningConfig { effort: Some("high".into()), max_tokens: None, exclude: None }),
..Default::default()
};
assert_eq!(max_output_tokens(&effort, &r), 9_000);
let off = CompletionParameters {
max_tokens: Some(1_000),
reasoning: Some(ReasoningConfig { effort: Some("none".into()), max_tokens: None, exclude: None }),
..Default::default()
};
assert_eq!(max_output_tokens(&off, &r), 1_000);
}
fn video(duration_secs: Option<f64>) -> Message {
parts(vec![ContentPart::Video {
video_url: crate::message::VideoUrl { url: "https://x/y.mp4".into(), duration_secs },
}])
}
fn audio(duration_secs: Option<f64>) -> Message {
parts(vec![ContentPart::Audio {
input_audio: AudioInput {
data: "AAAA".into(),
format: Some("wav".into()),
duration_secs,
},
}])
}
#[test]
fn audio_of_known_length_is_counted_by_the_second() {
let one_second = estimate_prompt_tokens(&[audio(Some(1.0))]);
assert_eq!(one_second.audio_tokens, AUDIO_TOKENS_PER_SECOND as u64);
assert_eq!(one_second.image_tokens, 0);
let sliver = estimate_prompt_tokens(&[audio(Some(0.01))]);
assert_eq!(sliver.audio_tokens, 1);
}
#[test]
fn a_video_is_counted_as_frames_and_a_soundtrack() {
let ten = estimate_prompt_tokens(&[video(Some(10.0))]);
assert_eq!(ten.image_tokens, 10 * u64::from(TOKENS_PER_VIDEO_FRAME), "one frame a second");
assert_eq!(ten.audio_tokens, 10 * AUDIO_TOKENS_PER_SECOND as u64, "its sound too");
let sliver = estimate_prompt_tokens(&[video(Some(0.1))]);
assert_eq!(sliver.image_tokens, u64::from(TOKENS_PER_VIDEO_FRAME));
}
#[test]
fn the_counts_reproduce_geminis_published_video_limits() {
const WINDOW: u64 = 1_000_000;
let hour = estimate_prompt_tokens(&[video(Some(3600.0))]);
assert_eq!(hour.image_tokens, 3600 * 258, "one frame a second");
assert_eq!(hour.audio_tokens, 3600 * 32, "its soundtrack");
assert!(hour.image_tokens < WINDOW, "an hour of silent video fits");
assert!(hour.total_tokens() > WINDOW, "an hour WITH audio does not fit");
let three_quarters = estimate_prompt_tokens(&[video(Some(45.0 * 60.0))]);
assert!(three_quarters.total_tokens() < WINDOW, "45 min with audio fits");
}
#[test]
fn media_of_unknown_length_is_assumed_not_refused() {
let assumed = DEFAULT_MEDIA_SECONDS as u64;
let a = estimate_prompt_tokens(&[audio(None)]);
assert_eq!(a.audio_tokens, assumed * AUDIO_TOKENS_PER_SECOND as u64);
let v = estimate_prompt_tokens(&[video(None)]);
assert_eq!(v.image_tokens, assumed * u64::from(TOKENS_PER_VIDEO_FRAME));
assert_eq!(v.audio_tokens, assumed * AUDIO_TOKENS_PER_SECOND as u64);
}
#[test]
fn media_with_a_nonsense_duration_falls_back_to_the_assumption() {
let assumed = estimate_prompt_tokens(&[audio(None)]).audio_tokens;
for bad in [f64::INFINITY, f64::NAN, f64::NEG_INFINITY, -1.0] {
assert_eq!(estimate_prompt_tokens(&[audio(Some(bad))]).audio_tokens, assumed, "{bad}");
assert_eq!(estimate_prompt_tokens(&[video(Some(bad))]).audio_tokens, assumed, "{bad}");
}
}
#[test]
fn an_absurdly_long_clip_counts_saturated_rather_than_overflowing() {
for absurd in [1e30, f64::MAX] {
let est = estimate_prompt_tokens(&[video(Some(absurd))]);
assert!(est.image_tokens > 0 && est.audio_tokens > 0, "{absurd} still counts");
}
let many = vec![
ContentPart::Video {
video_url: crate::message::VideoUrl {
url: "x".into(),
duration_secs: Some(f64::MAX),
},
};
64
];
let est = estimate_prompt_tokens(&[parts(many)]);
assert_eq!(est.image_tokens, u64::MAX, "saturated, not wrapped");
assert!(est.total_tokens() == u64::MAX, "and the total does not panic");
}
#[test]
fn an_absurdly_long_clip_still_prices_as_one_full_window() {
let window = 1_000u32;
let r = ModelRates {
price: TokenPrice::new(1.0, 0.0),
max_completion_tokens: Some(0),
context_length: window,
};
let params = CompletionParameters { max_tokens: Some(0), ..Default::default() };
let cost = estimate_cost_usd(&[video(Some(f64::MAX))], ¶ms, &r);
let full_window_of_audio = f64::from(window) * r.price.audio_rate() / 1e6;
assert!((cost - full_window_of_audio).abs() < 1e-12, "{cost}");
}
#[test]
fn a_prompt_too_large_to_send_is_priced_as_the_full_window() {
let window = 1_000u32;
let r = ModelRates {
price: TokenPrice::new(1.0, 0.0),
max_completion_tokens: Some(0),
context_length: window,
};
let params = CompletionParameters { max_tokens: Some(0), ..Default::default() };
let huge = estimate_prompt_tokens(&[video(Some(86_400.0))]);
assert!(huge.total_tokens() > u64::from(window) * 1_000);
let cost = estimate_cost_usd(&[video(Some(86_400.0))], ¶ms, &r);
let full_window_of_audio = f64::from(window) * r.price.audio_rate() / 1e6;
assert!((cost - full_window_of_audio).abs() < 1e-12, "{cost} vs {full_window_of_audio}");
}
#[test]
fn an_over_long_prompt_fills_the_window_with_its_dearest_tokens() {
let window = 100u32;
let price = TokenPrice::new(1.0, 0.0).with_media_rates(Some(1000.0), None);
let r = ModelRates { price, max_completion_tokens: Some(0), context_length: window };
let params = CompletionParameters { max_tokens: Some(0), ..Default::default() };
let long_text = "word ".repeat(500);
let messages = [audio(Some(10.0)), user(&long_text)];
let cost = estimate_cost_usd(&messages, ¶ms, &r);
let window_of_audio = f64::from(window) * 1000.0 / 1e6;
assert!((cost - window_of_audio).abs() < 1e-12, "audio must crowd out the cheap text");
let window_of_text = f64::from(window) * 1.0 / 1e6;
assert!(cost > window_of_text * 500.0);
}
#[test]
fn audio_is_billed_at_the_audio_rate_not_the_input_rate() {
let r = ModelRates {
price: TokenPrice::new(0.1, 0.2).with_media_rates(Some(100.0), None),
max_completion_tokens: Some(0),
context_length: 1_000,
};
let params = CompletionParameters { max_tokens: Some(0), ..Default::default() };
let cost = estimate_cost_usd(&[audio(Some(1.0))], ¶ms, &r);
let audio_only = 32.0 * 100.0 / 1e6;
assert!(cost > audio_only, "the envelope's text tokens count too: {cost}");
assert!(cost < audio_only * 1.01, "but audio dominates: {cost} vs {audio_only}");
let as_input = 32.0 * 0.1 / 1e6;
assert!(cost > as_input * 500.0, "audio must not be billed at the input rate");
}
#[test]
fn an_unpublished_media_rate_falls_back_to_a_premium_over_text() {
let plain = TokenPrice::new(3.0, 15.0);
assert_eq!(plain.image_rate(), 3.0, "image is billed exactly as text");
assert_eq!(plain.audio_rate(), 3.0 * AUDIO_RATE_FALLBACK_MULTIPLE);
assert!(plain.audio_rate() > plain.input_per_mtok, "audio always costs more");
let priced = plain.clone().with_media_rates(Some(30.0), Some(4.0));
assert_eq!(priced.audio_rate(), 30.0);
assert_eq!(priced.image_rate(), 4.0);
}
#[test]
fn a_still_image_is_charged_a_bounded_tile_cost_rather_than_nothing() {
let with_image = parts(vec![ContentPart::Image {
image_url: crate::message::ImageUrl { url: "https://x/y.png".into(), detail: None },
}]);
let est = estimate_prompt_tokens(&[with_image]);
assert_eq!(est.image_tokens, u64::from(TOKENS_PER_STILL_IMAGE), "a still is not free");
assert_eq!(est.audio_tokens, 0, "a still image has no soundtrack");
}
#[test]
fn cost_prices_the_bounded_output_at_the_output_rate() {
let r = ModelRates {
price: TokenPrice::new(3.0, 15.0),
max_completion_tokens: Some(1_000_000),
context_length: 2_000_000,
};
let p = CompletionParameters { max_tokens: Some(1_000_000), ..Default::default() };
let cost = estimate_cost_usd(&[user("hi")], &p, &r);
assert!((cost - 15.0).abs() < 0.001, "{cost}");
}
}