From 9cb966ef78b144605e041c7506d6afee4ebca15e Mon Sep 17 00:00:00 2001 From: LH-kevin Date: Sat, 22 Aug 2026 01:25:25 +0800 Subject: [PATCH] fix(cost): enforce threshold pricing for DB-backed PTU --- litellm/litellm_core_utils/ptu_pricing.py | 11 ++- .../model_management_endpoints.py | 15 +++- .../litellm_core_utils/test_ptu_pricing.py | 31 ++++++++ .../test_ptu_model_settings.py | 72 +++++++++++++++++++ 4 files changed, 124 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py index b599beb2dee..5df0870ac3c 100644 --- a/litellm/litellm_core_utils/ptu_pricing.py +++ b/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, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ec98d7d65f1..3a46aedf09d 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/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)) ) diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py index 673599b0e82..0aa3ce9e62f 100644 --- a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py +++ b/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.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index a1c38d26b9d..1df6ec1dd29 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/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_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."""