feat - move to using enums for roles

This commit is contained in:
Ishaan Jaff 2024-05-30 13:55:45 -07:00
parent fc54acc17a
commit b29adeb6c6

View file

@ -524,12 +524,14 @@ 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.value,
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.value
)
else:
allowed_routes = (
jwt_handler.litellm_jwtauth.admin_allowed_routes
@ -671,9 +673,11 @@ async def user_api_key_auth(
#### 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.value
)
else:
return UserAPIKeyAuth(user_role="proxy_admin")
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value)
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 +744,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.value
):
# update end-user params on valid token
valid_token.end_user_id = end_user_params.get("end_user_id")
@ -773,7 +777,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.value,
user_id=litellm_proxy_admin_name,
**end_user_params,
)
@ -1377,7 +1381,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.value,
**valid_token_dict,
)
elif (
_has_user_setup_sso()
@ -1398,15 +1404,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.value,
**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.value,
**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.value,
**valid_token_dict,
)
else:
raise Exception()
@ -3738,11 +3750,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.value,
query_type="update_data",
update_key_values={
"user_role": "proxy_admin",
},
update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
)
)
@ -6093,7 +6103,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.value
):
user_id = None # unless they're admin
@ -7888,7 +7898,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.value
):
teams_2 = await prisma_client.get_data(
table_name="team",
query_type="find_all",
@ -8716,7 +8729,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.value
): # don't restrict proxy admin
if (
data.tpm_limit is not None
@ -9322,7 +9335,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.value:
raise HTTPException(
status_code=401,
detail={
@ -9416,7 +9429,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.value
):
raise HTTPException(
status_code=401,
@ -9619,7 +9632,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.value:
raise HTTPException(
status_code=400,
detail={
@ -9684,7 +9697,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.value:
raise HTTPException(
status_code=400,
detail={
@ -9718,7 +9731,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.value:
raise HTTPException(
status_code=400,
detail={
@ -10696,7 +10709,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.value:
raise HTTPException(
status_code=400,
detail={
@ -10777,7 +10790,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.value:
# raise HTTPException(
# status_code=400,
# detail={"error": CommonProxyErrors.not_allowed_access.value},
@ -11235,12 +11248,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.value,
)
)
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.value, "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 +11648,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.value:
raise HTTPException(
status_code=400,
detail={
@ -11699,7 +11712,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.value:
raise HTTPException(
status_code=400,
detail={
@ -11811,7 +11824,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.value:
raise HTTPException(
status_code=400,
detail={
@ -12006,7 +12019,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.value:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.not_allowed_access.value},
@ -12080,7 +12093,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.value:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.not_allowed_access.value},
@ -12143,7 +12156,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.value:
raise HTTPException(
status_code=400,
detail={
@ -12218,7 +12231,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.value:
raise HTTPException(
status_code=400,
detail={