phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #3842 from BerriAI/litellm_set_budget_dur

Browse source 
[Fix] Set budget_duration on `/team/new` and `/team/update`
This commit is contained in:
Ishaan Jaff 2024-05-25 16:31:39 -07:00 • committed by GitHub
commit e4053b6732
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 116 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

18
litellm/proxy/proxy_server.py
View file

@ -8040,12 +8040,15 @@ async def new_team(
## ADD TO TEAM TABLE
complete_team_data = LiteLLM_TeamTable(
**data.json(),
max_parallel_requests=user_api_key_dict.max_parallel_requests,
budget_duration=user_api_key_dict.budget_duration,
budget_reset_at=user_api_key_dict.budget_reset_at,
model_id=_model_id,
)
# If budget_duration is set, set `budget_reset_at`
if complete_team_data.budget_duration is not None:
duration_s = _duration_in_seconds(duration=complete_team_data.budget_duration)
reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
complete_team_data.budget_reset_at = reset_at
team_row = await prisma_client.insert_data(
data=complete_team_data.json(exclude_none=True), table_name="team"
)
@ -8124,6 +8127,15 @@ async def update_team(
)
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_reset_at in DB
updated_kv["budget_reset_at"] = reset_at
team_row = await prisma_client.update_data(
update_key_values=updated_kv,
data=updated_kv,