From a03d13feddd92f67e40c97f20eaa7da6e595a5c7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 14:38:05 -0700 Subject: [PATCH] fix - allow setting budget duration and reset at --- litellm/proxy/proxy_server.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 189b6cb88..4bac6ff3c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8033,7 +8033,16 @@ async def new_team( _model_id = model_dict.id - ## ADD TO TEAM TABLE -> the user api key settings are overriding the data we passed i + ## ADD TO TEAM TABLE + # Check budget_duration and budget_reset_at + if data.budget_duration: + 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 + complete_team_data = LiteLLM_TeamTable( **data.json(), model_id=_model_id,