phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat(proxy_server.py): adds ui_access_mode to control access to proxy ui

Browse source 
allows admins to control who can access proxy UI - even when SSO is turned on
This commit is contained in:
Krrish Dholakia 2024-02-27 17:46:44 -08:00
parent 411787963b
commit b5af282820
2 changed files with 24 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -402,6 +402,9 @@ class ConfigGeneralSettings(LiteLLMBase):
None, None,
description="sends alerts if requests hang for 5min+", description="sends alerts if requests hang for 5min+",
) )
ui_access_mode: Optional[Literal["admin_only", "all"]] = Field(
"all", description="Control access to the Proxy UI"
)
class ConfigYAML(LiteLLMBase): class ConfigYAML(LiteLLMBase):

23
litellm/proxy/proxy_server.py
View file

@ -239,6 +239,7 @@ health_check_interval = None
health_check_results = {} health_check_results = {}
queue: List = [] queue: List = []
litellm_proxy_budget_name = "litellm-proxy-budget" litellm_proxy_budget_name = "litellm-proxy-budget"
ui_access_mode: Literal["admin", "all"] = "all"
### INITIALIZE GLOBAL LOGGING OBJECT ### ### INITIALIZE GLOBAL LOGGING OBJECT ###
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
### REDIS QUEUE ### ### REDIS QUEUE ###
@ -1406,7 +1407,7 @@ class ProxyConfig:
""" """
Load config values into proxy global state Load config values into proxy global state
""" """
global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, ui_access_mode
# Load existing config # Load existing config
config = await self.get_config(config_file_path=config_file_path) config = await self.get_config(config_file_path=config_file_path)
@ -1713,6 +1714,10 @@ class ProxyConfig:
) )
## COST TRACKING ## ## COST TRACKING ##
cost_tracking() cost_tracking()
## ADMIN UI ACCESS ##
ui_access_mode = general_settings.get(
"ui_access_mode", "all"
) # can be either ["admin_only" or "all"]
### BACKGROUND HEALTH CHECKS ### ### BACKGROUND HEALTH CHECKS ###
# Enable background health checks # Enable background health checks
use_background_health_checks = general_settings.get( use_background_health_checks = general_settings.get(
@ -5620,7 +5625,7 @@ def get_image():
@app.get("/sso/callback", tags=["experimental"]) @app.get("/sso/callback", tags=["experimental"])
async def auth_callback(request: Request): async def auth_callback(request: Request):
"""Verify login""" """Verify login"""
global general_settings global general_settings, ui_access_mode
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
@ -5851,6 +5856,20 @@ async def auth_callback(request: Request):
): ):
# checks if user is admin # checks if user is admin
user_role = "app_admin" user_role = "app_admin"
verbose_proxy_logger.debug(
f"user_role: {user_role}; ui_access_mode: {ui_access_mode}"
)
## CHECK IF ROLE ALLOWED TO USE PROXY ##
if ui_access_mode == "admin_only" and "admin" not in user_role:
verbose_proxy_logger.debug("EXCEPTION RAISED")
raise HTTPException(
status_code=401,
detail={
"error": f"User not allowed to access proxy. User role={user_role}, proxy mode={ui_access_mode}"
},
)
import jwt import jwt
jwt_token = jwt.encode( jwt_token = jwt.encode(