scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

fix(proxy): reserve budget for tiered pricing

Browse source 
Ensure tier-only models reserve their estimated request cost so concurrent requests cannot bypass exhausted budgets.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Shivam Rawat 2026-07-11 13:43:04 -07:00
parent 2c1d62ce2b
commit c52d23e171
2 changed files with 124 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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",

80
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -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,