From 58bdbecbc6b8240c330782de09b0a6c23a8c94b0 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 28 Mar 2024 15:59:35 -0700
Subject: [PATCH] fix(proxy_server.py): allow updating tpm/rpm limits via
 /team/update

---
 litellm/proxy/_types.py       |  7 +------
 litellm/proxy/proxy_server.py |  1 -
 tests/test_team.py            | 20 +++++++++-----------
 3 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 577231c26..757f89d47 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -427,13 +427,8 @@ class TeamMemberDeleteRequest(LiteLLMBase):
         return values
 
 
-class UpdateTeamRequest(LiteLLMBase):
+class UpdateTeamRequest(TeamBase):
     team_id: str  # required
-    team_alias: Optional[str] = None
-    admins: Optional[list] = None
-    members: Optional[list] = None
-    members_with_roles: Optional[List[Member]] = None
-    metadata: Optional[dict] = None
 
 
 class DeleteTeamRequest(LiteLLMBase):
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 6682686c6..9c4f04017 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -5826,7 +5826,6 @@ async def update_team(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
-    [BETA]
     [RECOMMENDED] - use `/team/member_add` to add new team members instead 
 
     You can now update team budget / rate limits via /team/update
diff --git a/tests/test_team.py b/tests/test_team.py
index f0ef0bb22..37ee69dd5 100644
--- a/tests/test_team.py
+++ b/tests/test_team.py
@@ -157,18 +157,10 @@ async def new_team(session, i, user_id=None, member_list=None, model_aliases=Non
         return await response.json()
 
 
-async def update_team(
-    session,
-    i,
-    team_id,
-    user_id=None,
-    member_list=None,
-):
+async def update_team(session, i, team_id, user_id=None, member_list=None, **kwargs):
     url = "http://0.0.0.0:4000/team/update"
     headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
-    data = {
-        "team_id": team_id,
-    }
+    data = {"team_id": team_id, **kwargs}
     if user_id is not None:
         data["members_with_roles"] = [{"role": "user", "user_id": user_id}]
     elif member_list is not None:
@@ -284,9 +276,15 @@ async def test_team_update():
             {"role": "user", "user_id": new_normal_user},
         ]
         team_data = await update_team(
-            session=session, i=0, member_list=member_list, team_id=team_data["team_id"]
+            session=session,
+            i=0,
+            member_list=member_list,
+            team_id=team_data["team_id"],
+            tpm_limit=100,
         )
 
+        assert team_data["data"]["tpm_limit"] == 100
+
 
 @pytest.mark.asyncio
 async def test_team_delete():