From ee51bd07e54ece754e7773b408d9da70ad9ec46e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:40:36 -0700 Subject: [PATCH] fix(cost): preserve rates across mixed token thresholds --- .../litellm_core_utils/llm_cost_calc/utils.py | 4 +-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 30 +++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e7987e6221d..0b56f1032df 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -696,7 +696,7 @@ def _get_token_base_cost( # Only sort the threshold keys (typically 1-2 keys instead of 66+) threshold: float | None = None - for key in sorted(threshold_keys, key=_parse_above_token_threshold, reverse=True): + for key in sorted(threshold_keys, key=_parse_above_token_threshold): value = model_info.get(key) if value is not None: try: @@ -770,8 +770,6 @@ def _get_token_base_cost( ) cache_read_cost = get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) - - break except (IndexError, ValueError): continue except Exception: diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3a3daf85a16..ac20bbfcafc 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2550,6 +2550,36 @@ def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None: assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003) +@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001)) +@pytest.mark.parametrize("higher_rate", (0.0, 0.009)) +def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, higher_rate: float) -> None: + from typing import Final, cast + from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost + + info: Final = cast( + ModelInfo, + { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.003, + "cache_creation_input_token_cost": 0.004, + "input_cost_per_token_above_200k_tokens": 0.005, + "output_cost_per_token_above_200k_tokens": 0.006, + "cache_creation_input_token_cost_above_200k_tokens": 0.007, + "cache_read_input_token_cost_above_512k_tokens": higher_rate, + "output_cost_per_token_above_512k_tokens": higher_rate, + }, + ) + rates: Final = _get_token_base_cost( + model_info=info, + usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3), + ) + assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001) + assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002) + assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004) + assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003) + + def test_threshold_keys_exclude_service_tier_variants(): from typing import cast