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

litellm/proxy/_types.py

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

litellm/proxy/proxy_server.py

@@ -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
         is_master_key_valid = secrets.compare_digest(api_key, master_key)
         if is_master_key_valid:
-            return UserAPIKeyAuth(api_key=master_key)
-
+            return UserAPIKeyAuth(api_key=master_key, user_role="proxy_admin")
         if isinstance(
             api_key, str
         ):  # if generated token, make sure it starts with sk-.
@@ -794,7 +793,9 @@ async def user_api_key_auth(
                 pass
             else:
                 if _is_user_proxy_admin(user_id_information):
-                    pass
+                    return UserAPIKeyAuth(
+                        api_key=api_key, user_role="proxy_admin", **valid_token_dict
+                    )
                 else:
                     raise Exception(
                         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:
         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 (
             data.tpm_limit is not None
             and user_api_key_dict.tpm_limit is not None