tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-15 18:24:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #3927 from BerriAI/litellm_use_enums_for_user_roles

Browse Source 
[Feat] LiteLLM Proxy - use enums for user roles
This commit is contained in:
Ishaan Jaff
2024-05-30 16:14:17 -07:00
committed by GitHub
parent 61454eaa0f a3b1493bab
commit f64d547542
7 changed files with 201 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+99 -3
View File
@@ -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)
litellm/proxy/auth/auth_checks.py
+8 -3
View File
@@ -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
litellm/proxy/proxy_server.py
+47 -46
View File
@@ -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={
litellm/proxy/utils.py
+3 -2
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
litellm/tests/test_add_update_models.py
+10 -4
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",
),
)
litellm/tests/test_key_generate_prisma.py
+30 -16
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,
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",
),
)
tests/test_keys.py
+4 -1
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,