From e6018a464ff92f6de68afc1b4f9d8eee53add1da Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 19 Sep 2024 17:13:58 -0700 Subject: [PATCH] [ Proxy - User Management]: If user assigned to a team don't show Default Team (#5791) * rename endpoint to ui_settings * ui allow DEFAULT_TEAM_DISABLED * fix logic * docs Set `default_team_disabled: true` on your litellm config.yaml --- docs/my-website/docs/proxy/ui.md | 16 +++++++++++++++ litellm/proxy/_types.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 2 +- litellm/proxy/proxy_config.yaml | 5 +++++ litellm/proxy/proxy_server.py | 20 ++++++++++++++++--- .../src/components/dashboard_default_team.tsx | 5 +++++ .../src/components/networking.tsx | 4 ++-- .../src/components/user_dashboard.tsx | 14 ++++++++----- 8 files changed, 56 insertions(+), 12 deletions(-) diff --git a/docs/my-website/docs/proxy/ui.md b/docs/my-website/docs/proxy/ui.md index 02d2afb7b..85cb2eb98 100644 --- a/docs/my-website/docs/proxy/ui.md +++ b/docs/my-website/docs/proxy/ui.md @@ -239,6 +239,22 @@ If you don't see all your keys this could be due to a cached token. So just re-l ::: +### Disable `Default Team` on Admin UI + +Use this if you want to hide the Default Team on the Admin UI + +The following logic will apply +- If team assigned don't show `Default Team` +- If no team assigned then they should see `Default Team` + +Set `default_team_disabled: true` on your litellm config.yaml + +```yaml +general_settings: + master_key: sk-1234 + default_team_disabled: true # OR you can set env var PROXY_DEFAULT_TEAM_DISABLED="true" +``` + ### Sign in with Username, Password when SSO is on If you need to access the UI via username/password when SSO is on navigate to `/fallback/login`. This route will allow you to sign in with your username/password credentials. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 57873e94a..ee375a875 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -289,7 +289,7 @@ class LiteLLMRoutes(enum.Enum): "/key/delete", "/global/spend/logs", "/global/predict/spend/logs", - "/sso/get/logout_url", + "/sso/get/ui_settings", ] management_routes = [ # key diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4a62321d9..d49d9843e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1038,7 +1038,7 @@ async def user_api_key_auth( # this token is only used for managing the ui allowed_routes = [ "/sso", - "/sso/get/logout_url", + "/sso/get/ui_settings", "/login", "/key/generate", "/key/update", diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index eaf94974b..1795adf71 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -21,6 +21,11 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app + +general_settings: + master_key: sk-1234 + default_team_disabled: true + litellm_settings: success_callback: ["prometheus"] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 613dacb57..8c6bc618c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -659,6 +659,11 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): # Set your Azure Key Vault URI KVUri = os.getenv("AZURE_KEY_VAULT_URI", None) + if KVUri is None: + raise Exception( + "Error when loading keys from Azure Key Vault: AZURE_KEY_VAULT_URI is not set." + ) + credential = DefaultAzureCredential() # Create the SecretClient using the credential @@ -8342,16 +8347,25 @@ async def login(request: Request): @app.get( - "/sso/get/logout_url", + "/sso/get/ui_settings", tags=["experimental"], include_in_schema=False, dependencies=[Depends(user_api_key_auth)], ) -async def get_logout_url(request: Request): +async def get_ui_settings(request: Request): _proxy_base_url = os.getenv("PROXY_BASE_URL", None) _logout_url = os.getenv("PROXY_LOGOUT_URL", None) - return {"PROXY_BASE_URL": _proxy_base_url, "PROXY_LOGOUT_URL": _logout_url} + default_team_disabled = general_settings.get("default_team_disabled", False) + if "PROXY_DEFAULT_TEAM_DISABLED" in os.environ: + if os.environ["PROXY_DEFAULT_TEAM_DISABLED"].lower() == "true": + default_team_disabled = True + + return { + "PROXY_BASE_URL": _proxy_base_url, + "PROXY_LOGOUT_URL": _logout_url, + "DEFAULT_TEAM_DISABLED": default_team_disabled, + } @app.get("/onboarding/get_token", include_in_schema=False) diff --git a/ui/litellm-dashboard/src/components/dashboard_default_team.tsx b/ui/litellm-dashboard/src/components/dashboard_default_team.tsx index 98f287ed4..7ba67ef3c 100644 --- a/ui/litellm-dashboard/src/components/dashboard_default_team.tsx +++ b/ui/litellm-dashboard/src/components/dashboard_default_team.tsx @@ -1,10 +1,12 @@ import React, { useState, useEffect } from "react"; import { Select, SelectItem, Text, Title } from "@tremor/react"; +import { ProxySettings } from "./user_dashboard"; interface DashboardTeamProps { teams: Object[] | null; setSelectedTeam: React.Dispatch>; userRole: string | null; + proxySettings: ProxySettings | null; } type TeamInterface = { @@ -17,6 +19,7 @@ const DashboardTeam: React.FC = ({ teams, setSelectedTeam, userRole, + proxySettings, }) => { const defaultTeam: TeamInterface = { models: [], @@ -31,6 +34,8 @@ const DashboardTeam: React.FC = ({ if (userRole === "App User") { // Non-Admin SSO users should only see their own team - they should not see "Default Team" updatedTeams = teams; + } else if (proxySettings && proxySettings.DEFAULT_TEAM_DISABLED === true) { + updatedTeams = teams ? [...teams] : [defaultTeam]; } else { updatedTeams = teams ? [...teams, defaultTeam] : [defaultTeam]; } diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 27d109699..635edaf4b 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2752,8 +2752,8 @@ export const getProxyBaseUrlAndLogoutUrl = async ( */ try { let url = proxyBaseUrl - ? `${proxyBaseUrl}/sso/get/logout_url` - : `/sso/get/logout_url`; + ? `${proxyBaseUrl}/sso/get/ui_settings` + : `sso/get/ui_settings`; //message.info("Requesting model data"); const response = await fetch(url, { diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 712e42676..a1866b6b8 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -28,6 +28,12 @@ type UserSpendData = { max_budget?: number | null; }; +export interface ProxySettings { + PROXY_BASE_URL: string | null; + PROXY_LOGOUT_URL: string | null; + DEFAULT_TEAM_DISABLED: boolean; +} + function getCookie(name: string) { console.log("COOKIES", document.cookie) const cookieValue = document.cookie @@ -46,8 +52,6 @@ interface UserDashboardProps { setUserEmail: React.Dispatch>; setTeams: React.Dispatch>; setKeys: React.Dispatch>; - setProxySettings: React.Dispatch>; - proxySettings: any; premiumUser: boolean; } @@ -67,8 +71,6 @@ const UserDashboard: React.FC = ({ setUserEmail, setTeams, setKeys, - setProxySettings, - proxySettings, premiumUser, }) => { const [userSpendData, setUserSpendData] = useState( @@ -87,6 +89,7 @@ const UserDashboard: React.FC = ({ const [accessToken, setAccessToken] = useState(null); const [teamSpend, setTeamSpend] = useState(null); const [userModels, setUserModels] = useState([]); + const [proxySettings, setProxySettings] = useState(null); const defaultTeam: TeamInterface = { models: [], team_alias: "Default Team", @@ -166,7 +169,7 @@ const UserDashboard: React.FC = ({ } else { const fetchData = async () => { try { - const proxy_settings = await getProxyBaseUrlAndLogoutUrl(accessToken); + const proxy_settings: ProxySettings = await getProxyBaseUrlAndLogoutUrl(accessToken); setProxySettings(proxy_settings); const response = await userInfoCall( @@ -347,6 +350,7 @@ const UserDashboard: React.FC = ({ teams={teams} setSelectedTeam={setSelectedTeam} userRole={userRole} + proxySettings={proxySettings} />