fix migrate to LitellmUserRoles

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
 

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",
         ),
     )
 

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",
         ),
     )
 

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,