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

fix(proxy/utils.py): fix reset budget logic

Browse source 
uses fewer clients - prevents read timeouts
This commit is contained in:
Krrish Dholakia 2024-03-19 18:02:09 -07:00
parent 5171d7689f
commit e0d3b18835
1 changed files with 10 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

35
litellm/proxy/utils.py
View file

@ -1820,35 +1820,20 @@ async def reset_budget(prisma_client: PrismaClient):
if prisma_client is not None:
### RESET KEY BUDGET ###
now = datetime.utcnow()
keys_to_reset = await prisma_client.get_data(
table_name="key", query_type="find_all", expires=now, reset_at=now
)
if keys_to_reset is not None and len(keys_to_reset) > 0:
for key in keys_to_reset:
key.spend = 0.0
duration_s = _duration_in_seconds(duration=key.budget_duration)
key.budget_reset_at = now + timedelta(seconds=duration_s)
await prisma_client.update_data(
query_type="update_many", data_list=keys_to_reset, table_name="key"
asyncio.create_task(
prisma_client.db.litellm_verificationtoken.update_many(
where={"budget_reset_at": {"lt": now}},
data={"spend": 0, "budget_reset_at": now},
)
)
### RESET USER BUDGET ###
now = datetime.utcnow()
users_to_reset = await prisma_client.get_data(
table_name="user", query_type="find_all", reset_at=now
)
if users_to_reset is not None and len(users_to_reset) > 0:
for user in users_to_reset:
user.spend = 0.0
duration_s = _duration_in_seconds(duration=user.budget_duration)
user.budget_reset_at = now + timedelta(seconds=duration_s)
await prisma_client.update_data(
query_type="update_many", data_list=users_to_reset, table_name="user"
asyncio.create_task(
prisma_client.db.litellm_usertable.update_many(
where={"budget_reset_at": {"lt": now}},
data={"spend": 0, "budget_reset_at": now},
)
)
async def update_spend(