diff --git a/litellm/_redis.py b/litellm/_redis.py index b1c2b16d7f1..ed9f3580162 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -395,9 +395,7 @@ def _get_redis_client_logic(**env_overrides): if _sentinel_password is not None: redis_kwargs["sentinel_password"] = _sentinel_password - _service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret( - "REDIS_SERVICE_NAME" - ) + _service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret("REDIS_SERVICE_NAME") if _service_name is not None: redis_kwargs["service_name"] = _service_name diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 65afa32d5c0..6fceebbea29 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4081,14 +4081,14 @@ async def _team_model_max_budget_check( """ if team_object is None or team_object.team_id is None or model is None: return - team_model_max_budget = team_object.model_max_budget + team_model_max_budget: Final = team_object.model_max_budget if not isinstance(team_model_max_budget, dict) or len(team_model_max_budget) == 0: return from litellm.proxy.proxy_server import model_max_budget_limiter - models = (model,) if isinstance(model, str) else tuple(model) - key_model_max_budget = valid_token.model_max_budget if valid_token is not None else None + models: Final = (model,) if isinstance(model, str) else tuple(model) + key_model_max_budget: Final = valid_token.model_max_budget if valid_token is not None else None for model_name in models: await model_max_budget_limiter.is_team_within_model_budget( team_id=team_object.team_id, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index de901e3b66d..ec84649fc88 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2337,7 +2337,9 @@ async def _run_centralized_common_checks( if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id - user_api_key_auth_obj.team_model_max_budget = team_object.model_max_budget if team_object is not None else None + user_api_key_auth_obj.team_model_max_budget = ( # rebind-ok: this hydration point mutates the caller's auth object by contract, like org_id above + team_object.model_max_budget if team_object is not None else None + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 38515d0b2ff..00d00aa6d1b 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -71,11 +71,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model=model, key_budget_config=_current_model_budget_info, ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): + if _current_spend is not None and _current_spend > _current_model_budget_info.max_budget: raise litellm.BudgetExceededError( message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", current_cost=_current_spend, @@ -137,11 +133,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model=model, key_budget_config=_current_model_budget_info, ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): + if _current_spend is not None and _current_spend > _current_model_budget_info.max_budget: raise litellm.BudgetExceededError( message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", current_cost=_current_spend, @@ -170,7 +162,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): """ if model in internal_model_max_budget: return model - stripped_model = self._get_model_without_custom_llm_provider(model) + stripped_model: Final = self._get_model_without_custom_llm_provider(model) if stripped_model in internal_model_max_budget: return stripped_model return next( @@ -195,7 +187,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): """ if not key_model_max_budget: return False - budget_config = self._get_request_model_budget_config( + budget_config: Final = self._get_request_model_budget_config( model=model, internal_model_max_budget=self._coerce_budget_configs(key_model_max_budget) ) return ( @@ -226,22 +218,22 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): verbose_proxy_logger.debug("Team model budget check skipped for model=%s: key has its own entry", model) return True - internal_model_max_budget = self._coerce_budget_configs(team_model_max_budget) + internal_model_max_budget: Final = self._coerce_budget_configs(team_model_max_budget) verbose_proxy_logger.debug("team internal_model_max_budget %s", internal_model_max_budget) - matched_model = self._get_matched_budget_model_name( + matched_model: Final = self._get_matched_budget_model_name( model=model, internal_model_max_budget=internal_model_max_budget ) if matched_model is None: verbose_proxy_logger.debug("Model %s not found in team_model_max_budget", model) return True - budget_config = internal_model_max_budget[matched_model] + budget_config: Final = internal_model_max_budget[matched_model] if not budget_config.max_budget or budget_config.max_budget <= 0: return True - current_spend = await self._get_team_spend_for_model( + current_spend: Final = await self._get_team_spend_for_model( team_id=team_id, matched_model=matched_model, budget_config=budget_config, @@ -276,7 +268,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): write always address one counter regardless of how the request spelled the model. """ - team_model_spend_cache_key = ( + team_model_spend_cache_key: Final = ( f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{matched_model}:{budget_config.budget_duration}" ) return await self.dual_cache.async_get_cache( @@ -304,20 +296,22 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if self._key_already_covers_model(key_model_max_budget, model): return - internal_model_max_budget = self._coerce_budget_configs(team_model_max_budget) - matched_model = self._get_matched_budget_model_name( + internal_model_max_budget: Final = self._coerce_budget_configs(team_model_max_budget) + matched_model: Final = self._get_matched_budget_model_name( model=model, internal_model_max_budget=internal_model_max_budget ) if matched_model is None: return - budget_config = internal_model_max_budget[matched_model] + budget_config: Final = internal_model_max_budget[matched_model] if not budget_config.budget_duration: return - team_spend_key = ( + team_spend_key: Final = ( f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{matched_model}:{budget_config.budget_duration}" ) - team_start_time_key = f"team_model_budget_start_time:{team_id}:{matched_model}:{budget_config.budget_duration}" + team_start_time_key: Final = ( + f"team_model_budget_start_time:{team_id}:{matched_model}:{budget_config.budget_duration}" + ) await self._increment_spend_for_key( budget_config=budget_config, spend_key=team_spend_key, @@ -382,14 +376,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): self, model: str, internal_model_max_budget: Mapping[str, BudgetConfig] ) -> BudgetConfig | None: """ - Get the budget config for the request model - - 1. Check if `model` is in `internal_model_max_budget` - 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` + Get the budget config for the request model via _get_matched_budget_model_name, + so request and config-entry names are both `{provider}/`-normalized and a + prefixed entry like `openai/gpt-4` binds a bare `gpt-4` request. """ - return internal_model_max_budget.get(model, None) or internal_model_max_budget.get( - self._get_model_without_custom_llm_provider(model), None + matched_model: Final = self._get_matched_budget_model_name( + model=model, internal_model_max_budget=internal_model_max_budget ) + return internal_model_max_budget[matched_model] if matched_model is not None else None @staticmethod def _get_model_without_custom_llm_provider(model: str) -> str: @@ -427,7 +421,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): user_api_key_end_user_model_max_budget: Final[dict | None] = _metadata.get( "user_api_key_end_user_model_max_budget", None ) - user_api_key_team_model_max_budget: Mapping[str, Mapping[str, str | float]] | None = _metadata.get( + user_api_key_team_model_max_budget: Final[Mapping[str, Mapping[str, str | float]] | None] = _metadata.get( "user_api_key_team_model_max_budget", None ) if ( @@ -503,7 +497,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): response_cost=response_cost, ) - user_api_key_team_id: str | None = _metadata.get("user_api_key_team_id", None) + user_api_key_team_id: Final[str | None] = _metadata.get("user_api_key_team_id", None) if ( user_api_key_team_id is not None and user_api_key_team_model_max_budget is not None diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6b8ce1762c3..8f0ced79149 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1744,7 +1744,9 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = ( user_api_key_dict.end_user_model_max_budget ) - data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget + data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = ( # rebind-ok: pre-call metadata injection + user_api_key_dict.team_model_max_budget + ) # User spend, budget - used by prometheus.py # Follow same pattern as team and API key budgets diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b11211e34a9..e607a2d8a0e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -425,7 +425,7 @@ def _team_key_model_max_budget_check( return if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return - team_member_object = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) + team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) if team_member_object is not None and team_member_object.role == "admin": return raise HTTPException( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 09cd6e6def6..1724d254ddb 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1056,11 +1056,11 @@ def _check_team_model_max_budget_update_authority( from litellm.types.utils import BudgetConfig - existing_configs = MappingProxyType( + existing_configs: Final = MappingProxyType( {_model: BudgetConfig(**_info) for _model, _info in existing_model_max_budget.items()} ) - requested_items = data.model_max_budget.items() if data.model_max_budget else () - requested_configs = MappingProxyType({_model: BudgetConfig(**_info) for _model, _info in requested_items}) + requested_items: Final = data.model_max_budget.items() if data.model_max_budget else () + requested_configs: Final = MappingProxyType({_model: BudgetConfig(**_info) for _model, _info in requested_items}) for model_name, existing_config in existing_configs.items(): requested_config = requested_configs.get(model_name) if requested_config is None: @@ -1111,11 +1111,11 @@ def _validate_team_model_max_budget( code="400", ) - normalized_names = tuple( + normalized_names: Final = tuple( _PROXY_VirtualKeyModelMaxBudgetLimiter._get_model_without_custom_llm_provider(entry_name) for entry_name in model_max_budget ) - colliding_entries = tuple( + colliding_entries: Final = tuple( sorted( entry_name for entry_name, normalized_name in zip(model_max_budget, normalized_names) @@ -1134,6 +1134,32 @@ def _validate_team_model_max_budget( code="400", ) + from types import MappingProxyType + + from litellm.types.utils import BudgetConfig + + entry_configs: Final = MappingProxyType( + {entry_name: BudgetConfig(**entry_info) for entry_name, entry_info in model_max_budget.items()} + ) + non_enforceable: Final = tuple( + sorted( + entry_name + for entry_name, config in entry_configs.items() + if config.max_budget is None or not math.isfinite(config.max_budget) or config.max_budget <= 0 + ) + ) + if non_enforceable: + raise ProxyException( + message=( + f"model_max_budget entries {non_enforceable} have no enforceable budget_limit; " + "enforcement skips non-positive caps, so a zero or negative value would leave the " + "model uncapped. Use a positive finite budget_limit, or remove the entry" + ), + type=ProxyErrorTypes.bad_request_error, + param="model_max_budget", + code="400", + ) + def _should_auto_add_team_creator( user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/router.py b/litellm/router.py index 39d5080a33d..cdff78e4987 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3410,9 +3410,7 @@ class Router: # Await the first task to complete successfully while pending_tasks: - done, pending_tasks = await asyncio.wait( - pending_tasks, return_when=asyncio.FIRST_COMPLETED - ) + done, pending_tasks = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED) for completed_task in done: result = await check_response(completed_task) @@ -5240,9 +5238,7 @@ class Router: # Update kwargs with the current model name or any other model-specific adjustments ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## if not custom_llm_provider: - _, custom_llm_provider, _, _ = get_llm_provider( - model=model - ) + _, custom_llm_provider, _, _ = get_llm_provider(model=model) new_kwargs: Final = safe_deep_copy(kwargs) self._update_kwargs_with_deployment( deployment=cast(dict, model_name), @@ -6029,9 +6025,7 @@ class Router: raise Exception( "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" ) - return await original_function( - custom_llm_provider=custom_llm_provider, client=client, **kwargs - ) + return await original_function(custom_llm_provider=custom_llm_provider, client=client, **kwargs) #### [END] ASSISTANTS API #### @@ -6364,9 +6358,7 @@ class Router: mask_sensitive_structure(fallback_model_group), ) if len(fallback_failure_exception_str) > 0: - original_exception.message += ( - f"\nError doing the fallback: {fallback_failure_exception_str}" - ) + original_exception.message += f"\nError doing the fallback: {fallback_failure_exception_str}" raise original_exception @@ -9124,9 +9116,7 @@ class Router: and model_info["supports_parallel_function_calling"] is True ): model_group_info.supports_parallel_function_calling = True - if ( - model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True - ): + if model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True: model_group_info.supports_vision = True if ( model_info.get("supports_function_calling", None) is not None @@ -9144,9 +9134,7 @@ class Router: ): model_group_info.supports_url_context = True - if ( - model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True - ): + if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True: model_group_info.supports_reasoning = True if ( model_info.get("supported_openai_params", None) is not None @@ -10876,9 +10864,7 @@ class Router: args=(e, traceback_exception), ).start() # log response # Handle any exceptions that might occur during streaming - asyncio.create_task( - logging_obj.async_failure_handler(e, traceback_exception) - ) + asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception)) raise e async def async_get_available_deployment_for_pass_through( @@ -11003,9 +10989,7 @@ class Router: target=logging_obj.failure_handler, args=(e, traceback_exception), ).start() - asyncio.create_task( - logging_obj.async_failure_handler(e, traceback_exception) - ) + asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception)) raise e async def _run_routing_plugins( diff --git a/litellm/utils.py b/litellm/utils.py index 19c3d10695a..de96489a816 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -766,16 +766,8 @@ def function_setup( if ( len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 - ) and len( - callback_list - ) == 0: - callback_list = list( - set( - litellm.input_callback - + litellm.success_callback - + litellm.failure_callback - ) - ) + ) and len(callback_list) == 0: + callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) get_set_callbacks: Final = getattr(sys.modules[__name__], "get_set_callbacks") get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS - safety net for callbacks added via direct append diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index cb241a815a8..ceda203aba4 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -61,6 +61,61 @@ def test_get_request_model_budget_config(budget_limiter): assert config is None +def test_get_request_model_budget_config_strips_entry_prefix(budget_limiter): + """Key/end-user entries keyed provider/model must bind bare (and cross-prefixed) + requests; only the request side was normalized before, so an openai/gpt-4 + entry never matched a gpt-4 request and the cap was inert (Greptile finding, + key-side sibling of the team matcher fix).""" + internal_budget = {"openai/gpt-4": GenericBudgetInfo(budget_limit=100.0, time_period="1d")} + matched = budget_limiter._get_request_model_budget_config( + model="gpt-4", internal_model_max_budget=internal_budget + ) + assert matched is not None and matched.max_budget == 100.0 + cross = budget_limiter._get_request_model_budget_config( + model="azure/gpt-4", internal_model_max_budget=internal_budget + ) + assert cross is not None and cross.max_budget == 100.0 + assert ( + budget_limiter._get_request_model_budget_config( + model="claude-3", internal_model_max_budget=internal_budget + ) + is None + ) + + +@pytest.mark.asyncio +async def test_is_key_within_model_budget_prefixed_entry_bare_request(budget_limiter): + """A key cap keyed openai/gpt-4 must enforce on a bare gpt-4 request.""" + user_api_key = UserAPIKeyAuth( + token="test-key", + key_alias="test-alias", + model_max_budget={"openai/gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, + ) + with patch.object(budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0): + with pytest.raises(litellm.BudgetExceededError): + await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") + + +@pytest.mark.asyncio +async def test_team_check_skipped_for_prefixed_key_entry_on_bare_request(budget_limiter): + """A key whose own entry is keyed openai/gpt-4 is exempt from the team cap for + a bare gpt-4 request; before the fix it charged, and was blocked by, the + shared team counter.""" + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=10_000.0) + ) as mock_team_spend: + assert ( + await budget_limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget={"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, + model="gpt-4", + key_model_max_budget={"openai/gpt-4": {"budget_limit": 50.0, "time_period": "1d"}}, + ) + is True + ) + mock_team_spend.assert_not_called() + + # Test is_key_within_model_budget @pytest.mark.asyncio async def test_is_key_within_model_budget(budget_limiter): diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ffd62995f49..39930d27b96 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11108,3 +11108,16 @@ class TestValidateTeamModelMaxBudget: with pytest.raises(ProxyException) as exc_info: self._validate({"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}) assert exc_info.value.code == "400" + + @pytest.mark.parametrize("bad_limit", [0.0, -5.0, float("inf")]) + def test_non_positive_cap_rejected(self, bad_limit): + """Enforcement skips non-positive caps, so accepting one silently uncaps the + model; a team admin could exploit that by lowering an imposed cap to zero + (Veria finding). Validation now rejects the value outright.""" + from litellm.proxy._types import ProxyException + + with pytest.raises(ProxyException) as exc_info: + self._validate({"gpt-4": {"budget_limit": bad_limit, "time_period": "1d"}}) + assert exc_info.value.code == "400" + assert "gpt-4" in exc_info.value.message + assert "budget_limit" in exc_info.value.message