fix(cost): enforce threshold pricing for DB-backed PTU

litellm/litellm_core_utils/ptu_pricing.py

@@ -56,6 +56,13 @@ _THRESHOLD_RATE_KEY: Final[re.Pattern[str]] = re.compile(
     r"cache_read_input_token_cost"
     r")_above_\d+k?_tokens$"
 )
+
+
+def is_threshold_rate_key(field: str) -> bool:
+    """Whether ``field`` is a custom above-threshold rate key the cost calculator reads."""
+    return _THRESHOLD_RATE_KEY.fullmatch(field) is not None
+
+
 PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()] | Mapping[str, float]]] = MappingProxyType(
     {
         **dict.fromkeys(PTU_ZEROED_PRICING_FIELDS, 0.0),
@@ -189,9 +196,7 @@ def zeroed_ptu_pricing(
         return None
     if not is_ptu_cost_attribution_enabled():
         return None
-    declared_threshold_rates: Final = frozenset(
-        key for key in declared if _THRESHOLD_RATE_KEY.fullmatch(key) is not None
-    )
+    declared_threshold_rates: Final = frozenset(key for key in declared if is_threshold_rate_key(key))
     return MappingProxyType(
         {
             **PTU_ZEROED_PRICING,

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -30,6 +30,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
     PTU_ZEROED_PRICING_FIELDS,
     PTU_ZEROED_TABLE_FIELDS,
     SEARCH_CONTEXT_SIZES,
+    is_threshold_rate_key,
     ptu_config_error,
 )
 from litellm.proxy._types import (
@@ -379,6 +380,9 @@ def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supp
         sorted(
             tuple(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field)))
             + tuple(field for field in _PTU_EMPTIED_PRICING_FIELDS if supplied.get(field))
+            + tuple(
+                field for field in supplied if is_threshold_rate_key(field) and _is_nonzero_price(supplied.get(field))
+            )
         )
     )
     if not priced:
@@ -418,9 +422,12 @@ def _ptu_zeroed_pricing(
     if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
         return _NO_PRICING_OVERRIDE
     _raise_if_ptu_deployment_is_priced(model_info=model_info, supplied=supplied)
+    candidates: Final = _CUSTOM_PRICING_FIELDS | frozenset(
+        field for field in (*model_info.keys(), *litellm_params.keys()) if is_threshold_rate_key(field)
+    )
     stored: Final = frozenset(
         field
-        for field in _CUSTOM_PRICING_FIELDS
+        for field in candidates
         if _is_nonzero_price(model_info.get(field)) or _is_nonzero_price(litellm_params.get(field))
     )
     return MappingProxyType(
@@ -460,7 +467,11 @@ def _ptu_pricing_delta(
         return _NO_PRICING_OVERRIDE, frozenset()
     return _NO_PRICING_OVERRIDE, frozenset(
         field
-        for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS, _PTU_EMPTIED_PRICING_FIELDS)
+        for field in _CUSTOM_PRICING_FIELDS.union(
+            _PTU_ZEROED_PRICING_FIELDS,
+            _PTU_EMPTIED_PRICING_FIELDS,
+            frozenset(field for field in (*model_info.keys(), *litellm_params.keys()) if is_threshold_rate_key(field)),
+        )
         if _is_zero_price(model_info.get(field)) or _is_zero_price(litellm_params.get(field))
     )
 

tests/test_litellm/litellm_core_utils/test_ptu_pricing.py

@@ -7,6 +7,7 @@ from unittest.mock import patch
 import pytest
 
 from litellm.litellm_core_utils.ptu_pricing import (
+    is_threshold_rate_key,
     ptu_config_error,
     CUSTOM_PRICING_FIELDS,
     PTU_EMPTIED_PRICING_FIELDS,
@@ -198,6 +199,36 @@ def test_threshold_like_non_pricing_params_are_left_alone(param_field):
     assert param_field not in override
 
 
+@pytest.mark.parametrize(
+    "field",
+    [
+        "input_cost_per_token_above_32k_tokens",
+        "output_cost_per_token_above_64k_tokens",
+        "cache_read_input_token_cost_above_128k_tokens",
+        "cache_creation_input_token_cost_above_200k_tokens",
+        "cache_creation_input_token_cost_above_1hr_above_200k_tokens",
+    ],
+)
+def test_is_threshold_rate_key_accepts_supported_threshold_rates(field):
+    assert is_threshold_rate_key(field)
+
+
+@pytest.mark.parametrize(
+    "field",
+    [
+        "api_key_above_32k_tokens",
+        "secret_above_32k_tokens",
+        "credential_above_32k_tokens",
+        "custom_provider_param_above_32k_tokens",
+        "input_cost_per_token_above_32k_tokens_priority",
+        "input_cost_per_token",
+        "tiered_pricing",
+    ],
+)
+def test_is_threshold_rate_key_rejects_non_pricing_fields(field):
+    assert not is_threshold_rate_key(field)
+
+
 def test_a_setting_that_is_not_a_charge_is_left_alone():
     """CustomPricingLiteLLMParams also carries configuration, and zeroing one of those
     would break the deployment rather than stop a charge."""

tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py

@@ -794,6 +794,15 @@ class TestPtuDeploymentsAreNotBilledPerToken:
         assert exc.value.status_code == 400
         assert "tiered_pricing" in str(exc.value.detail)
 
+    @pytest.mark.parametrize("field", ["input_cost_per_token_above_32k_tokens", "output_cost_per_token_above_32k_tokens"])
+    def test_an_arbitrary_threshold_rate_the_caller_supplies_is_refused(self, field):
+        """The router lets deployments declare arbitrary above-threshold rates; those bill a
+        PTU deployment past the threshold just as surely as a flat rate, so the refusal must
+        cover them too."""
+        with pytest.raises(HTTPException) as exc:
+            self._zeroed(model_info=self.PTU, supplied={field: 9e-06})
+        assert exc.value.status_code == 400
+        assert field in str(exc.value.detail)
     def test_a_search_context_price_the_caller_supplies_is_refused(self):
         """The rates sit in a table keyed by context size, so a guard that only reads numbers
         lets a per-request charge onto a deployment its reserved capacity already pays for."""
@@ -880,6 +889,12 @@ class TestPtuDeploymentsAreNotBilledPerToken:
         assert zeroed["input_cost_per_second"] == 0
         assert zeroed["input_cost_per_token"] == 0
 
+    @pytest.mark.parametrize("field", ["input_cost_per_token_above_32k_tokens", "output_cost_per_token_above_32k_tokens"])
+    def test_an_arbitrary_threshold_rate_already_on_the_row_is_zeroed(self, field):
+        """A row priced through a path this rule does not cover must heal on its next save."""
+        zeroed = self._zeroed(model_info={**self.PTU, field: 9e-06}, litellm_params={})
+        assert zeroed[field] == 0
+        assert zeroed["input_cost_per_token"] == 0
     @pytest.mark.asyncio
     async def test_a_refused_price_does_not_leave_the_team_changed(self):
         """The team ACL write autocommits, so the refusal has to run before it. Otherwise a
@@ -944,6 +959,32 @@ class TestPtuDeploymentsAreNotBilledPerToken:
         assert priced.litellm_params.get("regional_processing_uplift_multiplier_eu") == 1.15
         assert priced.litellm_params.get("input_cost_per_token") == 0
 
+    @pytest.mark.parametrize(
+        "param",
+        [
+            "api_key_above_32k_tokens",
+            "secret_above_32k_tokens",
+            "credential_above_32k_tokens",
+            "custom_provider_param_above_32k_tokens",
+        ],
+    )
+    def test_a_threshold_like_setting_that_is_not_a_price_is_left_alone(self, param):
+        """Only the cost calculator's base keys qualify as threshold rates: a deployment param
+        that merely ends in _above_<N>k_tokens is not a charge and must keep its value."""
+        priced = _ptu_priced_deployment(
+            Deployment(
+                model_name="settings",
+                litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", **{param: "keep-me"}),
+                model_info=ModelInfo(
+                    id="dep-settings",
+                    team_id="t",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                ),
+            )
+        )
+        assert priced.litellm_params.get(param) == "keep-me"
+        assert priced.litellm_params.get("input_cost_per_token") == 0
     def test_removing_ptu_config_releases_every_rate_it_zeroed(self):
         """The zeroing covers any stored rate, so a release that only spans the mirrored fields
         leaves a per-second deployment billing nothing for that dimension forever."""
@@ -976,6 +1017,37 @@ class TestPtuDeploymentsAreNotBilledPerToken:
         )
         assert "input_cost_per_second" not in json.loads(off["litellm_params"])
 
+    def test_removing_ptu_config_releases_a_zeroed_threshold_rate(self):
+        """The zeroing spans threshold rates too, so a release that only spans the enumeration
+        sets would leave one billing nothing past the threshold forever."""
+        on = update_db_model(
+            db_model=Deployment(
+                model_name="threshold",
+                litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", input_cost_per_token_above_32k_tokens=9e-06),
+                model_info=ModelInfo(id="dep-thr", team_id="t"),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(
+                    id="dep-thr",
+                    team_id="t",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                )
+            ),
+        )
+        assert json.loads(on["litellm_params"])["input_cost_per_token_above_32k_tokens"] == 0
+
+        off = update_db_model(
+            db_model=Deployment(
+                model_name="threshold",
+                litellm_params=LiteLLM_Params(**json.loads(on["litellm_params"])),
+                model_info=ModelInfo(**json.loads(on["model_info"])),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(id="dep-thr", ptu_count=None, cost_per_ptu_per_hour=None)
+            ),
+        )
+        assert "input_cost_per_token_above_32k_tokens" not in json.loads(off["litellm_params"])
     def test_removing_ptu_config_releases_a_zeroed_search_context_table(self):
         """The all-zero table exists only to stop the double charge, so a deployment taken off PTU
         has to give it up or it keeps serving grounded requests for free forever."""