phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +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": if query_type == "find_all":
response = await self.db.litellm_budgettable.find_many( response = await self.db.litellm_budgettable.find_many(
where={ # type:ignore 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 return response
@ -2499,8 +2507,19 @@ async def reset_budget(prisma_client: PrismaClient):
budget_id_list_to_reset_enduser = [] budget_id_list_to_reset_enduser = []
if budgets_to_reset is not None and len(budgets_to_reset) > 0: if budgets_to_reset is not None and len(budgets_to_reset) > 0:
for budget in budgets_to_reset: 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) 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) budget.budget_reset_at = now + timedelta(seconds=duration_s)
await prisma_client.update_data( await prisma_client.update_data(
query_type="update_many", query_type="update_many",

5
tests/proxy_unit_tests/test_proxy_utils.py
View file

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