phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat - move to using enums for roles

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-30 13:55:45 -07:00
parent fc54acc17a
commit b29adeb6c6
1 changed files with 46 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

79
litellm/proxy/proxy_server.py
View file

@ -524,12 +524,14 @@ async def user_api_key_auth(
if is_admin: if is_admin:
# check allowed admin routes # check allowed admin routes
is_allowed = allowed_routes_check( is_allowed = allowed_routes_check(
user_role="proxy_admin", user_role=LitellmUserRoles.PROXY_ADMIN.value,
user_route=route, user_route=route,
litellm_proxy_roles=jwt_handler.litellm_jwtauth, litellm_proxy_roles=jwt_handler.litellm_jwtauth,
) )
if is_allowed: if is_allowed:
return UserAPIKeyAuth(user_role="proxy_admin") return UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN.value
)
else: else:
allowed_routes = ( allowed_routes = (
jwt_handler.litellm_jwtauth.admin_allowed_routes jwt_handler.litellm_jwtauth.admin_allowed_routes
@ -671,9 +673,11 @@ async def user_api_key_auth(
#### ELSE #### #### ELSE ####
if master_key is None: if master_key is None:
if isinstance(api_key, str): 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: 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 elif api_key is None: # only require api key if master key is set
raise Exception("No api key passed in.") raise Exception("No api key passed in.")
elif api_key == "": elif api_key == "":
@ -740,7 +744,7 @@ async def user_api_key_auth(
if ( if (
valid_token is not None valid_token is not None
and isinstance(valid_token, UserAPIKeyAuth) 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 # update end-user params on valid token
valid_token.end_user_id = end_user_params.get("end_user_id") 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: if is_master_key_valid:
_user_api_key_obj = UserAPIKeyAuth( _user_api_key_obj = UserAPIKeyAuth(
api_key=master_key, api_key=master_key,
user_role="proxy_admin", user_role=LitellmUserRoles.PROXY_ADMIN.value,
user_id=litellm_proxy_admin_name, user_id=litellm_proxy_admin_name,
**end_user_params, **end_user_params,
) )
@ -1377,7 +1381,9 @@ async def user_api_key_auth(
user_id_information user_id_information
): ):
return UserAPIKeyAuth( 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 ( elif (
_has_user_setup_sso() _has_user_setup_sso()
@ -1398,15 +1404,21 @@ async def user_api_key_auth(
user_id_information user_id_information
): ):
return UserAPIKeyAuth( 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: elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value:
return UserAPIKeyAuth( 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: else:
return UserAPIKeyAuth( 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: else:
raise Exception() raise Exception()
@ -3738,11 +3750,9 @@ async def startup_event():
spend=0, spend=0,
token=master_key, token=master_key,
user_id=litellm_proxy_admin_name, user_id=litellm_proxy_admin_name,
user_role="proxy_admin", user_role=LitellmUserRoles.PROXY_ADMIN.value,
query_type="update_data", query_type="update_data",
update_key_values={ update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
"user_role": "proxy_admin",
},
) )
) )
@ -6093,7 +6103,7 @@ async def delete_key_fn(
) )
if ( if (
user_api_key_dict.user_role is not None 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 user_id = None # unless they're admin
@ -7888,7 +7898,10 @@ async def user_info(
user_id=user_api_key_dict.user_id user_id=user_api_key_dict.user_id
) )
# *NEW* get all teams in user 'teams' field # *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( teams_2 = await prisma_client.get_data(
table_name="team", table_name="team",
query_type="find_all", query_type="find_all",
@ -8716,7 +8729,7 @@ async def new_team(
if ( if (
user_api_key_dict.user_role is None 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 ): # don't restrict proxy admin
if ( if (
data.tpm_limit is not None data.tpm_limit is not None
@ -9322,7 +9335,7 @@ async def list_team(
""" """
global prisma_client 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( raise HTTPException(
status_code=401, status_code=401,
detail={ detail={
@ -9416,7 +9429,7 @@ async def new_organization(
if ( if (
user_api_key_dict.user_role is None 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( raise HTTPException(
status_code=401, status_code=401,
@ -9619,7 +9632,7 @@ async def budget_settings(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -9684,7 +9697,7 @@ async def list_budget(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -9718,7 +9731,7 @@ async def delete_budget(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -10696,7 +10709,7 @@ async def alerting_settings(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -10777,7 +10790,7 @@ async def alerting_settings(
# detail={"error": CommonProxyErrors.db_not_connected_error.value}, # 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( # raise HTTPException(
# status_code=400, # status_code=400,
# detail={"error": CommonProxyErrors.not_allowed_access.value}, # detail={"error": CommonProxyErrors.not_allowed_access.value},
@ -11235,12 +11248,12 @@ async def login(request: Request):
await user_update( await user_update(
data=UpdateUserRequest( data=UpdateUserRequest(
user_id=key_user_id, user_id=key_user_id,
user_role="proxy_admin", user_role=LitellmUserRoles.PROXY_ADMIN.value,
) )
) )
if os.getenv("DATABASE_URL") is not None: if os.getenv("DATABASE_URL") is not None:
response = await generate_key_helper_fn( 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: else:
raise ProxyException( raise ProxyException(
@ -11635,7 +11648,7 @@ async def new_invitation(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -11699,7 +11712,7 @@ async def invitation_info(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -11811,7 +11824,7 @@ async def invitation_delete(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -12006,7 +12019,7 @@ async def update_config_general_settings(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={"error": CommonProxyErrors.not_allowed_access.value}, 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}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={"error": CommonProxyErrors.not_allowed_access.value}, detail={"error": CommonProxyErrors.not_allowed_access.value},
@ -12143,7 +12156,7 @@ async def get_config_list(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
@ -12218,7 +12231,7 @@ async def delete_config_general_settings(
detail={"error": CommonProxyErrors.db_not_connected_error.value}, 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( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={