From c68da92debe5c38843aa5fa14bcf3b0112f29dfc Mon Sep 17 00:00:00 2001 From: LH-kevin Date: Sat, 22 Aug 2026 02:01:23 +0800 Subject: [PATCH] fix(cost): support custom service-tier thresholds --- .../litellm_core_utils/llm_cost_calc/utils.py | 20 +++-- litellm/router.py | 2 +- .../test_threshold_cost_selection.py | 75 +++++++++++++++++++ .../test_router_model_cost_isolation.py | 35 +++++++++ 4 files changed, 123 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 6c48381f0e0..3c382f29755 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -55,7 +55,7 @@ _ABOVE_TOKEN_THRESHOLD_COST_KEY: Final[re.Pattern[str]] = re.compile( r"output_cost_per_token|" r"cache_creation_input_token_cost(?:_above_1hr)?|" r"cache_read_input_token_cost" - r")_above_\d+k?_tokens$" + r")_above_\d+k?_tokens(?:_(?Ppriority|flex|ultrafast))?$" ) _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType( @@ -323,16 +323,20 @@ def _get_token_base_cost( cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key)) ## CHECK IF ABOVE THRESHOLD - selected_threshold_keys: dict[str, tuple[float, str]] = {} # mutable-ok: threshold accumulator + selected_threshold_keys: dict[str, tuple[float, str, str | None]] = {} # mutable-ok: threshold accumulator for key in model_info: - if "_above_" not in key or not key.endswith("_tokens"): + if "_above_" not in key or not key.endswith(("_tokens", "_priority", "_flex", "_ultrafast")): continue match = _ABOVE_TOKEN_THRESHOLD_COST_KEY.fullmatch(key) if match is None or model_info.get(key) is None: continue + threshold_key_tier = match.group("tier") + if threshold_key_tier is not None and threshold_key_tier != service_tier: + continue + try: threshold = _parse_above_token_threshold(key) except (IndexError, ValueError): @@ -348,6 +352,7 @@ def _get_token_base_cost( selected_threshold_keys[base_key] = ( threshold, key, + threshold_key_tier, ) costs_by_base_key = { # mutable-ok: selected-rate updates @@ -358,16 +363,15 @@ def _get_token_base_cost( "cache_read_input_token_cost": cache_read_cost, } - for base_key, (_, threshold_key) in selected_threshold_keys.items(): - tiered_key = _get_service_tier_cost_key( - threshold_key, - service_tier, + for base_key, (_, threshold_key, threshold_key_tier) in selected_threshold_keys.items(): + cost_key = ( + threshold_key if threshold_key_tier is not None else _get_service_tier_cost_key(threshold_key, service_tier) ) costs_by_base_key[base_key] = cast( float, _get_cost_per_unit( model_info, - tiered_key, + cost_key, costs_by_base_key[base_key], ), ) diff --git a/litellm/router.py b/litellm/router.py index 13eb0038dde..56ff872d2ed 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -293,7 +293,7 @@ _CUSTOM_PRICING_THRESHOLD_COST_KEY: Final[re.Pattern[str]] = re.compile( r"output_cost_per_token|" r"cache_creation_input_token_cost(?:_above_1hr)?|" r"cache_read_input_token_cost" - r")_above_\d+k?_tokens$" + r")_above_\d+k?_tokens(?:_(?:priority|flex|ultrafast))?$" ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_threshold_cost_selection.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_threshold_cost_selection.py index af893cccc49..5f4e1490b92 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_threshold_cost_selection.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_threshold_cost_selection.py @@ -68,3 +68,78 @@ def test_cost_types_select_thresholds_independently() -> None: assert rates[0] == pytest.approx(1e-6) assert rates[1] == pytest.approx(18e-6) + + +@pytest.mark.parametrize( + ("tier_field", "tier_value", "rate_index"), + [ + ("output_cost_per_token_above_200k_tokens_priority", 1.5e-6, 1), + ("cache_read_input_token_cost_above_200k_tokens_priority", 4e-7, 4), + ], +) +def test_tier_qualified_threshold_field_selects_tier_for_matching_service_tier( + tier_field: str, tier_value: float, rate_index: int +) -> None: + """A tier-qualified threshold key with no standard sibling must apply when the request's + service tier matches.""" + model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 5e-7, + tier_field: tier_value, + } + usage = Usage( + prompt_tokens=250_000, + completion_tokens=1_000, + total_tokens=251_000, + ) + + rates = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier="priority", + ) + + assert rates[rate_index] == pytest.approx(tier_value) + + +def test_tier_qualified_threshold_is_ignored_under_the_default_tier() -> None: + """service_tier=None must keep billing the flat rate: the priority key is inactive.""" + model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 5e-7, + "output_cost_per_token_above_200k_tokens_priority": 1.5e-6, + } + usage = Usage( + prompt_tokens=250_000, + completion_tokens=1_000, + total_tokens=251_000, + ) + + rates = _get_token_base_cost(model_info=model_info, usage=usage) + + assert rates[1] == pytest.approx(2e-6) + + +def test_tier_qualified_threshold_is_ignored_under_a_different_service_tier() -> None: + """A priority-qualified key must not bill a flex request.""" + model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 5e-7, + "output_cost_per_token_above_200k_tokens_priority": 1.5e-6, + } + usage = Usage( + prompt_tokens=250_000, + completion_tokens=1_000, + total_tokens=251_000, + ) + + rates = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier="flex", + ) + + assert rates[1] == pytest.approx(2e-6) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 20c4facaab2..8824eb4509e 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -73,6 +73,41 @@ def test_is_custom_pricing_field_recognizes_declared_and_arbitrary_threshold_fie assert not _is_custom_pricing_field("input_cost_per_token_above_32k_tokens_extra") +def test_is_custom_pricing_field_recognizes_tier_qualified_threshold_fields(): + assert _is_custom_pricing_field("input_cost_per_token_above_200k_tokens_priority") + assert _is_custom_pricing_field("output_cost_per_token_above_200k_tokens_flex") + assert _is_custom_pricing_field("cache_read_input_token_cost_above_200k_tokens_priority") + assert _is_custom_pricing_field("cache_creation_input_token_cost_above_1hr_above_200k_tokens_ultrafast") + assert not _is_custom_pricing_field("api_key_above_32k_tokens_priority") + assert not _is_custom_pricing_field("secret_above_32k_tokens_flex") + assert not _is_custom_pricing_field("credential_above_200k_tokens_ultrafast") + assert not _is_custom_pricing_field("custom_provider_param_above_32k_tokens_priority") + + +def test_custom_tier_qualified_threshold_key_is_registered(): + """A deployment declaring only a tier-qualified threshold rate (e.g. + output_cost_per_token_above_200k_tokens_priority, which is not enumerated in + CustomPricingLiteLLMParams) must keep it in its model_cost entry so matching-tier + requests bill it.""" + model_id = "custom-tier-only-threshold" + Router( + model_list=[ + { + "model_name": "tier-only-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key", + "input_cost_per_token": 1e-6, + "output_cost_per_token_above_32k_tokens_priority": 1.5e-6, + }, + "model_info": {"id": model_id}, + } + ] + ) + + registered = litellm.model_cost[model_id] + assert registered["output_cost_per_token_above_32k_tokens_priority"] == 1.5e-6 + def test_copy_custom_pricing_fields_preserves_declared_and_arbitrary_threshold_fields(): """Copy supported pricing fields without copying unrelated LiteLLM params.""" model_info = {}