fix - updating team

litellm/proxy/_types.py

@@ -612,7 +612,6 @@ class UpdateTeamRequest(LiteLLMBase):
     models: Optional[list] = None
     blocked: Optional[bool] = None
     budget_duration: Optional[str] = None
-    budget_reset_at: Optional[datetime] = None
 
 
 class DeleteTeamRequest(LiteLLMBase):

litellm/proxy/proxy_server.py

@@ -8121,16 +8121,16 @@ async def update_team(
             detail={"error": f"Team not found, passed team_id={data.team_id}"},
         )
 
+    updated_kv = data.json(exclude_none=True)
+
     # Check budget_duration and budget_reset_at
     if data.budget_duration is not None:
         duration_s = _duration_in_seconds(duration=data.budget_duration)
         reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
 
-        # set the budget duration and budget_reset_at in DB
-        data.budget_duration = duration_s
-        data.budget_reset_at = reset_at
+        # set the budget_reset_at in DB
+        updated_kv["budget_reset_at"] = reset_at
 
-    updated_kv = data.json(exclude_none=True)
     team_row = await prisma_client.update_data(
         update_key_values=updated_kv,
         data=updated_kv,

litellm/tests/test_key_generate_prisma.py

@@ -52,6 +52,7 @@ from litellm.proxy.proxy_server import (
     user_info,
     info_key_fn,
     new_team,
+    update_team,
     chat_completion,
     completion,
     embeddings,
@@ -73,6 +74,7 @@ from litellm.proxy._types import (
     UpdateKeyRequest,
     GenerateKeyRequest,
     NewTeamRequest,
+    UpdateTeamRequest,
     UserAPIKeyAuth,
     LiteLLM_UpperboundKeyGenerateParams,
 )
@@ -2185,3 +2187,31 @@ async def test_create_update_team(prisma_client):
     assert response["budget_reset_at"] is not None and isinstance(
         response["budget_reset_at"], datetime.datetime
     )
+
+    # updating team budget duration and reset at
+
+    response = await update_team(
+        UpdateTeamRequest(
+            team_id=_team_id,
+            max_budget=30,
+            budget_duration="2d",
+            tpm_limit=30,
+            rpm_limit=30,
+        ),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+        ),
+    )
+
+    print("RESPONSE from update_team", response)
+    _updated_info = response["data"]
+    _updated_info = dict(_updated_info)
+
+    assert _updated_info["team_id"] == _team_id
+    assert _updated_info["max_budget"] == 30
+    assert _updated_info["tpm_limit"] == 30
+    assert _updated_info["rpm_limit"] == 30
+    assert _updated_info["budget_duration"] == "2d"
+    assert _updated_info["budget_reset_at"] is not None and isinstance(
+        _updated_info["budget_reset_at"], datetime.datetime
+    )