diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b697b6e97..ca9926cef 100644 --- a/litellm/proxy/_types.py +++ b/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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1aab7bac2..7f8d7d4a3 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index fdb7649d5..a90c13803 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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"], diff --git a/tests/test_keys.py b/tests/test_keys.py index 39787eb97..f21c50c0d 100644 --- a/tests/test_keys.py +++ b/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=", diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 8dde3fb00..d7fb9c5eb 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -116,7 +116,7 @@ const CreateKey: React.FC = ({ wrapperCol={{ span: 16 }} labelAlign="left" > - {userRole === "App Owner" || userRole === "Admin" ? ( + {userRole === "App Owner" || userRole === "Admin" || userRole === "App User" ? ( <> = ({ Test Key - setPage("models")}> - - Models - - + { + userRole == "Admin" ? ( + setPage("models")}> + + Models + + + ) : null + } + {userRole == "Admin" ? ( setPage("teams")}> @@ -75,11 +80,18 @@ const Sidebar: React.FC = ({ ) : null} - setPage("usage")}> - - Usage - - + + { + userRole == "Admin" ? ( + setPage("usage")}> + + Usage + + + + ) : null + } + {userRole == "Admin" ? ( setPage("users")}> @@ -87,16 +99,27 @@ const Sidebar: React.FC = ({ ) : null} - setPage("settings")}> - - Integrations - - - setPage("general-settings")}> - - Settings - - + + { + userRole == "Admin" ? ( + setPage("settings")}> + + Integrations + + + ) : null + } + + { + userRole == "Admin" ? ( + setPage("general-settings")}> + + Settings + + + ) : null + } + {userRole == "Admin" ? ( setPage("admin-panel")}> diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 4b961ca34..96b6246f5 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/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}`;