From 4dcecde97af6049676359523d844096b1fa58713 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 19 Apr 2024 16:45:13 -0700 Subject: [PATCH 1/8] ui - non admin flow --- .../src/components/leftnav.tsx | 63 +++++++++++++------ 1 file changed, 43 insertions(+), 20 deletions(-) diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 04a99c41f..dd6ef0970 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -63,11 +63,16 @@ const Sidebar: React.FC = ({ 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")}> From b81d66639b3dfe70593d505fcf3b40c2f4c00821 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 19 Apr 2024 17:31:36 -0700 Subject: [PATCH 2/8] ui - allow App user to see their own info --- ui/litellm-dashboard/src/components/networking.tsx | 3 +++ 1 file changed, 3 insertions(+) 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}`; From 00a07a99cda2650ed9e65a47dd7050fea1913f92 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 19 Apr 2024 17:36:29 -0700 Subject: [PATCH 3/8] fix - backend logic for non admin flow --- litellm/proxy/proxy_server.py | 69 ++++++++++++++++++++--------------- 1 file changed, 40 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index db85b7ba1..14ff78200 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5713,6 +5713,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), @@ -8112,36 +8126,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 From abeadadf3ee5de9b668ce0ed8a4a227d73c6c5ab Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 19 Apr 2024 21:15:57 -0700 Subject: [PATCH 4/8] fix - create key user flow --- ui/litellm-dashboard/src/components/create_key_button.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" ? ( <> Date: Sat, 20 Apr 2024 11:09:34 -0700 Subject: [PATCH 5/8] (ci/cd) testing with team_id and /user/new --- litellm/proxy/proxy_server.py | 7 ++++ litellm/tests/test_key_generate_prisma.py | 48 ++++++++++++++++++++--- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 14ff78200..2aab7e453 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6532,6 +6532,13 @@ 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 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"], From a34f725db7bc6f176dd826ba60ff8b553e0b0b37 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 20 Apr 2024 11:43:43 -0700 Subject: [PATCH 6/8] fix - test keys --- tests/test_keys.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_keys.py b/tests/test_keys.py index 39787eb97..7a038bf35 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", @@ -680,6 +698,7 @@ async def test_key_delete(): key = key_gen["key"] # generate a admin UI key + generate_team(session=session) admin_ui_key = await generate_user(session=session, user_role="proxy_admin") print( "trying to delete key=", From fd282ea9325da9a1a4ef9e4b932b075a07d48e7c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 20 Apr 2024 11:48:41 -0700 Subject: [PATCH 7/8] fix testing fixes --- litellm/proxy/_types.py | 8 ++++++++ litellm/proxy/proxy_server.py | 14 +++++++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) 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 2aab7e453..ebeea120d 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" @@ -6545,7 +6557,7 @@ async def team_member_add( 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( From 7b24a74e772f20297c0762b4cfa4020929dd19bf Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 20 Apr 2024 12:30:58 -0700 Subject: [PATCH 8/8] test fix - test_key_delete_ui --- tests/test_keys.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_keys.py b/tests/test_keys.py index 7a038bf35..f21c50c0d 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -686,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" @@ -698,7 +698,8 @@ async def test_key_delete(): key = key_gen["key"] # generate a admin UI key - generate_team(session=session) + 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=",