From c59d2a0e2ba34d4faa1d09cb829e7f1242d68047 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 13:10:48 -0700 Subject: [PATCH] fix: fix general linting errors --- .../hooks/parallel_request_limiter_v3.py | 64 ++++++---- litellm/proxy/management_endpoints/ui_sso.py | 115 +++++++++++------- .../team_member_permission_checks.py | 2 +- 3 files changed, 110 insertions(+), 71 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 3e4cbec29f5..86d61b7c4df 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -303,18 +303,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if "{" in key and "}" in key: start = key.find("{") end = key.find("}", start) - hash_tag = key[start:end+1] + hash_tag = key[start : end + 1] else: # Fallback for keys without hash tags hash_tag = "no_hash_tag" - + if hash_tag not in groups: groups[hash_tag] = [] groups[hash_tag].append(key) - + return groups - async def _execute_redis_batch_rate_limiter_script( self, keys_to_fetch: List[str], @@ -332,10 +331,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if self.batch_rate_limiter_script is None: return [] - + key_groups = self._group_keys_by_hash_tag(keys_to_fetch) all_cache_values = [] - + for hash_tag, group_keys in key_groups.items(): try: group_cache_values = await self.batch_rate_limiter_script( @@ -354,7 +353,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_size=self.window_size, ) all_cache_values.extend(group_cache_values) - + return all_cache_values async def should_rate_limit( @@ -378,7 +377,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for descriptor in descriptors: descriptor_key = descriptor["key"] descriptor_value = descriptor["value"] - rate_limit: Optional[RateLimitDescriptorRateLimitObject] = descriptor.get("rate_limit", {}) or {} + rate_limit: Optional[RateLimitDescriptorRateLimitObject] = ( + descriptor.get("rate_limit") or RateLimitDescriptorRateLimitObject() + ) requests_limit = rate_limit.get("requests_per_unit") tokens_limit = rate_limit.get("tokens_per_unit") max_parallel_requests_limit = rate_limit.get("max_parallel_requests") @@ -632,26 +633,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for i, status in enumerate(response["statuses"]): if status["code"] == "OVER_LIMIT": descriptor = descriptors[floor(i / 2)] - + # Calculate reset time (window_start + window_size) now = datetime.now().timestamp() reset_time = now + self.window_size # Conservative estimate - reset_time_formatted = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") - + reset_time_formatted = datetime.fromtimestamp( + reset_time + ).strftime("%Y-%m-%d %H:%M:%S UTC") + # Handle negative remaining values more gracefully - remaining_display = max(0, status['limit_remaining']) - + remaining_display = max(0, status["limit_remaining"]) + # Create detailed error message - rate_limit_type = status['rate_limit_type'] - current_limit = status['current_limit'] - + rate_limit_type = status["rate_limit_type"] + current_limit = status["current_limit"] + detail = ( f"Rate limit exceeded for {descriptor['key']}: {descriptor['value']}. " f"Limit type: {rate_limit_type}. " f"Current limit: {current_limit}, Remaining: {remaining_display}. " f"Limit resets at: {reset_time_formatted}" ) - + raise HTTPException( status_code=429, detail=detail, @@ -693,7 +696,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return pipeline_operations - + async def _execute_token_increment_script( self, pipeline_operations: List["RedisPipelineIncrementOperation"], @@ -703,15 +706,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if self.token_increment_script is None: return - + # Group operations by hash tag for Redis cluster compatibility operation_keys = [op["key"] for op in pipeline_operations] key_groups = self._group_keys_by_hash_tag(operation_keys) - + for _hash_tag, group_keys in key_groups.items(): # Get operations for this hash tag group - group_operations = [op for op in pipeline_operations if op["key"] in group_keys] - + group_operations = [ + op for op in pipeline_operations if op["key"] in group_keys + ] + keys = [] args = [] @@ -731,7 +736,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): args=args, ) - async def async_increment_tokens_with_ttl_preservation( self, pipeline_operations: List["RedisPipelineIncrementOperation"], @@ -757,7 +761,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: await self._execute_token_increment_script(pipeline_operations) - + verbose_proxy_logger.debug( f"Successfully executed TTL-preserving increment for {len(pipeline_operations)} keys" ) @@ -811,7 +815,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # Get metadata from kwargs - litellm_metadata = kwargs["litellm_params"].get(get_metadata_variable_name_from_kwargs(kwargs), {}) + litellm_metadata = kwargs["litellm_params"].get( + get_metadata_variable_name_from_kwargs(kwargs), {} + ) if litellm_metadata is None: return user_api_key = litellm_metadata.get("user_api_key") @@ -825,7 +831,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get total tokens from response total_tokens = 0 # spot fix for /responses api - if (isinstance(response_obj, ModelResponse) or isinstance(response_obj, BaseLiteLLMOpenAIResponseObject)): + if isinstance(response_obj, ModelResponse) or isinstance( + response_obj, BaseLiteLLMOpenAIResponseObject + ): _usage = getattr(response_obj, "usage", None) if _usage and isinstance(_usage, Usage): if rate_limit_type == "output": @@ -943,7 +951,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): _get_parent_otel_span_from_kwargs(kwargs) ) litellm_metadata = kwargs["litellm_params"]["metadata"] - user_api_key = litellm_metadata.get("user_api_key") if litellm_metadata else None + user_api_key = ( + litellm_metadata.get("user_api_key") if litellm_metadata else None + ) pipeline_operations: List[RedisPipelineIncrementOperation] = [] if user_api_key: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index dc8b96656c4..59227aa5e76 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -10,7 +10,6 @@ Has all /sso/* routes import asyncio import os -from litellm._uuid import uuid from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast @@ -19,6 +18,7 @@ from fastapi.responses import RedirectResponse import litellm from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.caching import DualCache from litellm.constants import MAX_SPENDLOG_ROWS_TO_QUERY from litellm.llms.custom_httpx.http_handler import ( @@ -115,7 +115,10 @@ def process_sso_jwt_access_token( @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) async def google_login( - request: Request, source: Optional[str] = None, key: Optional[str] = None, existing_key: Optional[str] = None + request: Request, + source: Optional[str] = None, + key: Optional[str] = None, + existing_key: Optional[str] = None, ): # noqa: PLR0915 """ Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env @@ -664,17 +667,20 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: status_code=401, detail="Result not returned by SSO provider.", ) - if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): # Extract the key ID from the state key_id = state.split(":", 1)[1] - + # Get existing_key from query parameters if provided existing_key = request.query_params.get("existing_key") - - verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}") - return await cli_sso_callback(request=request, key=key_id, existing_key=existing_key, result=result) + + verbose_proxy_logger.info( + f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}" + ) + return await cli_sso_callback( + request=request, key=key_id, existing_key=existing_key, result=result + ) return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, @@ -685,30 +691,30 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) -async def _regenerate_cli_key(existing_key: str, new_key: str, user_id: Optional[str] = None) -> None: +async def _regenerate_cli_key( + existing_key: str, new_key: str, user_id: Optional[str] = None +) -> None: """Regenerate an existing CLI key with a new token""" from litellm.proxy._types import RegenerateKeyRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.key_management_endpoints import ( regenerate_key_fn, ) - + verbose_proxy_logger.info(f"Regenerating existing CLI key: {existing_key}") - + admin_user_dict = UserAPIKeyAuth.get_litellm_cli_user_api_key_auth() - + regenerate_request = RegenerateKeyRequest( key=existing_key, new_key=new_key, duration="24hr", user_id=user_id, ) - + await regenerate_key_fn( - key=existing_key, - data=regenerate_request, - user_api_key_dict=admin_user_dict + key=existing_key, data=regenerate_request, user_api_key_dict=admin_user_dict ) - + verbose_proxy_logger.info(f"Regenerated CLI key: {new_key}") @@ -720,9 +726,9 @@ async def _create_new_cli_key( from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) - + verbose_proxy_logger.info("Creating new CLI key") - + await generate_key_helper_fn( request_type="key", duration="24hr", @@ -734,13 +740,20 @@ async def _create_new_cli_key( table_name="key", token=key, ) - + verbose_proxy_logger.info(f"Created new CLI key: {key}") -async def cli_sso_callback(request: Request, key: Optional[str] = None, existing_key: Optional[str] = None, result: Optional[Union[OpenID, dict]] = None): +async def cli_sso_callback( + request: Request, + key: Optional[str] = None, + existing_key: Optional[str] = None, + result: Optional[Union[OpenID, dict]] = None, +): """CLI SSO callback - regenerates existing CLI key or creates new one""" - verbose_proxy_logger.info(f"CLI SSO callback for key: {key}, existing_key: {existing_key}") + verbose_proxy_logger.info( + f"CLI SSO callback for key: {key}, existing_key: {existing_key}" + ) from litellm.proxy.proxy_server import prisma_client @@ -754,8 +767,10 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None, existing raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - - parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result(result=result) + + parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( + result=result + ) verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}") try: @@ -783,7 +798,9 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None, existing except Exception as e: verbose_proxy_logger.error(f"Error with CLI key: {e}") - raise HTTPException(status_code=500, detail=f"Failed to process CLI key: {str(e)}") + raise HTTPException( + status_code=500, detail=f"Failed to process CLI key: {str(e)}" + ) @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) @@ -874,7 +891,7 @@ async def insert_sso_user( auto_create_key=False, ) - if result_openid: + if result_openid and isinstance(result_openid, OpenID): new_user_request.metadata = {"auth_provider": result_openid.provider} response = await new_user( @@ -1052,11 +1069,13 @@ class SSOAuthenticationHandler: # or a cryptographicly signed state that we can verify stateless # For simplification we are using a static state, this is not perfect but some # SSO providers do not allow stateless verification - redirect_params = SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=state, - generic_authorization_endpoint=generic_authorization_endpoint + redirect_params = ( + SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=state, + generic_authorization_endpoint=generic_authorization_endpoint, + ) ) - + return await generic_sso.get_login_redirect(**redirect_params) # type: ignore raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" @@ -1064,26 +1083,26 @@ class SSOAuthenticationHandler: @staticmethod def _get_generic_sso_redirect_params( - state: Optional[str] = None, - generic_authorization_endpoint: Optional[str] = None + state: Optional[str] = None, + generic_authorization_endpoint: Optional[str] = None, ) -> dict: """ Get redirect parameters for Generic SSO with proper state priority handling. - + Priority order: 1. CLI state (if provided) 2. GENERIC_CLIENT_STATE environment variable 3. Generated UUID for Okta (if Okta endpoint detected) - + Args: state: Optional state parameter (e.g., CLI state) generic_authorization_endpoint: Authorization endpoint URL - + Returns: dict: Redirect parameters for SSO login """ redirect_params = {} - + if state: # CLI state takes priority # the litellm proxy cli sends the "state" parameter to the proxy server for auth. We should maintain the state parameter for the cli if it is provided @@ -1092,8 +1111,13 @@ class SSOAuthenticationHandler: generic_client_state = os.getenv("GENERIC_CLIENT_STATE", None) if generic_client_state: redirect_params["state"] = generic_client_state - elif generic_authorization_endpoint and "okta" in generic_authorization_endpoint: - redirect_params["state"] = uuid.uuid4().hex # set state param for okta - required + elif ( + generic_authorization_endpoint + and "okta" in generic_authorization_endpoint + ): + redirect_params["state"] = ( + uuid.uuid4().hex + ) # set state param for okta - required return redirect_params @@ -1127,11 +1151,11 @@ class SSOAuthenticationHandler: redirect_url += sso_callback_route else: redirect_url += "/" + sso_callback_route - + # Append existing_key as query parameter if provided if existing_key: redirect_url += f"?existing_key={existing_key}" - + return redirect_url @staticmethod @@ -1314,7 +1338,9 @@ class SSOAuthenticationHandler: return team_request @staticmethod - def _get_cli_state(source: Optional[str], key: Optional[str], existing_key: Optional[str] = None) -> Optional[str]: + def _get_cli_state( + source: Optional[str], key: Optional[str], existing_key: Optional[str] = None + ) -> Optional[str]: """ Checks the request 'source' if a cli state token was passed in @@ -1374,7 +1400,7 @@ class SSOAuthenticationHandler: if user_email is not None and (user_id is None or len(user_id) == 0): user_id = user_email - + return ParsedOpenIDResult( user_email=user_email, user_id=user_id, @@ -1408,13 +1434,16 @@ class SSOAuthenticationHandler: ) # User is Authe'd in - generate key for the UI to access Proxy - parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result(result=result, generic_client_id=generic_client_id) + parsed_openid_result = ( + SSOAuthenticationHandler._get_user_email_and_id_from_result( + result=result, generic_client_id=generic_client_id + ) + ) user_email = parsed_openid_result.get("user_email") user_id = parsed_openid_result.get("user_id") user_role = parsed_openid_result.get("user_role") verbose_proxy_logger.info(f"SSO callback result: {result}") - user_info = None user_id_models: List = [] max_internal_user_budget = litellm.max_internal_user_budget diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 3b67472ed23..4642028b77c 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -171,7 +171,7 @@ class TeamMemberPermissionChecks: """ all_available_permissions = [] for route in LiteLLMRoutes.key_management_routes.value: - all_available_permissions.append(route.value) + all_available_permissions.append(route) return all_available_permissions @staticmethod