phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat - update budget duration

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-25 14:57:29 -07:00
parent a03d13fedd
commit 6f65df0b1a
2 changed files with 12 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

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

11
litellm/proxy/proxy_server.py
View file

@ -8035,7 +8035,7 @@ async def new_team(
## ADD TO TEAM TABLE ## ADD TO TEAM TABLE
# Check budget_duration and budget_reset_at # Check budget_duration and budget_reset_at
if data.budget_duration: if data.budget_duration is not None:
duration_s = _duration_in_seconds(duration=data.budget_duration) duration_s = _duration_in_seconds(duration=data.budget_duration)
reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
@ -8125,6 +8125,15 @@ async def update_team(
detail={"error": f"Team not found, passed team_id={data.team_id}"}, detail={"error": f"Team not found, passed team_id={data.team_id}"},
) )
# 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
updated_kv = data.json(exclude_none=True) updated_kv = data.json(exclude_none=True)
team_row = await prisma_client.update_data( team_row = await prisma_client.update_data(
update_key_values=updated_kv, update_key_values=updated_kv,