scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(auth): optimize budget checks by skipping for zero-cost models for custom auth workflow

Browse source
This commit is contained in:
suleimanelkhoury 2026-04-15 11:48:09 +02:00
parent 9b27d10220
commit fe79bbdf64
1 changed files with 57 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
litellm/proxy/auth/user_api_key_auth.py
View file

@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import (
_cache_key_object,
_delete_cache_key_object,
_get_user_role,
_is_model_cost_zero,
_is_user_proxy_admin,
_virtual_key_max_budget_alert_check,
_virtual_key_max_budget_check,
@ -879,8 +880,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
model = get_model_from_request(request_data, route)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
skip_budget_checks = _is_model_cost_zero(
model=model, llm_router=llm_router
)
@ -1286,8 +1285,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
model = get_model_from_request(request_data, route)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
skip_budget_checks = _is_model_cost_zero(
model=model, llm_router=llm_router
)
@ -1924,54 +1921,65 @@ async def _run_post_custom_auth_checks(
llm_router=llm_router,
)
# 2a. Check if model has zero cost - if so, skip all budget checks
model = get_model_from_request(request_data, route)
skip_budget_checks = False
if model is not None and llm_router is not None:
skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
if skip_budget_checks:
verbose_proxy_logger.info(
f"Skipping all budget checks for zero-cost model: {model}"
)
current_model = request_data.get("model", None)
# 3. Check key-level model_max_budget
max_budget_per_model = valid_token.model_max_budget
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and current_model is not None
and valid_token.token is not None
):
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
)
if not skip_budget_checks:
# 3. Check key-level model_max_budget
max_budget_per_model = valid_token.model_max_budget
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and current_model is not None
and valid_token.token is not None
):
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
)
# 4. Check end-user model_max_budget
end_user_mmb = valid_token.end_user_model_max_budget
if (
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
# 4. Check end-user model_max_budget
end_user_mmb = valid_token.end_user_model_max_budget
if (
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
# 4b. Check end-user max_budget
end_user_mb = valid_token.end_user_max_budget
# Fallback in case end-user max budget not set on token
if end_user_mb is None and end_user_object is not None:
budget_table = end_user_object.litellm_budget_table
end_user_mb = budget_table.max_budget if budget_table is not None else None
if (
end_user_mb is not None
and end_user_object is not None
and valid_token.end_user_id is not None
):
await proxy_logging_obj.max_budget_limiter.is_end_user_within_budget(
end_user_id=valid_token.end_user_id,
end_user_max_budget=end_user_mb,
end_user_spend=end_user_object.spend,
route=route,
)
# 4b. Check end-user max_budget
end_user_mb = valid_token.end_user_max_budget
# Fallback in case end-user max budget not set on token
if end_user_mb is None and end_user_object is not None:
budget_table = end_user_object.litellm_budget_table
end_user_mb = budget_table.max_budget if budget_table is not None else None
if (
end_user_mb is not None
and end_user_object is not None
and valid_token.end_user_id is not None
):
await proxy_logging_obj.max_budget_limiter.is_end_user_within_budget(
end_user_id=valid_token.end_user_id,
end_user_max_budget=end_user_mb,
end_user_spend=end_user_object.spend,
route=route,
)
# 5. Look up user object if user_id is set
user_object = None
@ -2046,7 +2054,7 @@ async def _run_post_custom_auth_checks(
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=False,
skip_budget_checks=skip_budget_checks,
project_object=_project_obj,
)