From fe79bbdf644bcc1638ebcccb752ae8bedb6f6ce1 Mon Sep 17 00:00:00 2001 From: suleimanelkhoury Date: Wed, 15 Apr 2026 11:48:09 +0200 Subject: [PATCH] fix(auth): optimize budget checks by skipping for zero-cost models for custom auth workflow --- litellm/proxy/auth/user_api_key_auth.py | 106 +++++++++++++----------- 1 file changed, 57 insertions(+), 49 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2b9aee775b3..c46d9182b90 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, )