feat(ui_sso.py): Allows users to use test key pane, and have team bud…

…get limits be enforced for their use-case (BerriAI#7666)
inker610566 · Jan 10, 2025 · 63926f4 · 63926f4
1 parent 6df8faf
commit 63926f4
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 2 deletions.
diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
@@ -24,4 +24,11 @@ model_list:
       custom_tokenizer: 
         identifier: deepseek-ai/DeepSeek-V3-Base
         revision: main
-        auth_token: os.environ/HUGGINGFACE_API_KEY
+        auth_token: os.environ/HUGGINGFACE_API_KEY
+
+# litellm_settings:
+#   key_generation_settings:
+#     personal_key_generation: # maps to 'Default Team' on UI 
+#       allowed_user_roles: ["proxy_admin"]
+
+
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
@@ -391,6 +391,17 @@ async def add_missing_team_member(user_info: NewUserResponse, sso_teams: List[st
         )
 
 
+def get_disabled_non_admin_personal_key_creation():
+    key_generation_settings = litellm.key_generation_settings
+    if key_generation_settings is None:
+        return False
+    personal_key_generation = (
+        key_generation_settings.get("personal_key_generation") or {}
+    )
+    allowed_user_roles = personal_key_generation.get("allowed_user_roles") or []
+    return bool("proxy_admin" in allowed_user_roles)
+
+
 @router.get("/sso/callback", tags=["experimental"], include_in_schema=False)
 async def auth_callback(request: Request):  # noqa: PLR0915
     """Verify login"""
@@ -640,6 +651,10 @@ async def auth_callback(request: Request):  # noqa: PLR0915
                 },
             )
 
+    disabled_non_admin_personal_key_creation = (
+        get_disabled_non_admin_personal_key_creation()
+    )
+
     import jwt
 
     jwt_token = jwt.encode(  # type: ignore
@@ -653,6 +668,7 @@ async def auth_callback(request: Request):  # noqa: PLR0915
             "auth_header_name": general_settings.get(
                 "litellm_key_header_name", "Authorization"
             ),
+            "disabled_non_admin_personal_key_creation": disabled_non_admin_personal_key_creation,
         },
         master_key,
         algorithm="HS256",

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
@@ -204,6 +204,9 @@ def generate_feedback_box():
 )
 from litellm.proxy.management_endpoints.team_endpoints import router as team_router
 from litellm.proxy.management_endpoints.team_endpoints import update_team
+from litellm.proxy.management_endpoints.ui_sso import (
+    get_disabled_non_admin_personal_key_creation,
+)
 from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router
 from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
 from litellm.proxy.openai_files_endpoints.files_endpoints import (
@@ -7295,6 +7298,9 @@ async def login(request: Request):  # noqa: PLR0915
         _user_row = await prisma_client.db.litellm_usertable.find_first(
             where={"user_email": {"equals": username}}
         )
+    disabled_non_admin_personal_key_creation = (
+        get_disabled_non_admin_personal_key_creation()
+    )
     """
     To login to Admin UI, we support the following 
     - Login with UI_USERNAME and UI_PASSWORD
@@ -7366,6 +7372,7 @@ async def login(request: Request):  # noqa: PLR0915
                 "auth_header_name": general_settings.get(
                     "litellm_key_header_name", "Authorization"
                 ),
+                "disabled_non_admin_personal_key_creation": disabled_non_admin_personal_key_creation,
             },
             master_key,
             algorithm="HS256",
@@ -7433,6 +7440,7 @@ async def login(request: Request):  # noqa: PLR0915
                     "auth_header_name": general_settings.get(
                         "litellm_key_header_name", "Authorization"
                     ),
+                    "disabled_non_admin_personal_key_creation": disabled_non_admin_personal_key_creation,
                 },
                 master_key,
                 algorithm="HS256",
@@ -7547,6 +7555,10 @@ async def onboarding(invite_link: str):
         litellm_dashboard_ui += "/ui/onboarding"
     import jwt
 
+    disabled_non_admin_personal_key_creation = (
+        get_disabled_non_admin_personal_key_creation()
+    )
+
     jwt_token = jwt.encode(  # type: ignore
         {
             "user_id": user_obj.user_id,
@@ -7558,6 +7570,7 @@ async def onboarding(invite_link: str):
             "auth_header_name": general_settings.get(
                 "litellm_key_header_name", "Authorization"
             ),
+            "disabled_non_admin_personal_key_creation": disabled_non_admin_personal_key_creation,
         },
         master_key,
         algorithm="HS256",

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
@@ -68,6 +68,7 @@ const CreateKeyPage = () => {
   const { Title, Paragraph } = Typography;
   const [userRole, setUserRole] = useState("");
   const [premiumUser, setPremiumUser] = useState(false);
+  const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false);
   const [userEmail, setUserEmail] = useState<null | string>(null);
   const [teams, setTeams] = useState<null | any[]>(null);
   const [keys, setKeys] = useState<null | any[]>(null);
@@ -117,6 +118,8 @@ const CreateKeyPage = () => {
         // set accessToken
         setAccessToken(decoded.key);
 
+        setDisabledPersonalKeyCreation(decoded.disabled_non_admin_personal_key_creation);
+
         // check if userRole is defined
         if (decoded.user_role) {
           const formattedUserRole = formatUserRole(decoded.user_role);
@@ -220,6 +223,7 @@ const CreateKeyPage = () => {
               userRole={userRole}
               token={token}
               accessToken={accessToken}
+              disabledPersonalKeyCreation={disabledPersonalKeyCreation}
             />
           ) : page == "users" ? (
             <ViewUserDashboard

diff --git a/ui/litellm-dashboard/src/components/chat_ui.tsx b/ui/litellm-dashboard/src/components/chat_ui.tsx
@@ -35,6 +35,7 @@ interface ChatUIProps {
   token: string | null;
   userRole: string | null;
   userID: string | null;
+  disabledPersonalKeyCreation: boolean;
 }
 
 async function generateModelResponse(
@@ -81,8 +82,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
   token,
   userRole,
   userID,
+  disabledPersonalKeyCreation,
 }) => {
-  const [apiKeySource, setApiKeySource] = useState<'session' | 'custom'>('session');
+  const [apiKeySource, setApiKeySource] = useState<'session' | 'custom'>(
+    disabledPersonalKeyCreation ? 'custom' : 'session'
+  );
   const [apiKey, setApiKey] = useState("");
   const [inputMessage, setInputMessage] = useState("");
   const [chatHistory, setChatHistory] = useState<{ role: string; content: string; model?: string }[]>([]);
@@ -240,6 +244,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     <Col>
                       <Text>API Key Source</Text>
                       <Select
+                        disabled={disabledPersonalKeyCreation}
                         defaultValue="session"
                         style={{ width: "100%" }}
                         onChange={(value) => setApiKeySource(value as "session" | "custom")}