From c52d23e17151ea984e4cc06e664c58883f389e2a Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Sat, 11 Jul 2026 13:43:04 -0700 Subject: [PATCH] fix(proxy): reserve budget for tiered pricing Ensure tier-only models reserve their estimated request cost so concurrent requests cannot bypass exhausted budgets. Co-authored-by: Cursor --- .../spend_tracking/budget_reservation.py | 49 ++++++++++-- .../proxy/test_budget_reservation.py | 80 +++++++++++++++++++ 2 files changed, 124 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index b577513fc0e..c72049e9c74 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -10,6 +10,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.llms.dashscope.cost_calculator import _calculate_tiered_cost from litellm.proxy._types import ( LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -927,9 +928,6 @@ def _estimate_request_input_cost_for_model( model_info = _get_model_cost_info(model=model, llm_router=llm_router) if model_info is None: return None - input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) - if input_cost_per_token is None: - return None input_tokens = _estimate_input_tokens( request_body=request_body, route=route, @@ -938,6 +936,16 @@ def _estimate_request_input_cost_for_model( ) if input_tokens is None: return None + tiered_pricing = model_info.get("tiered_pricing") + if isinstance(tiered_pricing, list) and tiered_pricing: + return _calculate_tiered_cost( + tokens=input_tokens, + tiered_pricing=tiered_pricing, + cost_key="input_cost_per_token", + ) + input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) + if input_cost_per_token is None: + return None return input_tokens * input_cost_per_token @@ -974,13 +982,25 @@ def _estimate_request_max_cost_for_model( if input_tokens is None or output_tokens is None: return None + output_multiplier = _get_output_multiplier(request_body=request_body) + tiered_pricing = model_info.get("tiered_pricing") + if isinstance(tiered_pricing, list) and tiered_pricing: + return _calculate_tiered_cost( + tokens=input_tokens, + tiered_pricing=tiered_pricing, + cost_key="input_cost_per_token", + ) + _calculate_tiered_cost( + tokens=output_tokens * output_multiplier, + tiered_pricing=tiered_pricing, + cost_key="output_cost_per_token", + ) + cost = 0.0 if input_cost_per_token is not None: cost += input_tokens * input_cost_per_token elif input_tokens > 0: return None - output_multiplier = _get_output_multiplier(request_body=request_body) if output_cost_per_token is not None: cost += output_tokens * output_multiplier * output_cost_per_token elif output_tokens > 0: @@ -1035,7 +1055,26 @@ def _get_model_cost_info( try: model_group_info = llm_router.get_model_group_info(model_group=model) if model_group_info is not None: - return model_group_info.model_dump() + model_group_cost_info = model_group_info.model_dump() + deployments = llm_router.get_model_list(model_name=model) or [] + for deployment in deployments: + model_id = deployment.get("model_info", {}).get("id") + backend_model = deployment.get("litellm_params", {}).get("model") + if not isinstance(model_id, str) or not isinstance(backend_model, str): + continue + deployment_model_info = llm_router.get_deployment_model_info( + model_id=model_id, + model_name=backend_model, + ) + if deployment_model_info is None: + continue + tiered_pricing = deployment_model_info.get("tiered_pricing") + if isinstance(tiered_pricing, list) and tiered_pricing: + return { + **model_group_cost_info, + "tiered_pricing": tiered_pricing, + } + return model_group_cost_info except Exception: verbose_proxy_logger.debug( "Unable to load router model group info for budget reservation", diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 75242af81f4..8e00db78a08 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -26,6 +26,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, ) from litellm.proxy.utils import ProxyLogging +from litellm.router import Router @pytest.fixture() @@ -744,6 +745,85 @@ async def test_should_clamp_reservation_to_default_when_output_cap_missing( await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_should_reserve_tiered_pricing_cost(spend_counter_state): + _, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + router = Router( + model_list=[ + { + "model_name": "dashscope/qwen3-max", + "litellm_params": { + "model": "dashscope/qwen3-max", + "api_key": "sk-fake", + }, + "model_info": { + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [0, 32000], + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [32000, 128000], + }, + ], + }, + } + ] + ) + request_body = { + "model": "dashscope/qwen3-max", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + } + + estimated_cost = estimate_request_max_cost( + request_body=request_body, + route="/chat/completions", + llm_router=router, + ) + assert estimated_cost is not None + assert estimated_cost > 0 + + valid_token = UserAPIKeyAuth( + token="key-tiered-pricing", + spend=0.0, + max_budget=estimated_cost, + ) + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=router, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(estimated_cost) + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=router, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequests( spend_counter_state,