Merge pull request #3184 from BerriAI/litellm_ui_non_admins_flow

[UI] - non admin flow - only Create + Test Key available

litellm/proxy/_types.py

@@ -87,6 +87,14 @@ class LiteLLMRoutes(enum.Enum):
         "/v2/key/info",
     ]
 
+    sso_only_routes: List = [
+        "/key/generate",
+        "/key/update",
+        "/key/delete",
+        "/global/spend/logs",
+        "/global/predict/spend/logs",
+    ]
+
     management_routes: List = [  # key
         "/key/generate",
         "/key/update",

litellm/proxy/proxy_server.py

@@ -1053,6 +1053,11 @@ async def user_api_key_auth(
                                 status_code=status.HTTP_403_FORBIDDEN,
                                 detail="key not allowed to access this team's info",
                             )
+                elif (
+                    _has_user_setup_sso()
+                    and route in LiteLLMRoutes.sso_only_routes.value
+                ):
+                    pass
                 else:
                     raise Exception(
                         f"Only master key can be used to generate, delete, update info for new keys/users/teams. Route={route}"
@@ -1102,6 +1107,13 @@ async def user_api_key_auth(
                     return UserAPIKeyAuth(
                         api_key=api_key, user_role="proxy_admin", **valid_token_dict
                     )
+                elif (
+                    _has_user_setup_sso()
+                    and route in LiteLLMRoutes.sso_only_routes.value
+                ):
+                    return UserAPIKeyAuth(
+                        api_key=api_key, user_role="app_owner", **valid_token_dict
+                    )
                 else:
                     raise Exception(
                         f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
@@ -5721,6 +5733,20 @@ async def new_user(data: NewUserRequest):
             "user"  # only create a user, don't create key if 'auto_create_key' set to False
         )
     response = await generate_key_helper_fn(**data_json)
+
+    # Admin UI Logic
+    # if team_id passed add this user to the team
+    if data_json.get("team_id", None) is not None:
+        await team_member_add(
+            data=TeamMemberAddRequest(
+                team_id=data_json.get("team_id", None),
+                member=Member(
+                    user_id=data_json.get("user_id", None),
+                    role="user",
+                    user_email=data_json.get("user_email", None),
+                ),
+            )
+        )
     return NewUserResponse(
         key=response.get("token", ""),
         expires=response.get("expires", None),
@@ -6526,13 +6552,20 @@ async def team_member_add(
     existing_team_row = await prisma_client.get_data(  # type: ignore
         team_id=data.team_id, table_name="team", query_type="find_unique"
     )
+    if existing_team_row is None:
+        raise HTTPException(
+            status_code=404,
+            detail={
+                "error": f"Team not found for team_id={getattr(data, 'team_id', None)}"
+            },
+        )
 
     new_member = data.member
 
     existing_team_row.members_with_roles.append(new_member)
 
     complete_team_data = LiteLLM_TeamTable(
-        **existing_team_row.model_dump(),
+        **_get_pydantic_json_dict(existing_team_row),
     )
 
     team_row = await prisma_client.update_data(
@@ -8120,36 +8153,33 @@ async def auth_callback(request: Request):
                 }
                 user_role = getattr(user_info, "user_role", None)
 
-            else:
-                ## check if user-email in db ##
-                user_info = await prisma_client.db.litellm_usertable.find_first(
-                    where={"user_email": user_email}
-                )
-                if user_info is not None:
-                    user_defined_values = {
-                        "models": getattr(user_info, "models", user_id_models),
-                        "user_id": getattr(user_info, "user_id", user_id),
-                        "user_email": getattr(user_info, "user_id", user_email),
-                        "user_role": getattr(user_info, "user_role", None),
-                    }
-                    user_role = getattr(user_info, "user_role", None)
+            ## check if user-email in db ##
+            user_info = await prisma_client.db.litellm_usertable.find_first(
+                where={"user_email": user_email}
+            )
+            if user_info is not None:
+                user_defined_values = {
+                    "models": getattr(user_info, "models", user_id_models),
+                    "user_id": getattr(user_info, "user_id", user_id),
+                    "user_email": getattr(user_info, "user_id", user_email),
+                    "user_role": getattr(user_info, "user_role", None),
+                }
+                user_role = getattr(user_info, "user_role", None)
 
-                    # update id
-                    await prisma_client.db.litellm_usertable.update_many(
-                        where={"user_email": user_email}, data={"user_id": user_id}  # type: ignore
-                    )
-                elif litellm.default_user_params is not None and isinstance(
-                    litellm.default_user_params, dict
-                ):
-                    user_defined_values = {
-                        "models": litellm.default_user_params.get(
-                            "models", user_id_models
-                        ),
-                        "user_id": litellm.default_user_params.get("user_id", user_id),
-                        "user_email": litellm.default_user_params.get(
-                            "user_email", user_email
-                        ),
-                    }
+                # update id
+                await prisma_client.db.litellm_usertable.update_many(
+                    where={"user_email": user_email}, data={"user_id": user_id}  # type: ignore
+                )
+            elif litellm.default_user_params is not None and isinstance(
+                litellm.default_user_params, dict
+            ):
+                user_defined_values = {
+                    "models": litellm.default_user_params.get("models", user_id_models),
+                    "user_id": litellm.default_user_params.get("user_id", user_id),
+                    "user_email": litellm.default_user_params.get(
+                        "user_email", user_email
+                    ),
+                }
     except Exception as e:
         pass
 

litellm/tests/test_key_generate_prisma.py

@@ -120,6 +120,15 @@ async def test_new_user_response(prisma_client):
         await litellm.proxy.proxy_server.prisma_client.connect()
         from litellm.proxy.proxy_server import user_api_key_cache
 
+        await new_team(
+            NewTeamRequest(
+                team_id="ishaan-special-team",
+            ),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+            ),
+        )
+
         _response = await new_user(
             data=NewUserRequest(
                 models=["azure-gpt-3.5"],
@@ -999,10 +1008,32 @@ def test_generate_and_update_key(prisma_client):
 
         async def test():
             await litellm.proxy.proxy_server.prisma_client.connect()
+
+            # create team "litellm-core-infra@gmail.com""
+            print("creating team litellm-core-infra@gmail.com")
+            await new_team(
+                NewTeamRequest(
+                    team_id="litellm-core-infra@gmail.com",
+                ),
+                user_api_key_dict=UserAPIKeyAuth(
+                    user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+                ),
+            )
+
+            await new_team(
+                NewTeamRequest(
+                    team_id="ishaan-special-team",
+                ),
+                user_api_key_dict=UserAPIKeyAuth(
+                    user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+                ),
+            )
+
             request = NewUserRequest(
-                metadata={"team": "litellm-team3", "project": "litellm-project3"},
+                metadata={"project": "litellm-project3"},
                 team_id="litellm-core-infra@gmail.com",
             )
+
             key = await new_user(request)
             print(key)
 
@@ -1015,7 +1046,6 @@ def test_generate_and_update_key(prisma_client):
             print("\n info for key=", result["info"])
             assert result["info"]["max_parallel_requests"] == None
             assert result["info"]["metadata"] == {
-                "team": "litellm-team3",
                 "project": "litellm-project3",
             }
             assert result["info"]["team_id"] == "litellm-core-infra@gmail.com"
@@ -1037,7 +1067,7 @@ def test_generate_and_update_key(prisma_client):
             # update the team id
             response2 = await update_key_fn(
                 request=Request,
-                data=UpdateKeyRequest(key=generated_key, team_id="ishaan"),
+                data=UpdateKeyRequest(key=generated_key, team_id="ishaan-special-team"),
             )
             print("response2=", response2)
 
@@ -1048,11 +1078,10 @@ def test_generate_and_update_key(prisma_client):
             print("\n info for key=", result["info"])
             assert result["info"]["max_parallel_requests"] == None
             assert result["info"]["metadata"] == {
-                "team": "litellm-team3",
                 "project": "litellm-project3",
             }
             assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"]
-            assert result["info"]["team_id"] == "ishaan"
+            assert result["info"]["team_id"] == "ishaan-special-team"
 
             # cleanup - delete key
             delete_key_request = KeyRequest(keys=[generated_key])
@@ -1941,6 +1970,15 @@ async def test_master_key_hashing(prisma_client):
         await litellm.proxy.proxy_server.prisma_client.connect()
         from litellm.proxy.proxy_server import user_api_key_cache
 
+        await new_team(
+            NewTeamRequest(
+                team_id="ishaan-special-team",
+            ),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+            ),
+        )
+
         _response = await new_user(
             data=NewUserRequest(
                 models=["azure-gpt-3.5"],

tests/test_keys.py

@@ -14,6 +14,24 @@ sys.path.insert(
 import litellm
 
 
+async def generate_team(session):
+    url = "http://0.0.0.0:4000/team/new"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {
+        "team_id": "litellm-dashboard",
+    }
+
+    async with session.post(url, headers=headers, json=data) as response:
+        status = response.status
+        response_text = await response.text()
+
+        print(f"Response (Status code: {status}):")
+        print(response_text)
+        print()
+        _json_response = await response.json()
+        return _json_response
+
+
 async def generate_user(
     session,
     user_role="app_owner",
@@ -668,7 +686,7 @@ async def test_key_rate_limit():
 
 
 @pytest.mark.asyncio
-async def test_key_delete():
+async def test_key_delete_ui():
     """
     Admin UI flow - DO NOT DELETE
     -> Create a key with user_id = "ishaan"
@@ -680,6 +698,8 @@ async def test_key_delete():
         key = key_gen["key"]
 
         # generate a admin UI key
+        team = await generate_team(session=session)
+        print("generated team: ", team)
         admin_ui_key = await generate_user(session=session, user_role="proxy_admin")
         print(
             "trying to delete key=",

ui/litellm-dashboard/src/components/create_key_button.tsx

@@ -116,7 +116,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
           wrapperCol={{ span: 16 }}
           labelAlign="left"
         >
-          {userRole === "App Owner" || userRole === "Admin" ? (
+          {userRole === "App Owner" || userRole === "Admin" || userRole === "App User" ? (
             <>
               <Form.Item 
                 label="Key Name" 

ui/litellm-dashboard/src/components/leftnav.tsx

@@ -63,11 +63,16 @@ const Sidebar: React.FC<SidebarProps> = ({
             Test Key
             </Text>
           </Menu.Item>
-          <Menu.Item key="2" onClick={() => setPage("models")}>
-          <Text>
-            Models
-            </Text>
-          </Menu.Item>
+          {
+            userRole == "Admin" ? (
+              <Menu.Item key="2" onClick={() => setPage("models")}>
+              <Text>
+                Models
+                </Text>
+              </Menu.Item>
+            ) : null
+          }
+
           {userRole == "Admin" ? (
             <Menu.Item key="6" onClick={() => setPage("teams")}>
               <Text>
@@ -75,11 +80,18 @@ const Sidebar: React.FC<SidebarProps> = ({
               </Text>
             </Menu.Item>
           ) : null}
-          <Menu.Item key="4" onClick={() => setPage("usage")}>
-          <Text>
-            Usage
-            </Text>
-          </Menu.Item>
+
+          {
+            userRole == "Admin" ? (
+              <Menu.Item key="4" onClick={() => setPage("usage")}>
+            <Text>
+              Usage
+              </Text>
+            </Menu.Item>
+
+              ) : null
+          }
+          
             {userRole == "Admin" ? (
             <Menu.Item key="5" onClick={() => setPage("users")}>
               <Text>
@@ -87,16 +99,27 @@ const Sidebar: React.FC<SidebarProps> = ({
               </Text>
             </Menu.Item>
           ) : null}
-          <Menu.Item key="8" onClick={() => setPage("settings")}>
-          <Text>
-            Integrations
-          </Text>
-          </Menu.Item>
-          <Menu.Item key="9" onClick={() => setPage("general-settings")}>
-          <Text>
-            Settings
-          </Text>
-          </Menu.Item>
+
+          {
+            userRole == "Admin" ? (
+              <Menu.Item key="8" onClick={() => setPage("settings")}>
+                <Text>
+                  Integrations
+                </Text>
+                </Menu.Item>
+              ) : null
+          }
+          
+          {
+            userRole == "Admin" ? (
+              <Menu.Item key="9" onClick={() => setPage("general-settings")}>
+              <Text>
+                Settings
+              </Text>
+              </Menu.Item>
+            ) : null
+          }
+
           {userRole == "Admin" ? (
             <Menu.Item key="7" onClick={() => setPage("admin-panel")}>
               <Text>

ui/litellm-dashboard/src/components/networking.tsx

@@ -296,6 +296,9 @@ export const userInfoCall = async (
     if (userRole == "App Owner" && userID) {
       url = `${url}?user_id=${userID}`;
     }
+    if (userRole == "App User" && userID) {
+      url = `${url}?user_id=${userID}`;
+    }
     console.log("in userInfoCall viewAll=", viewAll);
     if (viewAll && page_size && (page != null) && (page != undefined)) {
       url = `${url}?view_all=true&page=${page}&page_size=${page_size}`;