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

feat: adding fallback for already running applications with created budgets

Browse source
This commit is contained in:
Laurien Lummer 2025-02-12 15:15:14 +01:00
parent 9f12cba8bc
commit adfe462f3e
2 changed files with 24 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/utils.py
View file

@ -1489,7 +1489,15 @@ class PrismaClient:
if query_type == "find_all":
response = await self.db.litellm_budgettable.find_many(
where={ # type:ignore
"budget_reset_at": {"lt": reset_at}
"OR": [
{
"AND": [
{"budget_reset_at": None},
{"NOT": {"budget_duration": None}},
]
},
{"budget_reset_at": {"lt": reset_at}},
]
}
)
return response
@ -2499,8 +2507,19 @@ async def reset_budget(prisma_client: PrismaClient):
budget_id_list_to_reset_enduser = []
if budgets_to_reset is not None and len(budgets_to_reset) > 0:
for budget in budgets_to_reset:
budget_id_list_to_reset_enduser.append(budget.budget_id)
duration_s = duration_in_seconds(duration=budget.budget_duration)
# Fallback for existing budgets that do not have a budget_reset_at date set, ensuring the duration is taken into account
if (
budget.budget_reset_at is None
and budget.created_at + timedelta(seconds=duration_s) > now
):
budget.budget_reset_at = budget.created_at + timedelta(
seconds=duration_s
)
continue
budget_id_list_to_reset_enduser.append(budget.budget_id)
budget.budget_reset_at = now + timedelta(seconds=duration_s)
await prisma_client.update_data(
query_type="update_many",