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):

11
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"
@ -4523,6 +4524,10 @@ async def new_team(
if data.team_id is None: if data.team_id is None:
data.team_id = str(uuid.uuid4()) data.team_id = str(uuid.uuid4())
if (
user_api_key_dict.user_role is None
or user_api_key_dict.user_role != "proxy_admin"
): # don't restrict proxy admin
if ( if (
data.tpm_limit is not None data.tpm_limit is not None
and user_api_key_dict.tpm_limit is not None and user_api_key_dict.tpm_limit is not None