Merge pull request #20541 from ryan-crabbe/perf/cost-calculator-optimizations

Perf/cost calculator optimizations
This commit is contained in:
ryan-crabbe 2026-02-21 12:18:38 -08:00 committed by GitHub
commit 469951466e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 107 additions and 43 deletions

View file

@ -198,9 +198,25 @@ def _get_token_base_cost(
cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key))
## CHECK IF ABOVE THRESHOLD
# Optimization: collect threshold keys first to avoid sorting all model_info keys.
# Most models don't have threshold pricing, so we can return early.
threshold_keys = [
k for k in model_info if k.startswith("input_cost_per_token_above_")
]
if not threshold_keys:
return (
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
)
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
threshold: Optional[float] = None
for key, value in sorted(model_info.items(), reverse=True):
if key.startswith("input_cost_per_token_above_") and value is not None:
for key in sorted(threshold_keys, reverse=True):
value = model_info.get(key)
if value is not None:
try:
# Handle both formats: _above_128k_tokens and _above_128_tokens
threshold_str = key.split("_above_")[1].split("_tokens")[0]
@ -511,47 +527,52 @@ def _calculate_input_cost(
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost
### AUDIO COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_audio_token", prompt_tokens_details["audio_tokens"]
)
if prompt_tokens_details["audio_tokens"]:
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_audio_token", prompt_tokens_details["audio_tokens"]
)
### IMAGE TOKEN COST
# For image token costs:
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
image_token_cost_key = "input_cost_per_image_token"
if model_info.get(image_token_cost_key) is None:
image_token_cost_key = "input_cost_per_token"
prompt_cost += calculate_cost_component(
model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]
)
if prompt_tokens_details["image_tokens"]:
# For image token costs:
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
image_token_cost_key = "input_cost_per_image_token"
if model_info.get(image_token_cost_key) is None:
image_token_cost_key = "input_cost_per_token"
prompt_cost += calculate_cost_component(
model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]
)
### CACHE WRITING COST - Now uses tiered pricing
prompt_cost += calculate_cache_writing_cost(
cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"],
cache_creation_token_details=prompt_tokens_details[
"cache_creation_token_details"
],
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
cache_creation_cost=cache_creation_cost,
)
if prompt_tokens_details["cache_creation_tokens"] or prompt_tokens_details["cache_creation_token_details"] is not None:
prompt_cost += calculate_cache_writing_cost(
cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"],
cache_creation_token_details=prompt_tokens_details[
"cache_creation_token_details"
],
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
cache_creation_cost=cache_creation_cost,
)
### CHARACTER COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_character", prompt_tokens_details["character_count"]
)
if prompt_tokens_details["character_count"]:
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_character", prompt_tokens_details["character_count"]
)
### IMAGE COUNT COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_image", prompt_tokens_details["image_count"]
)
if prompt_tokens_details["image_count"]:
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_image", prompt_tokens_details["image_count"]
)
### VIDEO LENGTH COST
prompt_cost += calculate_cost_component(
model_info,
"input_cost_per_video_per_second",
prompt_tokens_details["video_length_seconds"],
)
if prompt_tokens_details["video_length_seconds"]:
prompt_cost += calculate_cost_component(
model_info,
"input_cost_per_video_per_second",
prompt_tokens_details["video_length_seconds"],
)
return prompt_cost
@ -676,18 +697,11 @@ def generic_cost_per_token( # noqa: PLR0915
## TEXT COST
completion_cost = float(text_tokens) * completion_base_cost
_output_cost_per_audio_token = _get_cost_per_unit(
model_info, "output_cost_per_audio_token", None
)
_output_cost_per_reasoning_token = _get_cost_per_unit(
model_info, "output_cost_per_reasoning_token", None
)
_output_cost_per_image_token = _get_cost_per_unit(
model_info, "output_cost_per_image_token", None
)
## AUDIO COST
if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0:
_output_cost_per_audio_token = _get_cost_per_unit(
model_info, "output_cost_per_audio_token", None
)
_output_cost_per_audio_token = (
_output_cost_per_audio_token
if _output_cost_per_audio_token is not None
@ -697,6 +711,9 @@ def generic_cost_per_token( # noqa: PLR0915
## REASONING COST
if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0:
_output_cost_per_reasoning_token = _get_cost_per_unit(
model_info, "output_cost_per_reasoning_token", None
)
_output_cost_per_reasoning_token = (
_output_cost_per_reasoning_token
if _output_cost_per_reasoning_token is not None
@ -706,6 +723,9 @@ def generic_cost_per_token( # noqa: PLR0915
## IMAGE COST
if not is_text_tokens_total and image_tokens and image_tokens > 0:
_output_cost_per_image_token = _get_cost_per_unit(
model_info, "output_cost_per_image_token", None
)
_output_cost_per_image_token = (
_output_cost_per_image_token
if _output_cost_per_image_token is not None

View file

@ -23,6 +23,8 @@ sys.path.insert(
) # Adds the parent directory to the system path
from litellm.litellm_core_utils.llm_cost_calc.utils import (
_calculate_input_cost,
PromptTokensDetailsResult,
calculate_cache_writing_cost,
generic_cost_per_token,
)
@ -559,6 +561,48 @@ def test_calculate_cache_writing_cost():
assert result_zero == 0.0
def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details():
"""
Regression test: when cache_creation_tokens is 0 but cache_creation_token_details
has non-zero ephemeral tokens, the cost must still be calculated.
This ensures the guard in _calculate_input_cost doesn't skip
calculate_cache_writing_cost when only ephemeral token details are present.
"""
cache_creation_cost = 3.75e-06
cache_creation_cost_above_1hr = 6e-06
prompt_tokens_details: PromptTokensDetailsResult = {
"cache_hit_tokens": 0,
"cache_creation_tokens": 0,
"cache_creation_token_details": CacheCreationTokenDetails(
ephemeral_5m_input_tokens=100,
ephemeral_1h_input_tokens=200,
),
"text_tokens": 0,
"audio_tokens": 0,
"image_tokens": 0,
"character_count": 0,
"image_count": 0,
"video_length_seconds": 0.0,
}
model_info: ModelInfo = {}
result = _calculate_input_cost(
prompt_tokens_details=prompt_tokens_details,
model_info=model_info,
prompt_base_cost=0.0,
cache_read_cost=0.0,
cache_creation_cost=cache_creation_cost,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
)
# Expected: (100 * 3.75e-06) + (200 * 6e-06) = 0.000375 + 0.0012 = 0.001575
expected = (100 * cache_creation_cost) + (200 * cache_creation_cost_above_1hr)
assert result > 0, "Cost should not be zero when ephemeral token details are present"
assert round(result, 6) == round(expected, 6)
def test_service_tier_flex_pricing():
"""Test that flex service tier uses correct pricing (approximately 50% of standard)."""
# Set up environment for local model cost map