diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a741869e5f..5e0a211906 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -202,21 +202,29 @@ async def common_checks( and general_settings["enforce_user_param"] is True ): # Get HTTP method from request - http_method = request.method if hasattr(request, 'method') else None - + http_method = request.method if hasattr(request, "method") else None + # Check if it's a POST request and if it's an OpenAI route but not MCP is_post_method = http_method and http_method.upper() == "POST" is_openai_route = RouteChecks.is_llm_api_route(route=route) - is_mcp_route = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + is_mcp_route = ( + route in LiteLLMRoutes.mcp_routes.value + or RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) ) - + # Enforce user param only for POST requests on OpenAI routes (excluding MCP routes) - if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body: + if ( + is_post_method + and is_openai_route + and not is_mcp_route + and "user" not in request_body + ): raise Exception( f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" ) - + # 6.1 [OPTIONAL] If 'reject_clientside_metadata_tags' enabled - reject request if it has client-side 'metadata.tags' if ( general_settings.get("reject_clientside_metadata_tags", None) is not None @@ -502,53 +510,51 @@ async def get_default_end_user_budget( ) -> Optional[LiteLLM_BudgetTable]: """ Fetches the default end user budget from the database if litellm.max_end_user_budget_id is configured. - + This budget is applied to end users who don't have an explicit budget_id set. Results are cached for performance. - + Args: prisma_client: Database client instance user_api_key_cache: Cache for storing/retrieving budget data parent_otel_span: Optional OpenTelemetry span for tracing - + Returns: LiteLLM_BudgetTable if configured and found, None otherwise """ if prisma_client is None or litellm.max_end_user_budget_id is None: return None - + cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}" - + # Check cache first cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) if cached_budget is not None: return LiteLLM_BudgetTable(**cached_budget) - + # Fetch from database try: budget_record = await prisma_client.db.litellm_budgettable.find_unique( where={"budget_id": litellm.max_end_user_budget_id} ) - + if budget_record is None: verbose_proxy_logger.warning( f"Default end user budget not found in database: {litellm.max_end_user_budget_id}" ) return None - + # Cache the budget for 60 seconds await user_api_key_cache.async_set_cache( - key=cache_key, + key=cache_key, value=budget_record.dict(), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - + return LiteLLM_BudgetTable(**budget_record.dict()) - + except Exception as e: - verbose_proxy_logger.error( - f"Error fetching default end user budget: {str(e)}" - ) + verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}") return None @@ -560,38 +566,38 @@ async def _apply_default_budget_to_end_user( ) -> LiteLLM_EndUserTable: """ Helper function to apply default budget to end user if they don't have a budget assigned. - + Args: end_user_obj: The end user object to potentially apply default budget to prisma_client: Database client instance user_api_key_cache: Cache for storing/retrieving data parent_otel_span: Optional OpenTelemetry span for tracing - + Returns: Updated end user object with default budget applied if applicable """ # If end user already has a budget assigned, no need to apply default if end_user_obj.litellm_budget_table is not None: return end_user_obj - + # If no default budget configured, return as-is if litellm.max_end_user_budget_id is None: return end_user_obj - + # Fetch and apply default budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) - + if default_budget is not None: # Apply default budget to end user object end_user_obj.litellm_budget_table = default_budget verbose_proxy_logger.debug( f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" ) - + return end_user_obj @@ -601,20 +607,20 @@ def _check_end_user_budget( ) -> None: """ Check if end user is within their budget limit. - + Args: end_user_obj: The end user object to check route: The request route - + Raises: litellm.BudgetExceededError: If end user has exceeded their budget """ if route in LiteLLMRoutes.info_routes.value: return - + if end_user_obj.litellm_budget_table is None: return - + end_user_budget = end_user_obj.litellm_budget_table.max_budget if end_user_budget is not None and end_user_obj.spend > end_user_budget: raise litellm.BudgetExceededError( @@ -635,8 +641,8 @@ async def get_end_user_object( ) -> Optional[LiteLLM_EndUserTable]: """ Returns end user object from database or cache. - - If end user exists but has no budget_id, applies the default budget + + If end user exists but has no budget_id, applies the default budget (if configured via litellm.max_end_user_budget_id). Args: @@ -646,7 +652,7 @@ async def get_end_user_object( route: The request route parent_otel_span: Optional OpenTelemetry span for tracing proxy_logging_obj: Optional proxy logging object - + Returns: LiteLLM_EndUserTable if found, None otherwise """ @@ -655,14 +661,14 @@ async def get_end_user_object( if end_user_id is None: return None - + _key = "end_user_id:{}".format(end_user_id) # Check cache first cached_user_obj = await user_api_key_cache.async_get_cache(key=_key) if cached_user_obj is not None: return_obj = LiteLLM_EndUserTable(**cached_user_obj) - + # Apply default budget if needed return_obj = await _apply_default_budget_to_end_user( end_user_obj=return_obj, @@ -670,10 +676,10 @@ async def get_end_user_object( user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) - + # Check budget limits _check_end_user_budget(end_user_obj=return_obj, route=route) - + return return_obj # Fetch from database @@ -688,7 +694,7 @@ async def get_end_user_object( # Convert to LiteLLM_EndUserTable object _response = LiteLLM_EndUserTable(**response.dict()) - + # Apply default budget if needed _response = await _apply_default_budget_to_end_user( end_user_obj=_response, @@ -696,18 +702,17 @@ async def get_end_user_object( user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) - + # Save to cache (always store as dict for consistency) await user_api_key_cache.async_set_cache( - key="end_user_id:{}".format(end_user_id), - value=_response.dict() + key="end_user_id:{}".format(end_user_id), value=_response.dict() ) - + # Check budget limits _check_end_user_budget(end_user_obj=_response, route=route) return _response - + except Exception as e: if isinstance(e, litellm.BudgetExceededError): raise e @@ -747,7 +752,6 @@ async def get_tag_objects_batch( tag_objects = {} uncached_tags = [] - # Try to get all tags from cache first for tag_name in tag_names: @@ -1138,7 +1142,6 @@ async def _cache_management_object( user_api_key_cache: DualCache, proxy_logging_obj: Optional[ProxyLogging], ): - await user_api_key_cache.async_set_cache( key=key, value=value, @@ -1459,9 +1462,7 @@ async def get_team_object_by_alias( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - "Error looking up team by alias: %s", team_alias - ) + verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias) raise HTTPException( status_code=500, detail={ @@ -1602,11 +1603,11 @@ class ExperimentalUIJWTToken: ) -> str: """ Generate a JWT token for CLI authentication with 24-hour expiration. - + Args: user_info: User information from the database team_id: Team ID for the user (optional, uses user's team if available) - + Returns: Encrypted JWT token string """ @@ -1800,7 +1801,7 @@ async def get_org_object( - Check if org id in proxy Org Table - if valid, return LiteLLM_OrganizationTable object - if not, then raise an error - + Args: org_id: Organization ID to look up prisma_client: Database client @@ -1820,7 +1821,7 @@ async def get_org_object( cache_key = "org_id:{}".format(org_id) if include_budget_table: cache_key = "org_id:{}:with_budget".format(org_id) - + # check if in cache cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key) if cached_org_obj is not None: @@ -1833,7 +1834,7 @@ async def get_org_object( query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - + response = await prisma_client.db.litellm_organizationtable.find_unique( **query_kwargs ) @@ -1844,7 +1845,9 @@ async def get_org_object( # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=response.model_dump() if hasattr(response, "model_dump") else response, + value=response.model_dump() + if hasattr(response, "model_dump") + else response, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2218,10 +2221,15 @@ async def _virtual_key_max_budget_alert_check( and valid_token.spend is not None and valid_token.spend > 0 ): - alert_threshold = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - + alert_threshold = ( + valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + ) + # Only alert if we've crossed the threshold but haven't exceeded max_budget yet - if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget: + if ( + valid_token.spend >= alert_threshold + and valid_token.spend < valid_token.max_budget + ): verbose_proxy_logger.debug( "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", valid_token.token, @@ -2274,7 +2282,7 @@ async def _check_team_member_budget( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - + if ( team_membership is not None and team_membership.litellm_budget_table is not None @@ -2282,8 +2290,8 @@ async def _check_team_member_budget( ): team_member_budget = team_membership.litellm_budget_table.max_budget team_member_spend = team_membership.spend or 0.0 - - if team_member_spend > team_member_budget: + + if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, @@ -2343,11 +2351,11 @@ async def _organization_max_budget_check( ): """ Check if the organization is over its max budget. - + This function checks the organization budget using: 1. First, tries to use valid_token.org_id (if key has organization_id set) 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) - + This ensures organization budget checks work even when keys don't have organization_id set directly, as long as their team belongs to an organization. @@ -2364,7 +2372,7 @@ async def _organization_max_budget_check( org_id = valid_token.org_id elif team_object is not None and team_object.organization_id is not None: org_id = team_object.organization_id - + # If no organization_id found, skip the check if org_id is None: return @@ -2655,4 +2663,4 @@ def _can_object_call_vector_stores( code=status.HTTP_401_UNAUTHORIZED, ) - return True \ No newline at end of file + return True