litellm/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py

96 lines
3.5 KiB
Python

"""
Enterprise Custom SSO Handler for LiteLLM Proxy
This module contains enterprise-specific custom SSO authentication functionality
that allows users to implement their own SSO handling logic by providing custom
handlers that process incoming request headers and return OpenID objects.
Use this when you have an OAuth proxy in front of LiteLLM (where the OAuth proxy
has already authenticated the user) and you need to extract user information from
custom headers or other request attributes.
"""
from typing import cast
from fastapi import Request
from fastapi.responses import RedirectResponse
class EnterpriseCustomSSOHandler:
"""
Enterprise Custom SSO Handler for LiteLLM Proxy
This class provides methods for handling custom SSO authentication flows
where users can implement their own authentication logic by processing
request headers and returning user information in OpenID format.
"""
@staticmethod
async def handle_custom_ui_sso_sign_in(
request: Request,
) -> RedirectResponse:
"""
Allow a user to execute their custom code to parse incoming request headers and return a OpenID object
Use this when you have an OAuth proxy in front of LiteLLM (where the OAuth proxy has already authenticated the user)
Args:
request: The FastAPI request object containing headers and other request data
Returns:
RedirectResponse: Redirect response that sends the user to the LiteLLM UI with authentication token
Raises:
ValueError: If custom_ui_sso_sign_in_handler is not configured
Example:
This method is typically called when a user has already been authenticated by an
external OAuth proxy and the proxy has added custom headers containing user information.
The custom handler extracts this information and converts it to an OpenID object.
"""
from fastapi_sso.sso.base import OpenID
from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler
from litellm.proxy.proxy_server import (
CommonProxyErrors,
general_settings,
premium_user,
user_custom_ui_sso_sign_in_handler,
)
from litellm.proxy.auth.trusted_proxy_utils import (
require_trusted_proxy_request,
)
if premium_user is not True:
raise ValueError(CommonProxyErrors.not_premium_user.value)
if user_custom_ui_sso_sign_in_handler is None:
raise ValueError(
"custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings."
)
require_trusted_proxy_request(
request=request,
general_settings=general_settings,
feature_name="Custom UI SSO",
)
custom_sso_login_handler = cast(
CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler
)
openid_response: OpenID = (
await custom_sso_login_handler.handle_custom_ui_sso_sign_in(
request=request,
)
)
# Import here to avoid circular imports
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=openid_response,
request=request,
received_response=None,
generic_client_id=None,
ui_access_mode=None,
)