use crate::generator::CompletionParameters;
use crate::generator::ModelRates;
use crate::message::{ContentPart, Message, MessageContent};
use crate::provider::bpe;
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_IMAGE_BASE: u32 = 85;
const TOKENS_PER_IMAGE_TILE: u32 = 170;
fn still_image_tokens(width: Option<u32>, height: Option<u32>) -> u32 {
let (Some(w), Some(h)) = (width, height) else {
return TOKENS_PER_STILL_IMAGE;
};
if w == 0 || h == 0 {
return TOKENS_PER_STILL_IMAGE;
}
let (mut w, mut h) = (f64::from(w), f64::from(h));
let fit = (2048.0 / w).min(2048.0 / h).min(1.0);
w *= fit;
h *= fit;
let shrink = (768.0 / w.min(h)).min(1.0);
w *= shrink;
h *= shrink;
let tiles = (w / 512.0).ceil() * (h / 512.0).ceil();
TOKENS_PER_IMAGE_BASE + TOKENS_PER_IMAGE_TILE * tiles as u32
}
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_url } => {
image_tokens +=
u64::from(still_image_tokens(image_url.width, image_url.height))
}
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 a_declared_resolution_refines_a_stills_price_under_the_flat_bound() {
assert_eq!(
still_image_tokens(Some(400), Some(300)),
TOKENS_PER_IMAGE_BASE + TOKENS_PER_IMAGE_TILE
);
let huge = still_image_tokens(Some(8000), Some(6000));
assert!(huge < TOKENS_PER_STILL_IMAGE, "{huge}");
assert!(huge > still_image_tokens(Some(400), Some(300)));
assert_eq!(still_image_tokens(None, None), TOKENS_PER_STILL_IMAGE);
assert_eq!(
still_image_tokens(Some(0), Some(10)),
TOKENS_PER_STILL_IMAGE
);
}
#[test]
fn a_metadata_keeping_wire_round_trip_preserves_the_estimate() {
use crate::message::{messages_to_payload, ImageData, VideoData};
let originals = vec![parts(vec![
ContentPart::text("describe all of this"),
ContentPart::audio(
&crate::message::AudioData::from_bytes(&[0u8; 4], "mp3").with_duration(90.0),
),
ContentPart::video(&VideoData::from_url("https://x/y.mp4").with_duration(30.0)),
ContentPart::image(&ImageData::from_url("https://x/y.png").with_dimensions(800, 600)),
])];
let wire = serde_json::Value::Array(messages_to_payload(&originals, true));
let round_tripped: Vec<Message> = serde_json::from_value(wire).expect("wire parses back");
let r = rates(Some(1000), 200_000);
let params = CompletionParameters::new().with_max_tokens(500);
let direct = estimate_cost_usd(&originals, ¶ms, &r);
let re_parsed = estimate_cost_usd(&round_tripped, ¶ms, &r);
assert!(direct > 0.0);
assert_eq!(
direct, re_parsed,
"the wire bytes carry everything the estimator uses"
);
let strict = serde_json::Value::Array(messages_to_payload(&originals, false));
let strict_msgs: Vec<Message> = serde_json::from_value(strict).expect("parses");
assert_ne!(estimate_cost_usd(&strict_msgs, ¶ms, &r), direct);
}
#[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,
width: None,
height: None,
},
}])
}
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),
width: None,
height: None,
},
};
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,
width: None,
height: 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}");
}
}