tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 12:23:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

init changes to add PKCE for OKTA

Browse Source
This commit is contained in:
Ishaan Jaffer
2025-10-16 12:13:54 -07:00
parent 71b9becfdd
commit da8f88beba
2 changed files with 240 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/model_prices_and_context_window_backup.json
+98 -2
View File
@@ -9604,6 +9604,54 @@
"supports_vision": true,
"supports_web_search": true
},
"gemini-2.5-flash-image": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-language-models",
"max_audio_length_hours": 8.4,
"max_audio_per_prompt": 1,
"max_images_per_prompt": 3000,
"max_input_tokens": 32768,
"max_output_tokens": 32768,
"max_tokens": 32768,
"max_pdf_size_mb": 30,
"max_video_length": 1,
"max_videos_per_prompt": 10,
"mode": "image_generation",
"output_cost_per_image": 0.039,
"output_cost_per_reasoning_token": 2.5e-06,
"output_cost_per_token": 2.5e-06,
"rpm": 100000,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_audio_output": false,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
"supports_vision": true,
"supports_web_search": true,
"tpm": 8000000
},
"gemini-2.5-flash-image-preview": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 1e-06,
@@ -11057,6 +11105,54 @@
"supports_web_search": true,
"tpm": 8000000
},
"gemini/gemini-2.5-flash-image": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-language-models",
"max_audio_length_hours": 8.4,
"max_audio_per_prompt": 1,
"max_images_per_prompt": 3000,
"max_input_tokens": 32768,
"max_output_tokens": 32768,
"max_tokens": 32768,
"max_pdf_size_mb": 30,
"max_video_length": 1,
"max_videos_per_prompt": 10,
"mode": "image_generation",
"output_cost_per_image": 0.039,
"output_cost_per_reasoning_token": 2.5e-06,
"output_cost_per_token": 2.5e-06,
"rpm": 100000,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_audio_output": false,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
"supports_vision": true,
"supports_web_search": true,
"tpm": 8000000
},
"gemini/gemini-2.5-flash-image-preview": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 1e-06,
@@ -13315,11 +13411,11 @@
"text"
],
"supports_function_calling": true,
"supports_native_streaming": false,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": false,
"supports_tool_choice": true,
litellm/proxy/management_endpoints/ui_sso.py
+142 -6
View File
@@ -9,7 +9,10 @@ Has all /sso/* routes
"""
import asyncio
import base64
import hashlib
import os
import secrets
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
@@ -381,7 +384,10 @@ async def get_generic_sso_response(
try:
result = await generic_sso.verify_and_process(
request,
params={"include_client_id": generic_include_client_id},
params=SSOAuthenticationHandler.prepare_token_exchange_parameters(
request=request,
generic_include_client_id=generic_include_client_id,
),
headers=additional_generic_sso_headers_dict,
)
@@ -1072,14 +1078,66 @@ 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 = (
redirect_params, code_verifier = (
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
# Separate PKCE params from state params (fastapi-sso doesn't accept code_challenge)
pkce_params = {}
state_only_params = {}
for key, value in redirect_params.items():
if key in ("code_challenge", "code_challenge_method"):
pkce_params[key] = value
else:
state_only_params[key] = value
# Get the redirect response from fastapi-sso with only state param
redirect_response = await generic_sso.get_login_redirect(**state_only_params) # type: ignore
# If PKCE is enabled, add PKCE parameters to the redirect URL
if code_verifier and "state" in redirect_params:
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
from litellm.proxy.proxy_server import user_api_key_cache
# Store code_verifier in cache (10 min TTL)
cache_key = f"pkce_verifier:{redirect_params['state']}"
user_api_key_cache.set_cache(
key=cache_key,
value=code_verifier,
ttl=600,
)
# Add PKCE parameters to the authorization URL
if pkce_params:
parsed_url = urlparse(str(redirect_response.headers["location"]))
query_params = parse_qs(parsed_url.query)
# Add PKCE parameters
for key, value in pkce_params.items():
query_params[key] = [value]
# Reconstruct the URL with PKCE parameters
new_query = urlencode(query_params, doseq=True)
new_url = urlunparse((
parsed_url.scheme,
parsed_url.netloc,
parsed_url.path,
parsed_url.params,
new_query,
parsed_url.fragment
))
# Update the redirect response
redirect_response.headers["location"] = new_url
verbose_proxy_logger.debug(
"PKCE parameters added to authorization URL"
)
return redirect_response
raise ValueError(
"Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso"
)
@@ -1088,9 +1146,10 @@ class SSOAuthenticationHandler:
def _get_generic_sso_redirect_params(
state: Optional[str] = None,
generic_authorization_endpoint: Optional[str] = None,
) -> dict:
) -> Tuple[dict, Optional[str]]:
"""
Get redirect parameters for Generic SSO with proper state priority handling.
Optionally generates PKCE parameters if GENERIC_CLIENT_USE_PKCE is enabled.
Priority order:
1. CLI state (if provided)
@@ -1102,9 +1161,12 @@ class SSOAuthenticationHandler:
generic_authorization_endpoint: Authorization endpoint URL
Returns:
dict: Redirect parameters for SSO login
Tuple[dict, Optional[str]]:
- Redirect parameters for SSO login (may include PKCE params)
- code_verifier (if PKCE is enabled, None otherwise)
"""
redirect_params = {}
code_verifier: Optional[str] = None
if state:
# CLI state takes priority
@@ -1122,7 +1184,18 @@ class SSOAuthenticationHandler:
"state"
] = uuid.uuid4().hex # set state param for okta - required
return redirect_params
# Handle PKCE (Proof Key for Code Exchange) if enabled
# Set GENERIC_CLIENT_USE_PKCE=true to enable PKCE for enhanced OAuth security
use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"
if use_pkce:
code_verifier, code_challenge = SSOAuthenticationHandler.generate_pkce_params()
redirect_params["code_challenge"] = code_challenge
redirect_params["code_challenge_method"] = "S256"
verbose_proxy_logger.debug(
"PKCE enabled - code_challenge added to authorization request"
)
return redirect_params, code_verifier
@staticmethod
def should_use_sso_handler(
@@ -1606,6 +1679,69 @@ class SSOAuthenticationHandler:
redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token)
return redirect_response
@staticmethod
def prepare_token_exchange_parameters(
request: Request,
generic_include_client_id: bool,
) -> dict:
"""
Prepare token exchange parameters for Generic SSO.
Args:
request: Request object
generic_include_client_id: Generic OAuth Client ID
Returns:
dict: Token exchange parameters
"""
# Prepare token exchange parameters
token_params = {"include_client_id": generic_include_client_id}
# Retrieve PKCE code_verifier if PKCE was used in authorization
query_params = dict(request.query_params)
state = query_params.get("state")
if state:
from litellm.proxy.proxy_server import user_api_key_cache
cache_key = f"pkce_verifier:{state}"
code_verifier = user_api_key_cache.get_cache(key=cache_key)
if code_verifier:
# Add code_verifier to token exchange parameters
token_params["code_verifier"] = code_verifier
verbose_proxy_logger.debug(
"PKCE code_verifier retrieved and will be included in token exchange"
)
# Clean up the cache entry (single-use verifier)
user_api_key_cache.delete_cache(key=cache_key)
return token_params
@staticmethod
def generate_pkce_params() -> Tuple[str, str]:
"""
Generate PKCE (Proof Key for Code Exchange) parameters for OAuth 2.0.
Returns:
Tuple[str, str]: (code_verifier, code_challenge)
- code_verifier: Random 43-128 character string (we use 43 for efficiency)
- code_challenge: Base64-URL-encoded SHA256 hash of the code_verifier
Reference: https://datatracker.ietf.org/doc/html/rfc7636
"""
# Generate a cryptographically random code_verifier (43 characters)
# Using 32 random bytes which becomes 43 characters when base64-url-encoded
code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode('utf-8').rstrip('=')
# Generate code_challenge using S256 method (SHA256)
code_challenge_bytes = hashlib.sha256(code_verifier.encode('utf-8')).digest()
code_challenge = base64.urlsafe_b64encode(code_challenge_bytes).decode('utf-8').rstrip('=')
return code_verifier, code_challenge
class MicrosoftSSOHandler: