diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 555254a633..f1bcdf6c84 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -7,6 +7,75 @@ import uuid, json, sys, os from litellm.types.router import UpdateRouterConfig from litellm.types.utils import ProviderField + +class LitellmUserRoles(str, enum.Enum): + """ + 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 + + """ + + # 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" + + def __str__(self): + return str(self.value) + + @property + def description(self): + """ + Descriptions for the enum values + """ + descriptions = { + "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", + } + return descriptions.get(self.value, "") + + @property + def ui_label(self): + """ + UI labels for the enum values + """ + 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", "llm_too_slow", @@ -498,7 +567,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 = ( @@ -517,7 +595,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) @@ -1062,7 +1149,16 @@ class UserAPIKeyAuth( """ api_key: Optional[str] = None - user_role: Optional[Literal["proxy_admin", "app_owner", "app_user"]] = 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 allowed_model_region: Optional[Literal["eu"]] = None @root_validator(pre=True) 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 f0b98d0be8..45c2abdf02 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -507,7 +507,7 @@ 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) if general_settings.get("enable_jwt_auth", False) == True: is_jwt = jwt_handler.is_jwt(token=api_key) @@ -524,12 +524,12 @@ 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, 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) else: allowed_routes = ( jwt_handler.litellm_jwtauth.admin_allowed_routes @@ -552,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, ) @@ -664,16 +664,18 @@ 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, user_id=user_id, org_id=org_id, ) #### 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 + ) else: - return UserAPIKeyAuth(user_role="proxy_admin") + 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 == "": @@ -740,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 == "proxy_admin" + 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") @@ -773,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="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN, user_id=litellm_proxy_admin_name, **end_user_params, ) @@ -1377,7 +1379,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, + **valid_token_dict, ) elif ( _has_user_setup_sso() @@ -1398,15 +1402,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, + **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, + **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, + **valid_token_dict, ) else: raise Exception() @@ -3738,11 +3748,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, query_type="update_data", - update_key_values={ - "user_role": "proxy_admin", - }, + update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN}, ) ) @@ -6093,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 == "proxy_admin" + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ): user_id = None # unless they're admin @@ -7689,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) @@ -7888,7 +7886,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 + ): teams_2 = await prisma_client.get_data( table_name="team", query_type="find_all", @@ -8716,7 +8717,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 ): # don't restrict proxy admin if ( data.tpm_limit is not None @@ -9322,7 +9323,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: raise HTTPException( status_code=401, detail={ @@ -9416,7 +9417,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 ): raise HTTPException( status_code=401, @@ -9619,7 +9620,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: raise HTTPException( status_code=400, detail={ @@ -9684,7 +9685,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: raise HTTPException( status_code=400, detail={ @@ -9718,7 +9719,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: raise HTTPException( status_code=400, detail={ @@ -10696,7 +10697,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: raise HTTPException( status_code=400, detail={ @@ -10777,7 +10778,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: # raise HTTPException( # status_code=400, # detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -11235,12 +11236,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, ) ) 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, "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 +11636,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: raise HTTPException( status_code=400, detail={ @@ -11699,7 +11700,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: raise HTTPException( status_code=400, detail={ @@ -11811,7 +11812,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: raise HTTPException( status_code=400, detail={ @@ -12006,7 +12007,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: raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -12080,7 +12081,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: raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -12143,7 +12144,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: raise HTTPException( status_code=400, detail={ @@ -12218,7 +12219,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: raise HTTPException( status_code=400, detail={ 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..c818b4f4a2 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, + api_key="sk-1234", + user_id="1234", ), ) @@ -206,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 @@ -215,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 @@ -363,7 +366,8 @@ 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), ) print("new_team_response", new_team_response) created_team_id = new_team_response["team_id"] @@ -922,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 = "proxy_admin" + 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 @@ -972,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 = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -1044,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 = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -1078,7 +1082,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, + api_key="sk-1234", + user_id="1234", ), ) @@ -1088,7 +1094,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, + api_key="sk-1234", + user_id="1234", ), ) @@ -1158,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 = "proxy_admin" + result.user_role = LitellmUserRoles.PROXY_ADMIN result_delete_key = await delete_key_fn( data=delete_key_request, user_api_key_dict=result @@ -2038,7 +2046,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, + api_key="sk-1234", + user_id="1234", ), ) @@ -2076,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 == "proxy_admin" or admin gets rejected + 3. Test that non-master key with role == LitellmUserRoles.PROXY_ADMIN or admin gets rejected """ print("prisma client=", prisma_client) @@ -2121,10 +2131,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 or admin _response = await new_user( data=NewUserRequest( - user_role="proxy_admin", + user_role=LitellmUserRoles.PROXY_ADMIN, tpm_limit=20, ) ) @@ -2174,7 +2184,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, + api_key="sk-1234", + user_id="1234", ), ) @@ -2200,7 +2212,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, + 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,