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

fix: enable proxy admin to set budgets for teams on ui

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-27 11:20:21 -08:00
parent ac9a443fdf
commit 7250c86fbe
2 changed files with 50 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -477,6 +477,7 @@ class UserAPIKeyAuth(
""" """
api_key: Optional[str] = None api_key: Optional[str] = None
user_role: Optional[Literal["proxy_admin", "app_owner", "app_user"]] = None
@root_validator(pre=True) @root_validator(pre=True)
def check_api_key(cls, values): def check_api_key(cls, values):

93
litellm/proxy/proxy_server.py
View file

@ -332,8 +332,7 @@ async def user_api_key_auth(
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
is_master_key_valid = secrets.compare_digest(api_key, master_key) is_master_key_valid = secrets.compare_digest(api_key, master_key)
if is_master_key_valid: if is_master_key_valid:
return UserAPIKeyAuth(api_key=master_key) return UserAPIKeyAuth(api_key=master_key, user_role="proxy_admin")
if isinstance( if isinstance(
api_key, str api_key, str
): # if generated token, make sure it starts with sk-. ): # if generated token, make sure it starts with sk-.
@ -794,7 +793,9 @@ async def user_api_key_auth(
pass pass
else: else:
if _is_user_proxy_admin(user_id_information): if _is_user_proxy_admin(user_id_information):
pass return UserAPIKeyAuth(
api_key=api_key, user_role="proxy_admin", **valid_token_dict
)
else: else:
raise Exception( raise Exception(
f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed" f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
@ -4524,50 +4525,54 @@ async def new_team(
data.team_id = str(uuid.uuid4()) data.team_id = str(uuid.uuid4())
if ( if (
data.tpm_limit is not None user_api_key_dict.user_role is None
and user_api_key_dict.tpm_limit is not None or user_api_key_dict.user_role != "proxy_admin"
and data.tpm_limit > user_api_key_dict.tpm_limit ): # don't restrict proxy admin
): if (
raise HTTPException( data.tpm_limit is not None
status_code=400, and user_api_key_dict.tpm_limit is not None
detail={ and data.tpm_limit > user_api_key_dict.tpm_limit
"error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}" ):
}, raise HTTPException(
) status_code=400,
detail={
"error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}"
},
)
if ( if (
data.rpm_limit is not None data.rpm_limit is not None
and user_api_key_dict.rpm_limit is not None and user_api_key_dict.rpm_limit is not None
and data.rpm_limit > user_api_key_dict.rpm_limit and data.rpm_limit > user_api_key_dict.rpm_limit
): ):
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
"error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}" "error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}"
}, },
) )
if ( if (
data.max_budget is not None data.max_budget is not None
and user_api_key_dict.max_budget is not None and user_api_key_dict.max_budget is not None
and data.max_budget > user_api_key_dict.max_budget and data.max_budget > user_api_key_dict.max_budget
): ):
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
"error": f"max budget higher than user max. User max budget={user_api_key_dict.max_budget}" "error": f"max budget higher than user max. User max budget={user_api_key_dict.max_budget}"
}, },
) )
if data.models is not None: if data.models is not None:
for m in data.models: for m in data.models:
if m not in user_api_key_dict.models: if m not in user_api_key_dict.models:
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
detail={ detail={
"error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}" "error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}"
}, },
) )
if user_api_key_dict.user_id is not None: if user_api_key_dict.user_id is not None:
creating_user_in_list = False creating_user_in_list = False