From fc54acc17ab0dd19df58bb44110a93077d6b6271 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 13:51:58 -0700 Subject: [PATCH 01/11] fix - use LitellmUserRoles --- litellm/proxy/_types.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 555254a633..e4195c196c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -7,6 +7,30 @@ import uuid, json, sys, os from litellm.types.router import UpdateRouterConfig from litellm.types.utils import ProviderField + +class LitellmUserRoles(enum.Enum): + """ + proxy_admin: admin over the platform + proxy_admin_viewer: can login, view their own keys, view their spend + internal_user: can login, view/create/delete their own keys, view their spend + + """ + + # Admin Roles + PROXY_ADMIN = "proxy_admin" + PROXY_ADMIN_VIEW_ONLY = "proxy_admin_view_only" + + # Internal User Roles + INTERNAL_USER = "internal_user" + INTERNAL_USER_VIEW_ONLY = "internal_user_view_only" + + # Team Roles + TEAM = "team" + + # Customer Roles - External users of proxy + CUSTOMER = "customer" + + AlertType = Literal[ "llm_exceptions", "llm_too_slow", From b29adeb6c6b198b62169cf1059b0ee4704cb7888 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 13:55:45 -0700 Subject: [PATCH 02/11] feat - move to using enums for roles --- litellm/proxy/proxy_server.py | 79 ++++++++++++++++++++--------------- 1 file changed, 46 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f0b98d0be8..726a6c51e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -524,12 +524,14 @@ async def user_api_key_auth( if is_admin: # check allowed admin routes is_allowed = allowed_routes_check( - user_role="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, user_route=route, litellm_proxy_roles=jwt_handler.litellm_jwtauth, ) if is_allowed: - return UserAPIKeyAuth(user_role="proxy_admin") + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value + ) else: allowed_routes = ( jwt_handler.litellm_jwtauth.admin_allowed_routes @@ -671,9 +673,11 @@ async def user_api_key_auth( #### ELSE #### if master_key is None: if isinstance(api_key, str): - return UserAPIKeyAuth(api_key=api_key, user_role="proxy_admin") + return UserAPIKeyAuth( + api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN.value + ) else: - return UserAPIKeyAuth(user_role="proxy_admin") + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value) elif api_key is None: # only require api key if master key is set raise Exception("No api key passed in.") elif api_key == "": @@ -740,7 +744,7 @@ async def user_api_key_auth( if ( valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.user_role == "proxy_admin" + and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN.value ): # update end-user params on valid token valid_token.end_user_id = end_user_params.get("end_user_id") @@ -773,7 +777,7 @@ async def user_api_key_auth( if is_master_key_valid: _user_api_key_obj = UserAPIKeyAuth( api_key=master_key, - user_role="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id=litellm_proxy_admin_name, **end_user_params, ) @@ -1377,7 +1381,9 @@ async def user_api_key_auth( user_id_information ): return UserAPIKeyAuth( - api_key=api_key, user_role="proxy_admin", **valid_token_dict + api_key=api_key, + user_role=LitellmUserRoles.PROXY_ADMIN.value, + **valid_token_dict, ) elif ( _has_user_setup_sso() @@ -1398,15 +1404,21 @@ async def user_api_key_auth( user_id_information ): return UserAPIKeyAuth( - api_key=api_key, user_role="proxy_admin", **valid_token_dict + api_key=api_key, + user_role=LitellmUserRoles.PROXY_ADMIN.value, + **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 + api_key=api_key, + user_role=LitellmUserRoles.INTERNAL_USER.value, + **valid_token_dict, ) else: return UserAPIKeyAuth( - api_key=api_key, user_role="app_owner", **valid_token_dict + api_key=api_key, + user_role=LitellmUserRoles.INTERNAL_USER.value, + **valid_token_dict, ) else: raise Exception() @@ -3738,11 +3750,9 @@ async def startup_event(): spend=0, token=master_key, user_id=litellm_proxy_admin_name, - user_role="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, query_type="update_data", - update_key_values={ - "user_role": "proxy_admin", - }, + update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) ) @@ -6093,7 +6103,7 @@ async def delete_key_fn( ) if ( user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == "proxy_admin" + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ): user_id = None # unless they're admin @@ -7888,7 +7898,10 @@ async def user_info( user_id=user_api_key_dict.user_id ) # *NEW* get all teams in user 'teams' field - if getattr(caller_user_info, "user_role", None) == "proxy_admin": + if ( + getattr(caller_user_info, "user_role", None) + == LitellmUserRoles.PROXY_ADMIN.value + ): teams_2 = await prisma_client.get_data( table_name="team", query_type="find_all", @@ -8716,7 +8729,7 @@ async def new_team( if ( user_api_key_dict.user_role is None - or user_api_key_dict.user_role != "proxy_admin" + or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): # don't restrict proxy admin if ( data.tpm_limit is not None @@ -9322,7 +9335,7 @@ async def list_team( """ global prisma_client - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=401, detail={ @@ -9416,7 +9429,7 @@ async def new_organization( if ( user_api_key_dict.user_role is None - or user_api_key_dict.user_role != "proxy_admin" + or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): raise HTTPException( status_code=401, @@ -9619,7 +9632,7 @@ async def budget_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -9684,7 +9697,7 @@ async def list_budget( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -9718,7 +9731,7 @@ async def delete_budget( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -10696,7 +10709,7 @@ async def alerting_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -10777,7 +10790,7 @@ async def alerting_settings( # detail={"error": CommonProxyErrors.db_not_connected_error.value}, # ) -# if user_api_key_dict.user_role != "proxy_admin": +# if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: # raise HTTPException( # status_code=400, # detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -11235,12 +11248,12 @@ async def login(request: Request): await user_update( data=UpdateUserRequest( user_id=key_user_id, - user_role="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, ) ) if os.getenv("DATABASE_URL") is not None: response = await generate_key_helper_fn( - **{"user_role": "proxy_admin", "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"} # type: ignore + **{"user_role": LitellmUserRoles.PROXY_ADMIN.value, "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"} # type: ignore ) else: raise ProxyException( @@ -11635,7 +11648,7 @@ async def new_invitation( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -11699,7 +11712,7 @@ async def invitation_info( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -11811,7 +11824,7 @@ async def invitation_delete( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -12006,7 +12019,7 @@ async def update_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -12080,7 +12093,7 @@ async def get_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -12143,7 +12156,7 @@ async def get_config_list( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ @@ -12218,7 +12231,7 @@ async def delete_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != "proxy_admin": + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=400, detail={ From 6166381579002c603b66cc29de46e2725867aba6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 13:57:50 -0700 Subject: [PATCH 03/11] fix migrate to LitellmUserRoles --- litellm/proxy/utils.py | 5 +-- litellm/tests/test_add_update_models.py | 14 +++++--- litellm/tests/test_key_generate_prisma.py | 44 +++++++++++++++-------- tests/test_keys.py | 5 ++- 4 files changed, 47 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b96c469ec0..86131991bf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -15,6 +15,7 @@ from litellm.proxy._types import ( WebhookEvent, AlertType, ResetTeamBudgetRequest, + LitellmUserRoles, ) from litellm.caching import DualCache, RedisCache from litellm.router import Deployment, ModelInfo, LiteLLM_Params @@ -2637,7 +2638,7 @@ def _is_user_proxy_admin(user_id_information: Optional[list]): _user = user_id_information[0] if ( _user.get("user_role", None) is not None - and _user.get("user_role") == "proxy_admin" + and _user.get("user_role") == LitellmUserRoles.PROXY_ADMIN.value ): return True @@ -2650,7 +2651,7 @@ def _is_user_proxy_admin(user_id_information: Optional[list]): if ( _user.get("user_role", None) is not None - and _user.get("user_role") == "proxy_admin" + and _user.get("user_role") == LitellmUserRoles.PROXY_ADMIN.value ): return True diff --git a/litellm/tests/test_add_update_models.py b/litellm/tests/test_add_update_models.py index ec9ab33b6f..21124a4ba5 100644 --- a/litellm/tests/test_add_update_models.py +++ b/litellm/tests/test_add_update_models.py @@ -14,7 +14,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import pytest, logging, asyncio import litellm, asyncio -from litellm.proxy.proxy_server import add_new_model, update_model +from litellm.proxy.proxy_server import add_new_model, update_model, LitellmUserRoles from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -90,7 +90,9 @@ async def test_add_new_model(prisma_client): ), ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -137,7 +139,9 @@ async def test_add_update_model(prisma_client): ), ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -166,7 +170,9 @@ async def test_add_update_model(prisma_client): ), ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 375a1c85d8..7ceab3663d 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -61,6 +61,7 @@ from litellm.proxy.proxy_server import ( audio_transcriptions, moderations, model_list, + LitellmUserRoles, ) from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend from litellm._logging import verbose_proxy_logger @@ -137,7 +138,9 @@ async def test_new_user_response(prisma_client): team_id=_team_id, ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -363,7 +366,10 @@ async def test_call_with_valid_model_using_all_models(prisma_client): ) new_team_response = await new_team( - data=team_request, user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin") + data=team_request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value + ), ) print("new_team_response", new_team_response) created_team_id = new_team_response["team_id"] @@ -922,7 +928,7 @@ def test_delete_key(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN.value # delete the key result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -972,7 +978,7 @@ def test_delete_key_auth(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN.value result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -1044,7 +1050,7 @@ def test_generate_and_call_key_info(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN.value result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -1078,7 +1084,9 @@ def test_generate_and_update_key(prisma_client): team_id=_team_1, ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -1088,7 +1096,9 @@ def test_generate_and_update_key(prisma_client): team_id=_team_2, ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -1158,7 +1168,7 @@ def test_generate_and_update_key(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN.value result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -2038,7 +2048,9 @@ async def test_master_key_hashing(prisma_client): await new_team( NewTeamRequest(team_id=_team_id), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -2076,7 +2088,7 @@ async def test_reset_spend_authentication(prisma_client): """ 1. Test master key can access this route -> ONLY MASTER KEY SHOULD BE ABLE TO RESET SPEND 2. Test that non-master key gets rejected - 3. Test that non-master key with role == "proxy_admin" or admin gets rejected + 3. Test that non-master key with role == LitellmUserRoles.PROXY_ADMIN.value or admin gets rejected """ print("prisma client=", prisma_client) @@ -2121,10 +2133,10 @@ async def test_reset_spend_authentication(prisma_client): in e.message ) - # Test 3 - Non-Master Key with role == "proxy_admin" or admin + # Test 3 - Non-Master Key with role == LitellmUserRoles.PROXY_ADMIN.value or admin _response = await new_user( data=NewUserRequest( - user_role="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, tpm_limit=20, ) ) @@ -2174,7 +2186,9 @@ async def test_create_update_team(prisma_client): rpm_limit=20, ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) @@ -2200,7 +2214,9 @@ async def test_create_update_team(prisma_client): rpm_limit=30, ), user_api_key_dict=UserAPIKeyAuth( - user_role="proxy_admin", api_key="sk-1234", user_id="1234" + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + user_id="1234", ), ) diff --git a/tests/test_keys.py b/tests/test_keys.py index 11961e2a20..bbe4cc35eb 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -12,6 +12,7 @@ sys.path.insert( 0, os.path.abspath("../") ) # Adds the parent directory to the system path import litellm +from litellm.proxy._types import LitellmUserRoles async def generate_team( @@ -731,7 +732,9 @@ async def test_key_delete_ui(): # generate a admin UI key team = await generate_team(session=session) - admin_ui_key = await generate_user(session=session, user_role="proxy_admin") + admin_ui_key = await generate_user( + session=session, user_role=LitellmUserRoles.PROXY_ADMIN.value + ) print( "trying to delete key=", key, From 4861ff2fd4cf3b081fb80b8f386398cbdc061678 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 14:10:08 -0700 Subject: [PATCH 04/11] fix user_role test --- litellm/proxy/_types.py | 11 ++++++++++- litellm/proxy/proxy_server.py | 6 ++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4195c196c..1806fabc77 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1086,7 +1086,16 @@ class UserAPIKeyAuth( """ api_key: Optional[str] = None - user_role: Optional[Literal["proxy_admin", "app_owner", "app_user"]] = None + user_role: Optional[ + Literal[ + "proxy_admin", + "proxy_admin_view_only", + "internal_user", + "internal_user_view_only", + "team", + "customer", + ] + ] = None allowed_model_region: Optional[Literal["eu"]] = None @root_validator(pre=True) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 726a6c51e2..629ab3dd3b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -507,7 +507,9 @@ async def user_api_key_auth( if route in LiteLLMRoutes.public_routes.value: # check if public endpoint - return UserAPIKeyAuth(user_role="app_owner") + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + ) if general_settings.get("enable_jwt_auth", False) == True: is_jwt = jwt_handler.is_jwt(token=api_key) @@ -666,7 +668,7 @@ async def user_api_key_auth( team_object.rpm_limit if team_object is not None else None ), team_models=team_object.models if team_object is not None else [], - user_role="app_owner", + user_role=LitellmUserRoles.INTERNAL_USER.value, user_id=user_id, org_id=org_id, ) From a4b6a959d807ee50dfa1cf7a05e41a69e438d8de Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 14:28:53 -0700 Subject: [PATCH 05/11] fix literal usage --- litellm/proxy/_types.py | 34 ++++++++--- litellm/proxy/auth/auth_checks.py | 11 +++- litellm/proxy/proxy_server.py | 72 +++++++++++------------ litellm/tests/test_key_generate_prisma.py | 30 +++++----- 4 files changed, 82 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1806fabc77..984e65d694 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -522,7 +522,16 @@ class LiteLLM_ModelTable(LiteLLMBase): class NewUserRequest(GenerateKeyRequest): max_budget: Optional[float] = None user_email: Optional[str] = None - user_role: Optional[str] = None + user_role: Optional[ + Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + LitellmUserRoles.CUSTOMER, + ] + ] = None teams: Optional[list] = None organization_id: Optional[str] = None auto_create_key: bool = ( @@ -541,7 +550,16 @@ class UpdateUserRequest(GenerateRequestBase): user_email: Optional[str] = None spend: Optional[float] = None metadata: Optional[dict] = None - user_role: Optional[str] = None + user_role: Optional[ + Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + LitellmUserRoles.CUSTOMER, + ] + ] = None max_budget: Optional[float] = None @root_validator(pre=True) @@ -1088,12 +1106,12 @@ class UserAPIKeyAuth( api_key: Optional[str] = None user_role: Optional[ Literal[ - "proxy_admin", - "proxy_admin_view_only", - "internal_user", - "internal_user_view_only", - "team", - "customer", + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + LitellmUserRoles.CUSTOMER, ] ] = None allowed_model_region: Optional[Literal["eu"]] = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a6e97960e5..e4b8e6c8a8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,6 +15,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLMRoutes, LiteLLM_OrganizationTable, + LitellmUserRoles, ) from typing import Optional, Literal, Union from litellm.proxy.utils import PrismaClient @@ -133,7 +134,11 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: def allowed_routes_check( - user_role: Literal["proxy_admin", "team", "user"], + user_role: Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.TEAM, + LitellmUserRoles.INTERNAL_USER, + ], user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth, ) -> bool: @@ -141,14 +146,14 @@ def allowed_routes_check( Check if user -> not admin - allowed to access these routes """ - if user_role == "proxy_admin": + if user_role == LitellmUserRoles.PROXY_ADMIN: is_allowed = _allowed_routes_check( user_route=user_route, allowed_routes=litellm_proxy_roles.admin_allowed_routes, ) return is_allowed - elif user_role == "team": + elif user_role == LitellmUserRoles.TEAM: if litellm_proxy_roles.team_allowed_routes is None: """ By default allow a team to call openai + info routes diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 629ab3dd3b..f3a9b56924 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -507,9 +507,7 @@ async def user_api_key_auth( if route in LiteLLMRoutes.public_routes.value: # check if public endpoint - return UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - ) + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) if general_settings.get("enable_jwt_auth", False) == True: is_jwt = jwt_handler.is_jwt(token=api_key) @@ -526,14 +524,12 @@ async def user_api_key_auth( if is_admin: # check allowed admin routes is_allowed = allowed_routes_check( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, user_route=route, litellm_proxy_roles=jwt_handler.litellm_jwtauth, ) if is_allowed: - return UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) else: allowed_routes = ( jwt_handler.litellm_jwtauth.admin_allowed_routes @@ -556,7 +552,7 @@ async def user_api_key_auth( if team_id is not None: # check allowed team routes is_allowed = allowed_routes_check( - user_role="team", + user_role=LitellmUserRoles.TEAM, user_route=route, litellm_proxy_roles=jwt_handler.litellm_jwtauth, ) @@ -668,7 +664,7 @@ async def user_api_key_auth( team_object.rpm_limit if team_object is not None else None ), team_models=team_object.models if team_object is not None else [], - user_role=LitellmUserRoles.INTERNAL_USER.value, + user_role=LitellmUserRoles.INTERNAL_USER, user_id=user_id, org_id=org_id, ) @@ -676,10 +672,10 @@ async def user_api_key_auth( if master_key is None: if isinstance(api_key, str): return UserAPIKeyAuth( - api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN.value + api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN ) else: - return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value) + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) elif api_key is None: # only require api key if master key is set raise Exception("No api key passed in.") elif api_key == "": @@ -746,7 +742,7 @@ async def user_api_key_auth( if ( valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN.value + and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN ): # update end-user params on valid token valid_token.end_user_id = end_user_params.get("end_user_id") @@ -779,7 +775,7 @@ async def user_api_key_auth( if is_master_key_valid: _user_api_key_obj = UserAPIKeyAuth( api_key=master_key, - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, user_id=litellm_proxy_admin_name, **end_user_params, ) @@ -1384,7 +1380,7 @@ async def user_api_key_auth( ): return UserAPIKeyAuth( api_key=api_key, - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, **valid_token_dict, ) elif ( @@ -1407,19 +1403,19 @@ async def user_api_key_auth( ): return UserAPIKeyAuth( api_key=api_key, - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.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=LitellmUserRoles.INTERNAL_USER.value, + user_role=LitellmUserRoles.INTERNAL_USER, **valid_token_dict, ) else: return UserAPIKeyAuth( api_key=api_key, - user_role=LitellmUserRoles.INTERNAL_USER.value, + user_role=LitellmUserRoles.INTERNAL_USER, **valid_token_dict, ) else: @@ -3752,9 +3748,9 @@ async def startup_event(): spend=0, token=master_key, user_id=litellm_proxy_admin_name, - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, query_type="update_data", - update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, + update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN}, ) ) @@ -6105,7 +6101,7 @@ async def delete_key_fn( ) if ( user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ): user_id = None # unless they're admin @@ -7902,7 +7898,7 @@ async def user_info( # *NEW* get all teams in user 'teams' field if ( getattr(caller_user_info, "user_role", None) - == LitellmUserRoles.PROXY_ADMIN.value + == LitellmUserRoles.PROXY_ADMIN ): teams_2 = await prisma_client.get_data( table_name="team", @@ -8731,7 +8727,7 @@ async def new_team( if ( user_api_key_dict.user_role is None - or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): # don't restrict proxy admin if ( data.tpm_limit is not None @@ -9337,7 +9333,7 @@ async def list_team( """ global prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=401, detail={ @@ -9431,7 +9427,7 @@ async def new_organization( if ( user_api_key_dict.user_role is None - or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): raise HTTPException( status_code=401, @@ -9634,7 +9630,7 @@ async def budget_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -9699,7 +9695,7 @@ async def list_budget( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -9733,7 +9729,7 @@ async def delete_budget( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -10711,7 +10707,7 @@ async def alerting_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -10792,7 +10788,7 @@ async def alerting_settings( # detail={"error": CommonProxyErrors.db_not_connected_error.value}, # ) -# if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: +# if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: # raise HTTPException( # status_code=400, # detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -11250,12 +11246,12 @@ async def login(request: Request): await user_update( data=UpdateUserRequest( user_id=key_user_id, - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, ) ) if os.getenv("DATABASE_URL") is not None: response = await generate_key_helper_fn( - **{"user_role": LitellmUserRoles.PROXY_ADMIN.value, "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"} # type: ignore + **{"user_role": LitellmUserRoles.PROXY_ADMIN, "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"} # type: ignore ) else: raise ProxyException( @@ -11650,7 +11646,7 @@ async def new_invitation( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -11714,7 +11710,7 @@ async def invitation_info( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -11826,7 +11822,7 @@ async def invitation_delete( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -12021,7 +12017,7 @@ async def update_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -12095,7 +12091,7 @@ async def get_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -12158,7 +12154,7 @@ async def get_config_list( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ @@ -12233,7 +12229,7 @@ async def delete_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=400, detail={ diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 7ceab3663d..6be720ffd1 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -138,7 +138,7 @@ async def test_new_user_response(prisma_client): team_id=_team_id, ), user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234", ), @@ -367,9 +367,7 @@ async def test_call_with_valid_model_using_all_models(prisma_client): new_team_response = await new_team( data=team_request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value - ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) print("new_team_response", new_team_response) created_team_id = new_team_response["team_id"] @@ -928,7 +926,7 @@ def test_delete_key(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN.value + result.user_role = LitellmUserRoles.PROXY_ADMIN # delete the key result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -978,7 +976,7 @@ def test_delete_key_auth(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN.value + result.user_role = LitellmUserRoles.PROXY_ADMIN result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -1050,7 +1048,7 @@ def test_generate_and_call_key_info(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN.value + result.user_role = LitellmUserRoles.PROXY_ADMIN result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -1084,7 +1082,7 @@ def test_generate_and_update_key(prisma_client): team_id=_team_1, ), user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234", ), @@ -1096,7 +1094,7 @@ def test_generate_and_update_key(prisma_client): team_id=_team_2, ), user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234", ), @@ -1168,7 +1166,7 @@ def test_generate_and_update_key(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN.value + result.user_role = LitellmUserRoles.PROXY_ADMIN result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -2048,7 +2046,7 @@ async def test_master_key_hashing(prisma_client): await new_team( NewTeamRequest(team_id=_team_id), user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234", ), @@ -2088,7 +2086,7 @@ async def test_reset_spend_authentication(prisma_client): """ 1. Test master key can access this route -> ONLY MASTER KEY SHOULD BE ABLE TO RESET SPEND 2. Test that non-master key gets rejected - 3. Test that non-master key with role == LitellmUserRoles.PROXY_ADMIN.value or admin gets rejected + 3. Test that non-master key with role == LitellmUserRoles.PROXY_ADMIN or admin gets rejected """ print("prisma client=", prisma_client) @@ -2133,10 +2131,10 @@ async def test_reset_spend_authentication(prisma_client): in e.message ) - # Test 3 - Non-Master Key with role == LitellmUserRoles.PROXY_ADMIN.value or admin + # Test 3 - Non-Master Key with role == LitellmUserRoles.PROXY_ADMIN or admin _response = await new_user( data=NewUserRequest( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, tpm_limit=20, ) ) @@ -2186,7 +2184,7 @@ async def test_create_update_team(prisma_client): rpm_limit=20, ), user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234", ), @@ -2214,7 +2212,7 @@ async def test_create_update_team(prisma_client): rpm_limit=30, ), user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234", ), From 0a3a37a196632c7fda8bc82dee9368fef2062f15 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 14:41:47 -0700 Subject: [PATCH 06/11] fix - /user/new --- litellm/proxy/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 984e65d694..9ee14788c9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -8,7 +8,7 @@ from litellm.types.router import UpdateRouterConfig from litellm.types.utils import ProviderField -class LitellmUserRoles(enum.Enum): +class LitellmUserRoles(str, enum.Enum): """ proxy_admin: admin over the platform proxy_admin_viewer: can login, view their own keys, view their spend From b03f00062eb57dd3a8fa7e7e79aa818e455eacf5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 14:47:04 -0700 Subject: [PATCH 07/11] add better details about LitellmUserRoles --- litellm/proxy/_types.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9ee14788c9..e502a5369e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -10,9 +10,21 @@ from litellm.types.utils import ProviderField class LitellmUserRoles(str, enum.Enum): """ - proxy_admin: admin over the platform - proxy_admin_viewer: can login, view their own keys, view their spend - internal_user: can login, view/create/delete their own keys, view their spend + Admin Roles: + PROXY_ADMIN: admin over the platform + PROXY_ADMIN_VIEW_ONLY: can login, view all own keys, view all spend + + Internal User Roles: + INTERNAL_USER: can login, view/create/delete their own keys, view their spend + INTERNAL_USER_VIEW_ONLY: can login, view their own keys, view their own spend + + + Team Roles: + TEAM: used for JWT auth + + + Customer Roles: + CUSTOMER: External users -> these are customers """ From 21831905ba74fbc9b90d7e309ef2d02289d3144e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 15:10:16 -0700 Subject: [PATCH 08/11] fix test_generate_and_call_with_valid_key --- litellm/tests/test_key_generate_prisma.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 6be720ffd1..c818b4f4a2 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -209,7 +209,7 @@ def test_generate_and_call_with_valid_key(prisma_client, api_route): await litellm.proxy.proxy_server.prisma_client.connect() from litellm.proxy.proxy_server import user_api_key_cache - request = NewUserRequest(user_role="app_owner") + request = NewUserRequest(user_role=LitellmUserRoles.INTERNAL_USER) key = await new_user(request) print(key) user_id = key.user_id @@ -218,7 +218,7 @@ def test_generate_and_call_with_valid_key(prisma_client, api_route): new_user_info = await user_info(user_id=user_id) new_user_info = new_user_info["user_info"] print("new_user_info=", new_user_info) - assert new_user_info.user_role == "app_owner" + assert new_user_info.user_role == LitellmUserRoles.INTERNAL_USER assert new_user_info.user_id == user_id generated_key = key.key From 3d4cda3373bdb40395318918e2f5a4bf2a5aa7ce Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 15:33:23 -0700 Subject: [PATCH 09/11] fix /user/new --- litellm/proxy/proxy_server.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f3a9b56924..45c2abdf02 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7697,16 +7697,6 @@ async def new_user(data: NewUserRequest): - max_budget: (float|None) Max budget for given user. """ data_json = data.json() # type: ignore - if "user_role" in data_json: - user_role = data_json["user_role"] - if user_role is not None: - if user_role not in ["proxy_admin", "app_owner", "app_user"]: - raise ProxyException( - message=f"Invalid user role, passed in {user_role}. Must be one of 'admin', 'app_owner', 'app_user'", - type="invalid_user_role", - param="user_role", - code=status.HTTP_400_BAD_REQUEST, - ) if "user_id" in data_json and data_json["user_id"] is None: data_json["user_id"] = str(uuid.uuid4()) auto_create_key = data_json.pop("auto_create_key", True) From 338fc82ce89a1ed1a9570647ad2bf7cc3af4b678 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 15:35:33 -0700 Subject: [PATCH 10/11] fix LitellmUserRoles --- litellm/proxy/_types.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e502a5369e..70bc945000 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -42,6 +42,33 @@ class LitellmUserRoles(str, enum.Enum): # Customer Roles - External users of proxy CUSTOMER = "customer" + def __str__(self): + return str(self.value) + + @property + def description(self): + descriptions = { + "proxy_admin": "admin over the platform", + "proxy_admin_view_only": "can login, view all own keys, view all spend", + "internal_user": "internal user can login, view/create/delete their own keys, view their spend", + "internal_user_view_only": "internal user can login, view their own keys, view their own spend", + "team": "team scope used for JWT auth", + "customer": "customer", + } + return descriptions.get(self.value, "") + + @property + def ui_label(self): + ui_labels = { + "proxy_admin": "Admin", + "proxy_admin_view_only": "Admin - View Only", + "internal_user": "Internal User", + "internal_user_view_only": "Internal User - View Only", + "team": "Team", + "customer": "Customer", + } + return ui_labels.get(self.value, "") + AlertType = Literal[ "llm_exceptions", From a3b1493bab6e6e31ba7e5f4e851110d595e52684 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 30 May 2024 16:10:18 -0700 Subject: [PATCH 11/11] add doc string for enum --- litellm/proxy/_types.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 70bc945000..f1bcdf6c84 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -47,11 +47,14 @@ class LitellmUserRoles(str, enum.Enum): @property def description(self): + """ + Descriptions for the enum values + """ descriptions = { - "proxy_admin": "admin over the platform", - "proxy_admin_view_only": "can login, view all own keys, view all spend", - "internal_user": "internal user can login, view/create/delete their own keys, view their spend", - "internal_user_view_only": "internal user can login, view their own keys, view their own spend", + "proxy_admin": "admin over litellm proxy, has all permissions", + "proxy_admin_view_only": "view all keys, view all spend", + "internal_user": "view/create/delete their own keys, view their own spend", + "internal_user_view_only": "view their own keys, view their own spend", "team": "team scope used for JWT auth", "customer": "customer", } @@ -59,6 +62,9 @@ class LitellmUserRoles(str, enum.Enum): @property def ui_label(self): + """ + UI labels for the enum values + """ ui_labels = { "proxy_admin": "Admin", "proxy_admin_view_only": "Admin - View Only",