fix(cost): enforce PTU zeroing for tier thresholds

litellm/litellm_core_utils/ptu_pricing.py

@@ -54,7 +54,7 @@ _THRESHOLD_RATE_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/test_ptu_pricing.py

@@ -180,6 +180,25 @@ def test_arbitrary_threshold_rate_the_deployment_declares_is_zeroed_too(threshol
     assert override.get(threshold_field, 9e-06) == 0.0
 
 
+@pytest.mark.parametrize(
+    "threshold_field",
+    [
+        "output_cost_per_token_above_32k_tokens_priority",
+        "cache_read_input_token_cost_above_32k_tokens_priority",
+    ],
+)
+def test_service_tier_qualified_threshold_rate_declared_is_zeroed_too(threshold_field):
+    """A tier-qualified threshold rate the deployment declares must be zeroed like every
+    other declared rate, or a PTU deployment bills per token past the threshold."""
+    assert threshold_field not in CUSTOM_PRICING_FIELDS
+    assert threshold_field not in PTU_ZEROED_PRICING_FIELDS
+
+    override = _with_flag(_VALID, declared={threshold_field: 9e-06})
+
+    assert override is not None
+    assert override.get(threshold_field, 9e-06) == 0.0
+
+
 @pytest.mark.parametrize(
     "param_field",
     [
@@ -213,6 +232,20 @@ def test_is_threshold_rate_key_accepts_supported_threshold_rates(field):
     assert is_threshold_rate_key(field)
 
 
+@pytest.mark.parametrize(
+    "field",
+    [
+        "input_cost_per_token_above_32k_tokens_priority",
+        "output_cost_per_token_above_32k_tokens_flex",
+        "cache_read_input_token_cost_above_200k_tokens_priority",
+        "cache_creation_input_token_cost_above_32k_tokens_ultrafast",
+        "cache_creation_input_token_cost_above_1hr_above_200k_tokens_priority",
+    ],
+)
+def test_is_threshold_rate_key_accepts_service_tier_qualified_threshold_rates(field):
+    assert is_threshold_rate_key(field)
+
+
 @pytest.mark.parametrize(
     "field",
     [
@@ -220,7 +253,10 @@ def test_is_threshold_rate_key_accepts_supported_threshold_rates(field):
         "secret_above_32k_tokens",
         "credential_above_32k_tokens",
         "custom_provider_param_above_32k_tokens",
-        "input_cost_per_token_above_32k_tokens_priority",
+        "api_key_above_32k_tokens_priority",
+        "secret_above_32k_tokens_flex",
+        "credential_above_200k_tokens_ultrafast",
+        "custom_provider_param_above_32k_tokens_priority",
         "input_cost_per_token",
         "tiered_pricing",
     ],

tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py

@@ -794,7 +794,14 @@ 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"])
+    @pytest.mark.parametrize(
+        "field",
+        [
+            "input_cost_per_token_above_32k_tokens",
+            "output_cost_per_token_above_32k_tokens",
+            "input_cost_per_token_above_32k_tokens_priority",
+        ],
+    )
     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
@@ -889,7 +896,14 @@ 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"])
+    @pytest.mark.parametrize(
+        "field",
+        [
+            "input_cost_per_token_above_32k_tokens",
+            "output_cost_per_token_above_32k_tokens",
+            "output_cost_per_token_above_32k_tokens_priority",
+        ],
+    )
     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={})
@@ -966,6 +980,8 @@ class TestPtuDeploymentsAreNotBilledPerToken:
             "secret_above_32k_tokens",
             "credential_above_32k_tokens",
             "custom_provider_param_above_32k_tokens",
+            "api_key_above_32k_tokens_priority",
+            "custom_provider_param_above_32k_tokens_flex",
         ],
     )
     def test_a_threshold_like_setting_that_is_not_a_price_is_left_alone(self, param):
@@ -1048,6 +1064,42 @@ class TestPtuDeploymentsAreNotBilledPerToken:
             ),
         )
         assert "input_cost_per_token_above_32k_tokens" not in json.loads(off["litellm_params"])
+
+    def test_removing_ptu_config_releases_a_zeroed_tier_qualified_threshold_rate(self):
+        """The zeroing spans tier-qualified 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="tier-threshold",
+                litellm_params=LiteLLM_Params(
+                    model="openai/gpt-4o-mini", output_cost_per_token_above_32k_tokens_priority=9e-06
+                ),
+                model_info=ModelInfo(id="dep-tier-thr", team_id="t"),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(
+                    id="dep-tier-thr",
+                    team_id="t",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                )
+            ),
+        )
+        assert json.loads(on["litellm_params"])["output_cost_per_token_above_32k_tokens_priority"] == 0
+
+        off = update_db_model(
+            db_model=Deployment(
+                model_name="tier-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-tier-thr", ptu_count=None, cost_per_ptu_per_hour=None)
+            ),
+        )
+        assert "output_cost_per_token_above_32k_tokens_priority" 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."""