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

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",

5
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -3,7 +3,7 @@ import json
import os
import sys
import uuid
from datetime import datetime
from datetime import datetime, timezone
from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@ -1639,6 +1639,7 @@ def test_provider_specific_header():
},
}
async def create_budget(session, data):
url = "http://0.0.0.0:4000/budget/new"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
@ -1831,7 +1832,7 @@ async def test_reset_budget_for_endusers(prisma_client, budget_and_enduser_setup
updated_budgets = await prisma_client.get_data(
table_name="budget",
query_type="find_all",
reset_at=datetime.now() + timedelta(days=31),
reset_at=datetime.now(timezone.utc) + timedelta(days=31),
)
# Assertions for end users