diff --git a/litellm/main.py b/litellm/main.py index f471cae2656..1494ff2226c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1502,9 +1502,10 @@ def _token_rates_cleared_for_per_second_overrides(custom: Mapping[str, object]) def _tiers_with_flat_overrides(catalog: Mapping[str, object], custom: Mapping[str, object]) -> Mapping[str, object]: """A ``tiered_pricing`` table wins over flat rates in the cost calculator, so a flat rate the caller set is written - into every tier for that direction, and the other direction keeps each tier's own rate.""" + into every tier of the table in force (the caller's own, else the inherited one) for that direction, and the other + direction keeps each tier's own rate.""" overrides: Final = {key: custom[key] for key in _FLAT_TOKEN_RATE_FIELDS if key in custom} - tiers: Final = catalog.get("tiered_pricing") + tiers: Final = custom.get("tiered_pricing", catalog.get("tiered_pricing")) if not overrides or not isinstance(tiers, (list, tuple)) or not tiers: return {} return {"tiered_pricing": [{**tier, **overrides} for tier in _TIERS.validate_python(tiers)]} diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 06822916fd9..afff995aa69 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -1667,3 +1667,35 @@ def test_direct_call_with_only_a_per_second_rate_on_a_mapped_model_is_billed_per finally: litellm.model_cost.pop(model_key, None) _invalidate_model_cost_lowercase_map() + + +def test_direct_call_flat_rate_applies_to_the_callers_own_tiers_not_the_catalogs(): + model_key: Final = "openai/lit-caller-tiers-model" + litellm.register_model( + { + model_key: { + "litellm_provider": "openai", + "mode": "chat", + "tiered_pricing": [ + {"range": [0, 1000000], "input_cost_per_token": 3e-07, "output_cost_per_token": 6e-07} + ], + } + }, + persist_across_reloads=False, + ) + try: + response: Final = litellm.completion( + model=model_key, + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + mock_response="ok", + input_cost_per_token=1e-05, + tiered_pricing=[{"range": [0, 1000000], "input_cost_per_token": 5e-06, "output_cost_per_token": 9e-06}], + ) + assert response._hidden_params["response_cost"] == pytest.approx( + response.usage.prompt_tokens * 1e-05 + response.usage.completion_tokens * 9e-06 + ) + assert litellm.model_cost[model_key]["tiered_pricing"][0]["output_cost_per_token"] == 6e-07 + finally: + litellm.model_cost.pop(model_key, None) + _invalidate_model_cost_lowercase_map()