use crate::config::models::gateway::GatewayPricingConfig;
use crate::core::budget::{UnifiedBudgetLimits, UnifiedBudgetReservation};
use crate::core::models::openai::requests::ChatCompletionRequest;
use crate::core::models::openai::{
ChatMessage, ContentPart, Function, FunctionCall, MessageContent, ResponseFormat, Tool,
};
#[cfg(test)]
use crate::core::pricing_service::PricingService;
use crate::core::providers::unified_provider::ProviderError;
use crate::utils::ai::counter::token_counter::{TokenCounter, TokenizerIdentity};
use crate::utils::error::gateway_error::{GatewayError, Result as GatewayResult};
const IMAGE_PROMPT_BASE_TOKENS: u32 = 85;
pub(in crate::server::routes::ai) const IMAGE_HIGH_DETAIL_PROMPT_TOKENS: u32 = 1_105;
const AUDIO_PROMPT_BASE_TOKENS: u32 = 100;
const DOCUMENT_PROMPT_BASE_TOKENS: u32 = 1_000;
const TOOL_RESULT_BASE_TOKENS: u32 = 50;
const TOOL_USE_BASE_TOKENS: u32 = 100;
#[derive(Clone, Copy)]
pub(in crate::server::routes::ai) struct ChatCompletionBudgetRequest<'a> {
messages: &'a [ChatMessage],
tools: Option<&'a [Tool]>,
functions: Option<&'a [Function]>,
function_call: Option<&'a FunctionCall>,
response_format: Option<&'a ResponseFormat>,
max_tokens: Option<u32>,
max_completion_tokens: Option<u32>,
n: Option<u32>,
retained_prompt_tokens: u32,
}
impl<'a> ChatCompletionBudgetRequest<'a> {
pub(in crate::server::routes::ai) fn with_retained_prompt_tokens(
mut self,
tokens: u32,
) -> Self {
self.retained_prompt_tokens = tokens;
self
}
pub(in crate::server::routes::ai) fn with_output_limits(
mut self,
max_tokens: Option<u32>,
max_completion_tokens: Option<u32>,
) -> Self {
self.max_tokens = max_tokens;
self.max_completion_tokens = max_completion_tokens;
self
}
}
impl<'a> From<&'a ChatCompletionRequest> for ChatCompletionBudgetRequest<'a> {
fn from(request: &'a ChatCompletionRequest) -> Self {
Self {
messages: &request.messages,
tools: request.tools.as_deref(),
functions: request.functions.as_deref(),
function_call: request.function_call.as_ref(),
response_format: request.response_format.as_ref(),
max_tokens: request.max_tokens,
max_completion_tokens: request.max_completion_tokens,
n: request.n,
retained_prompt_tokens: 0,
}
}
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn reserve_completion_budget(
budget_limits: &UnifiedBudgetLimits,
provider: &str,
model: &str,
estimated_prompt_tokens: u32,
max_output_tokens: Option<u32>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
reserve_completion_budget_with_pricing(
super::default_spend_pricing_service(),
budget_limits,
provider,
model,
estimated_prompt_tokens,
max_output_tokens,
)
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn reserve_completion_budget_with_pricing(
pricing_service: &PricingService,
budget_limits: &UnifiedBudgetLimits,
provider: &str,
model: &str,
estimated_prompt_tokens: u32,
max_output_tokens: Option<u32>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
reserve_completion_budget_with_policy(
pricing_service,
&GatewayPricingConfig::default(),
budget_limits,
provider,
model,
estimated_prompt_tokens,
max_output_tokens,
)
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn reserve_completion_budget_with_policy(
pricing_service: &PricingService,
pricing_config: &GatewayPricingConfig,
budget_limits: &UnifiedBudgetLimits,
provider: &str,
model: &str,
estimated_prompt_tokens: u32,
max_output_tokens: Option<u32>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
reserve_completion_budget_with_split_pricing(
pricing_service,
pricing_config,
budget_limits,
provider,
model,
provider,
model,
estimated_prompt_tokens,
max_output_tokens,
)
}
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
pub(in crate::server::routes::ai) fn reserve_completion_budget_with_split_pricing(
pricing_service: &PricingService,
pricing_config: &GatewayPricingConfig,
budget_limits: &UnifiedBudgetLimits,
budget_provider: &str,
budget_model: &str,
pricing_provider: &str,
pricing_model: &str,
estimated_prompt_tokens: u32,
max_output_tokens: Option<u32>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
let estimate = match pricing_service.estimate_loaded_completion_cost_for_provider(
pricing_provider,
pricing_model,
estimated_prompt_tokens,
max_output_tokens,
) {
Ok(estimate) => estimate,
Err(error) => {
tracing::error!(
"cost estimation failed for pricing provider '{pricing_provider}' model \
'{pricing_model}' budget provider '{budget_provider}' model '{budget_model}': {error}; \
applying unpriced model policy"
);
return super::unpriced::reserve_unpriced_completion_budget(
pricing_config,
budget_limits,
budget_provider,
budget_model,
estimated_prompt_tokens,
max_output_tokens,
error,
);
}
};
if estimate.max_cost <= 0.0 {
super::ensure_budget_available(budget_limits, budget_provider, budget_model)?;
return Ok(None);
}
budget_limits
.reserve_spend(budget_provider, budget_model, estimate.max_cost)
.map(Some)
.map_err(|error| {
super::reservation_error_to_provider_error(error, budget_provider, budget_model)
})
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn reserve_chat_completion_budget(
budget_limits: &UnifiedBudgetLimits,
provider: &str,
model: &str,
request: &ChatCompletionRequest,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
reserve_chat_completion_budget_with_pricing(
super::default_spend_pricing_service(),
budget_limits,
provider,
model,
request,
)
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn reserve_chat_completion_budget_with_pricing(
pricing_service: &PricingService,
budget_limits: &UnifiedBudgetLimits,
provider: &str,
model: &str,
request: &ChatCompletionRequest,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
reserve_chat_completion_budget_with_policy(
pricing_service,
&GatewayPricingConfig::default(),
budget_limits,
provider,
model,
request,
)
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn reserve_chat_completion_budget_with_policy(
pricing_service: &PricingService,
pricing_config: &GatewayPricingConfig,
budget_limits: &UnifiedBudgetLimits,
provider: &str,
model: &str,
request: &ChatCompletionRequest,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
reserve_chat_completion_budget_with_split_pricing(
pricing_service,
pricing_config,
budget_limits,
provider,
model,
provider,
model,
request,
)
}
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
pub(in crate::server::routes::ai) fn reserve_chat_completion_budget_with_split_pricing<'a>(
pricing_service: &PricingService,
pricing_config: &GatewayPricingConfig,
budget_limits: &UnifiedBudgetLimits,
budget_provider: &str,
budget_model: &str,
pricing_provider: &str,
pricing_model: &str,
request: impl Into<ChatCompletionBudgetRequest<'a>>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
let request = request.into();
let token_identity = if pricing_provider == "openai" {
TokenizerIdentity::exact_openai(pricing_model)
} else {
TokenizerIdentity::approximate(pricing_provider, pricing_model)
};
let prompt_tokens = estimate_chat_prompt_tokens(
&token_identity,
request.messages,
request.tools,
request.functions,
request.function_call,
request.response_format,
);
reserve_completion_budget_with_split_pricing(
pricing_service,
pricing_config,
budget_limits,
budget_provider,
budget_model,
pricing_provider,
pricing_model,
prompt_tokens,
reservation_output_tokens(
pricing_service,
pricing_provider,
pricing_model,
prompt_tokens,
provider_effective_max_output_tokens_for_budget(budget_provider, budget_model, request),
request.n.unwrap_or(1),
),
)
}
pub(in crate::server::routes::ai) fn reserve_chat_completion_budget_with_request_pricing<'a>(
request_pricing: &super::pricing::RequestPricing,
pricing_config: &GatewayPricingConfig,
budget_limits: &UnifiedBudgetLimits,
budget_provider: &str,
budget_model: &str,
request: impl Into<ChatCompletionBudgetRequest<'a>>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
let request = request.into();
let prompt_tokens = try_estimate_chat_prompt_tokens(
request_pricing.token_identity(),
request.messages,
request.tools,
request.functions,
request.function_call,
request.response_format,
)
.map_err(|error| {
super::token_count_error(
budget_provider,
budget_model,
request_pricing.token_identity(),
error,
)
})?;
let prompt_tokens = prompt_tokens
.checked_add(request.retained_prompt_tokens)
.ok_or_else(|| {
ProviderError::invalid_request(
"responses",
"Retained context exceeds the token accounting range",
)
})?;
let max_output_tokens = request
.max_completion_tokens
.or(request.max_tokens)
.map(|tokens| tokens.saturating_mul(request.n.unwrap_or(1)));
reserve_completion_budget_with_counted_input(
request_pricing,
pricing_config,
budget_limits,
budget_provider,
budget_model,
prompt_tokens,
max_output_tokens,
)
}
#[allow(clippy::too_many_arguments)]
pub(in crate::server::routes::ai) fn reserve_completion_budget_with_counted_input(
request_pricing: &super::pricing::RequestPricing,
pricing_config: &GatewayPricingConfig,
budget_limits: &UnifiedBudgetLimits,
budget_provider: &str,
budget_model: &str,
prompt_tokens: u32,
max_output_tokens: Option<u32>,
) -> Result<Option<UnifiedBudgetReservation>, ProviderError> {
let estimate = match request_pricing.estimate_completion(prompt_tokens, max_output_tokens) {
Ok(estimate) => estimate,
Err(error) => {
return super::unpriced::reserve_unpriced_completion_budget(
pricing_config,
budget_limits,
budget_provider,
budget_model,
prompt_tokens,
max_output_tokens,
error,
);
}
};
if estimate.max_cost <= 0.0 {
super::ensure_budget_available(budget_limits, budget_provider, budget_model)?;
return Ok(None);
}
budget_limits
.reserve_spend(budget_provider, budget_model, estimate.max_cost)
.map(Some)
.map_err(|error| {
super::reservation_error_to_provider_error(error, budget_provider, budget_model)
})
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn estimate_chat_prompt_tokens(
identity: &TokenizerIdentity,
messages: &[ChatMessage],
tools: Option<&[Tool]>,
functions: Option<&[Function]>,
function_call: Option<&FunctionCall>,
response_format: Option<&ResponseFormat>,
) -> u32 {
let counter = TokenCounter::new();
let message_tokens = counter.count_chat_tokens(identity, messages).map_or_else(
|_| fallback_message_tokens(messages),
|value| value.input_tokens,
);
let multimodal_tokens = conservative_multimodal_prompt_extra(messages);
let tool_tokens = serialized_prompt_tokens_auto(&counter, identity, tools, |value| {
value.len().saturating_mul(256)
});
let function_tokens = serialized_prompt_tokens_auto(&counter, identity, functions, |value| {
value.len().saturating_mul(256)
});
let function_call_tokens =
serialized_prompt_tokens_auto(&counter, identity, function_call, |_| 64);
let response_format_tokens =
serialized_prompt_tokens_auto(&counter, identity, response_format, |_| 128);
message_tokens
.saturating_add(multimodal_tokens)
.saturating_add(tool_tokens)
.saturating_add(function_tokens)
.saturating_add(function_call_tokens)
.saturating_add(response_format_tokens)
}
pub(in crate::server::routes::ai) fn try_estimate_chat_prompt_tokens(
identity: &TokenizerIdentity,
messages: &[ChatMessage],
tools: Option<&[Tool]>,
functions: Option<&[Function]>,
function_call: Option<&FunctionCall>,
response_format: Option<&ResponseFormat>,
) -> GatewayResult<u32> {
let counter = TokenCounter::new();
let message_estimate = counter.count_chat_tokens(identity, messages)?;
observe_approximate_estimate(identity, "messages", &message_estimate);
let message_tokens = message_estimate.input_tokens;
let multimodal_tokens = conservative_multimodal_prompt_extra(messages);
let tool_tokens = serialized_prompt_tokens(&counter, identity, tools, "tools")?;
let function_tokens = serialized_prompt_tokens(&counter, identity, functions, "functions")?;
let function_call_tokens =
serialized_prompt_tokens(&counter, identity, function_call, "function_call")?;
let response_format_tokens =
serialized_prompt_tokens(&counter, identity, response_format, "response_format")?;
Ok(message_tokens
.saturating_add(multimodal_tokens)
.saturating_add(tool_tokens)
.saturating_add(function_tokens)
.saturating_add(function_call_tokens)
.saturating_add(response_format_tokens))
}
#[cfg(test)]
fn reservation_output_tokens(
pricing_service: &PricingService,
provider: &str,
model: &str,
prompt_tokens: u32,
requested_max_output_tokens: Option<u32>,
choice_count: u32,
) -> Option<u32> {
let counter = TokenCounter::new();
let identity = if provider == "openai" {
TokenizerIdentity::exact_openai(model)
} else {
TokenizerIdentity::approximate(provider, model)
};
let choice_count = choice_count.max(1);
let output_tokens = if let Some(requested) = requested_max_output_tokens {
Some(requested)
} else {
catalog_max_output_tokens_with_pricing(pricing_service, provider, model).or_else(|| {
counter
.estimate_output_tokens(None, prompt_tokens, &identity)
.ok()
})
};
output_tokens.map(|tokens| tokens.saturating_mul(choice_count))
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn catalog_max_output_tokens(
provider: &str,
model: &str,
) -> Option<u32> {
catalog_max_output_tokens_with_pricing(super::default_spend_pricing_service(), provider, model)
}
#[cfg(test)]
fn catalog_max_output_tokens_with_pricing(
pricing_service: &PricingService,
provider: &str,
model: &str,
) -> Option<u32> {
pricing_service.max_output_tokens_for_provider(provider, model)
}
#[cfg(test)]
pub(in crate::server::routes::ai) fn provider_effective_max_output_tokens(
provider: &str,
model: &str,
request: &ChatCompletionRequest,
) -> Option<u32> {
provider_effective_max_output_tokens_for_budget(provider, model, request.into())
}
#[cfg(test)]
fn provider_effective_max_output_tokens_for_budget(
provider: &str,
model: &str,
request: ChatCompletionBudgetRequest<'_>,
) -> Option<u32> {
let provider = crate::core::pricing::normalize_pricing_provider(provider);
match provider.as_str() {
"openai" | "azure" | "azure_ai" | "openai_like" | "openrouter" | "xai" | "groq"
| "deepseek" | "moonshot" | "minimax" | "zhipuai" | "xiaomi_mimo" | "amazon_nova"
| "baseten" | "huggingface" | "zai" | "together_ai" | "fireworks_ai" | "aiml" => {
request.max_completion_tokens.or(request.max_tokens)
}
"anthropic" => Some(request.max_tokens.unwrap_or(4096)),
"bedrock" => bedrock_effective_max_output_tokens(model, request),
"cohere" | "replicate" => request.max_tokens.or(request.max_completion_tokens),
_ => request.max_tokens,
}
}
#[cfg(test)]
fn bedrock_effective_max_output_tokens(
model: &str,
request: ChatCompletionBudgetRequest<'_>,
) -> Option<u32> {
use crate::core::providers::bedrock::BedrockApiType;
let Ok(config) = crate::core::providers::bedrock::get_model_config_for_model_id(model) else {
return request.max_tokens;
};
match config.api_type {
BedrockApiType::Converse | BedrockApiType::ConverseStream => {
request.max_completion_tokens.or(request.max_tokens)
}
BedrockApiType::Invoke | BedrockApiType::InvokeStream => request.max_tokens,
}
}
fn conservative_multimodal_prompt_extra(messages: &[ChatMessage]) -> u32 {
messages
.iter()
.filter_map(|message| message.content.as_ref())
.flat_map(|content| match content {
MessageContent::Text(_) => [].as_slice(),
MessageContent::Parts(parts) => parts.as_slice(),
})
.fold(0u32, |total, part| {
total.saturating_add(conservative_content_part_extra(part))
})
}
fn conservative_content_part_extra(part: &ContentPart) -> u32 {
match part {
ContentPart::ImageUrl { image_url } => {
image_prompt_floor(image_url.detail.as_deref()).saturating_sub(IMAGE_PROMPT_BASE_TOKENS)
}
ContentPart::Image {
source,
detail,
image_url,
} => {
let detail = detail
.as_deref()
.or_else(|| image_url.as_ref().and_then(|url| url.detail.as_deref()));
image_prompt_floor(detail)
.max(encoded_media_tokens(&source.data))
.saturating_sub(IMAGE_PROMPT_BASE_TOKENS)
}
ContentPart::Audio { audio } => encoded_media_tokens(&audio.data)
.max(AUDIO_PROMPT_BASE_TOKENS)
.saturating_sub(AUDIO_PROMPT_BASE_TOKENS),
ContentPart::Document { source, .. } => encoded_media_tokens(&source.data)
.max(DOCUMENT_PROMPT_BASE_TOKENS)
.saturating_sub(DOCUMENT_PROMPT_BASE_TOKENS),
ContentPart::ToolResult { .. } => {
serialized_content_part_tokens(part).saturating_sub(TOOL_RESULT_BASE_TOKENS)
}
ContentPart::ToolUse { .. } => {
serialized_content_part_tokens(part).saturating_sub(TOOL_USE_BASE_TOKENS)
}
ContentPart::Text { .. } => 0,
}
}
fn serialized_content_part_tokens(part: &ContentPart) -> u32 {
let Ok(json) = serde_json::to_string(part) else {
return u32::MAX;
};
u32::try_from(json.chars().count().div_ceil(4)).unwrap_or(u32::MAX)
}
fn image_prompt_floor(detail: Option<&str>) -> u32 {
if detail.is_some_and(|detail| detail.eq_ignore_ascii_case("low")) {
IMAGE_PROMPT_BASE_TOKENS
} else {
IMAGE_HIGH_DETAIL_PROMPT_TOKENS
}
}
fn encoded_media_tokens(data: &str) -> u32 {
u32::try_from(data.chars().count().div_ceil(4)).unwrap_or(u32::MAX)
}
#[cfg(test)]
fn serialized_prompt_tokens_auto<T, F>(
counter: &TokenCounter,
identity: &TokenizerIdentity,
value: Option<&T>,
fallback_units: F,
) -> u32
where
T: serde::Serialize + ?Sized,
F: FnOnce(&T) -> usize,
{
let Some(value) = value else {
return 0;
};
let Ok(json) = serde_json::to_string(value) else {
return u32::try_from(fallback_units(value)).unwrap_or(u32::MAX);
};
counter
.count_completion_tokens(identity, &json)
.map_or_else(
|_| u32::try_from(json.chars().count().div_ceil(4)).unwrap_or(u32::MAX),
|estimate| estimate.input_tokens,
)
}
#[cfg(test)]
fn fallback_message_tokens(messages: &[ChatMessage]) -> u32 {
let chars = messages
.iter()
.filter_map(|message| message.content.as_ref())
.map(|content| match content {
MessageContent::Text(text) => text.chars().count(),
MessageContent::Parts(parts) => serde_json::to_string(parts)
.map(|text| text.chars().count())
.unwrap_or_default(),
})
.sum::<usize>();
let overhead = messages.len().saturating_mul(4).saturating_add(8);
u32::try_from(chars.div_ceil(4).saturating_add(overhead)).unwrap_or(u32::MAX)
}
fn serialized_prompt_tokens<T>(
counter: &TokenCounter,
identity: &TokenizerIdentity,
value: Option<&T>,
input: &str,
) -> GatewayResult<u32>
where
T: serde::Serialize + ?Sized,
{
let Some(value) = value else {
return Ok(0);
};
let json = serde_json::to_string(value)
.map_err(|error| GatewayError::Config(format!("failed to serialize {input}: {error}")))?;
let estimate = counter.count_completion_tokens(identity, &json)?;
observe_approximate_estimate(identity, input, &estimate);
Ok(estimate.input_tokens)
}
fn observe_approximate_estimate(
identity: &TokenizerIdentity,
input: &str,
estimate: &crate::utils::ai::counter::types::TokenEstimate,
) {
if estimate.is_approximate {
tracing::warn!(
token_provider = identity.provider(),
token_model = identity.model(),
input,
is_approximate = true,
confidence = estimate.confidence,
"explicit approximate token count used"
);
}
}
#[cfg(test)]
mod budget_request_tests {
use super::*;
#[test]
fn selected_exact_openai_identity_fails_when_tokenizer_is_unavailable() {
let pricing = PricingService::new(None);
pricing.add_custom_model(
"gpt-audio-1.5".to_string(),
crate::core::pricing_service::LiteLLMModelInfo {
max_tokens: Some(4_096),
max_input_tokens: Some(4_096),
max_output_tokens: Some(1_024),
input_cost_per_token: Some(0.01),
output_cost_per_token: Some(0.02),
input_cost_per_character: None,
output_cost_per_character: None,
cost_per_second: None,
litellm_provider: "openai".to_string(),
mode: "chat".to_string(),
supports_function_calling: None,
supports_vision: None,
supports_streaming: None,
supports_parallel_function_calling: None,
supports_system_message: None,
extra: std::collections::HashMap::new(),
},
);
let request_pricing =
super::super::pricing::RequestPricing::from_exact(&pricing, "openai", "gpt-audio-1.5");
assert_eq!(
request_pricing.priced_parts(),
Some(("openai", "gpt-audio-1.5"))
);
let messages = vec![ChatMessage {
role: crate::core::models::openai::MessageRole::User,
content: Some(MessageContent::Text("hello".to_string())),
name: None,
function_call: None,
tool_calls: None,
tool_call_id: None,
audio: None,
}];
let request = ChatCompletionBudgetRequest {
messages: &messages,
tools: None,
functions: None,
function_call: None,
response_format: None,
max_tokens: Some(16),
max_completion_tokens: None,
n: None,
retained_prompt_tokens: 0,
};
let error = match reserve_chat_completion_budget_with_request_pricing(
&request_pricing,
&GatewayPricingConfig::default(),
&UnifiedBudgetLimits::new(),
"selected-openai",
"wire-audio-deployment",
request,
) {
Err(error) => error,
Ok(_) => {
panic!("missing exact tokenizer must fail before pricing or upstream execution")
}
};
assert!(matches!(
error,
ProviderError::InvalidRequest {
provider: "token_count",
..
}
));
let message = error.to_string();
assert!(message.contains("selected-openai/wire-audio-deployment"));
assert!(message.contains("openai/gpt-audio-1.5"));
}
#[test]
fn budget_request_view_borrows_large_request_parts() {
let request = ChatCompletionRequest {
messages: vec![ChatMessage {
role: crate::core::models::openai::MessageRole::User,
content: Some(MessageContent::Text("x".repeat(64 * 1024))),
name: None,
function_call: None,
tool_calls: None,
tool_call_id: None,
audio: None,
}],
tools: Some(vec![Tool {
tool_type: "function".to_string(),
function: crate::core::models::openai::Function {
name: "lookup".to_string(),
description: None,
parameters: Some(serde_json::json!({"type": "object"})),
},
}]),
max_tokens: Some(1024),
..Default::default()
};
let view =
ChatCompletionBudgetRequest::from(&request).with_output_limits(Some(128), Some(64));
assert!(std::ptr::eq(
view.messages.as_ptr(),
request.messages.as_ptr()
));
assert!(std::ptr::eq(
view.tools.expect("tools").as_ptr(),
request.tools.as_ref().expect("tools").as_ptr()
));
assert_eq!(view.max_tokens, Some(128));
assert_eq!(view.max_completion_tokens, Some(64));
}
}