diff --git a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py index fb0f130a6cf..0cb0f7004cc 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py @@ -25,6 +25,17 @@ def _coerce_cost_per_token(value: float | str | None) -> float: return float(value) +def _resolve_tier_cost_per_token( + tier: dict, + cost_key: str, + fallback_cost_key: str | None, +) -> float: + primary: Final = tier.get(cost_key) + if primary is not None: + return _coerce_cost_per_token(primary) + return _coerce_cost_per_token(tier.get(fallback_cost_key, 0)) + + def calculate_tiered_cost( tokens: int, tiered_pricing: list[dict], @@ -84,8 +95,7 @@ def calculate_tiered_cost( if tier_end > tier_start: tokens_in_tier = tier_end - tier_start - cost_per_token = tier.get(cost_key) or tier.get(fallback_cost_key, 0) - total_cost += tokens_in_tier * _coerce_cost_per_token(cost_per_token) + total_cost += tokens_in_tier * _resolve_tier_cost_per_token(tier, cost_key, fallback_cost_key) tokens_processed = tier_end # After loop, check if any tokens remain (i.e., tokens > highest tier's end range) @@ -93,8 +103,7 @@ def calculate_tiered_cost( if tokens_processed < tokens and sorted_tiers: last_tier: Final = sorted_tiers[-1] remaining_tokens: Final = tokens - tokens_processed - cost_per_token = last_tier.get(cost_key) or last_tier.get(fallback_cost_key, 0) - total_cost += remaining_tokens * _coerce_cost_per_token(cost_per_token) + total_cost += remaining_tokens * _resolve_tier_cost_per_token(last_tier, cost_key, fallback_cost_key) return total_cost @@ -135,5 +144,4 @@ def tier_rate( fallback_cost_key: str | None = None, ) -> float: """Read a per-token rate from a tier, coercing YAML string costs to float.""" - raw: Final = tier.get(cost_key) or tier.get(fallback_cost_key, 0) - return _coerce_cost_per_token(raw) + return _resolve_tier_cost_per_token(tier, cost_key, fallback_cost_key) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tiered_pricing.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tiered_pricing.py new file mode 100644 index 00000000000..9fba8cb1fab --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tiered_pricing.py @@ -0,0 +1,72 @@ +import os +import sys + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( + calculate_tiered_cost, + tier_rate, +) + + +def test_calculate_tiered_cost_honors_explicit_zero_primary_rate(): + tiers = [ + { + "range": [0, 100000], + "input_cost_per_token": 1e-6, + "cache_read_input_token_cost": 0.0, + } + ] + + cost = calculate_tiered_cost( + tokens=10000, + tiered_pricing=tiers, + cost_key="cache_read_input_token_cost", + fallback_cost_key="input_cost_per_token", + ) + + assert cost == 0.0 + + +def test_calculate_tiered_cost_honors_explicit_zero_primary_rate_in_overflow_tier(): + tiers = [ + { + "range": [0, 100000], + "input_cost_per_token": 1e-6, + "cache_read_input_token_cost": 0.0, + } + ] + + cost = calculate_tiered_cost( + tokens=150000, + tiered_pricing=tiers, + cost_key="cache_read_input_token_cost", + fallback_cost_key="input_cost_per_token", + ) + + assert cost == 0.0 + + +def test_calculate_tiered_cost_falls_back_when_primary_rate_is_missing(): + tiers = [{"range": [0, 100000], "input_cost_per_token": 1e-6}] + + cost = calculate_tiered_cost( + tokens=10000, + tiered_pricing=tiers, + cost_key="cache_read_input_token_cost", + fallback_cost_key="input_cost_per_token", + ) + + assert cost == 10000 * 1e-6 + + +def test_tier_rate_honors_explicit_zero_primary_rate(): + tier = {"cache_read_input_token_cost": 0.0, "input_cost_per_token": 1e-6} + + assert tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token") == 0.0 + + +def test_tier_rate_falls_back_when_primary_rate_is_missing(): + tier = {"input_cost_per_token": 1e-6} + + assert tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token") == 1e-6