Added LiteLLM to the stack
This commit is contained in:
@@ -0,0 +1,10 @@
|
||||
"""
|
||||
Enterprise Authentication Module for LiteLLM Proxy
|
||||
|
||||
This module contains enterprise-specific authentication functionality,
|
||||
including custom SSO handlers and advanced authentication features.
|
||||
"""
|
||||
|
||||
from .custom_sso_handler import EnterpriseCustomSSOHandler
|
||||
|
||||
__all__ = ["EnterpriseCustomSSOHandler"]
|
||||
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
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 TYPE_CHECKING, Dict, Optional, Union, cast
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.responses import RedirectResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_sso.sso.base import OpenID
|
||||
else:
|
||||
from typing import Any as OpenID
|
||||
|
||||
from litellm.proxy.management_endpoints.types import CustomOpenID
|
||||
|
||||
|
||||
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,
|
||||
premium_user,
|
||||
user_custom_ui_sso_sign_in_handler,
|
||||
)
|
||||
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.")
|
||||
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
|
||||
class EnterpriseRouteChecks:
|
||||
@staticmethod
|
||||
def is_llm_api_route_disabled() -> bool:
|
||||
"""
|
||||
Check if llm api route is disabled
|
||||
"""
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
## Check if DISABLE_LLM_API_ENDPOINTS is set
|
||||
if "DISABLE_LLM_API_ENDPOINTS" in os.environ:
|
||||
if not premium_user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"🚨🚨🚨 DISABLING LLM API ENDPOINTS is an Enterprise feature\n🚨 {CommonProxyErrors.not_premium_user.value}",
|
||||
)
|
||||
|
||||
return get_secret_bool("DISABLE_LLM_API_ENDPOINTS") is True
|
||||
|
||||
@staticmethod
|
||||
def is_management_routes_disabled() -> bool:
|
||||
"""
|
||||
Check if management route is disabled
|
||||
"""
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
if "DISABLE_ADMIN_ENDPOINTS" in os.environ:
|
||||
if not premium_user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"🚨🚨🚨 DISABLING LLM API ENDPOINTS is an Enterprise feature\n🚨 {CommonProxyErrors.not_premium_user.value}",
|
||||
)
|
||||
|
||||
return get_secret_bool("DISABLE_ADMIN_ENDPOINTS") is True
|
||||
|
||||
@staticmethod
|
||||
def should_call_route(route: str):
|
||||
"""
|
||||
Check if management route is disabled and raise exception
|
||||
"""
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
||||
if (
|
||||
RouteChecks.is_management_route(route=route)
|
||||
and EnterpriseRouteChecks.is_management_routes_disabled()
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Management routes are disabled for this instance.",
|
||||
)
|
||||
elif (
|
||||
RouteChecks.is_llm_api_route(route=route)
|
||||
and EnterpriseRouteChecks.is_llm_api_route_disabled()
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="LLM API routes are disabled for this instance.",
|
||||
)
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
|
||||
|
||||
async def enterprise_custom_auth(
|
||||
request: Request, api_key: str, user_custom_auth: Optional[Any]
|
||||
) -> Optional[UserAPIKeyAuth]:
|
||||
from litellm_enterprise.proxy.proxy_server import custom_auth_settings
|
||||
|
||||
if user_custom_auth is None:
|
||||
return None
|
||||
|
||||
if custom_auth_settings is None:
|
||||
return await user_custom_auth(request, api_key)
|
||||
|
||||
if custom_auth_settings["mode"] == "on":
|
||||
return await user_custom_auth(request, api_key)
|
||||
elif custom_auth_settings["mode"] == "off":
|
||||
return None
|
||||
elif custom_auth_settings["mode"] == "auto":
|
||||
try:
|
||||
return await user_custom_auth(request, api_key)
|
||||
except ProxyException as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error in custom auth, checking litellm auth: {e}"
|
||||
)
|
||||
return None
|
||||
else:
|
||||
raise ValueError(f"Invalid mode: {custom_auth_settings['mode']}")
|
||||
Reference in New Issue
Block a user