mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix: fix general linting errors
This commit is contained in:
Krrish Dholakia
parent
a69746735c
commit
c59d2a0e2b
3 changed files with 110 additions and 71 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue