diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 11e358a2f8..0ea6fff8e8 100644 --- a/litellm/proxy/_types.py +++ b/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): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6c6a01f45b..8bfc3314a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/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, diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 3619ce8f49..dda5c365cb 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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 + )