diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b96c469ec..86131991b 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 ec9ab33b6..21124a4ba 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 375a1c85d..7ceab3663 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 11961e2a2..bbe4cc35e 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,