phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix migrate to LitellmUserRoles

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-30 13:57:50 -07:00
parent b29adeb6c6
commit 6166381579
4 changed files with 47 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/utils.py
View file

@ -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

14
litellm/tests/test_add_update_models.py
View file

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

44
litellm/tests/test_key_generate_prisma.py
View file

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

5
tests/test_keys.py
View file

@ -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,