fix(cost): support custom service-tier thresholds

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(?:_(?P<tier>priority|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],
             ),
         )

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))?$"
 )
 
 

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)

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 = {}