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

add reset enduser spend logic to new reset_budget function

Browse source
This commit is contained in:
Laurien Lummer 2025-03-19 17:24:12 +01:00
parent 7dc0d29fe9
commit 7285534903
2 changed files with 81 additions and 112 deletions
Download patch file Download diff file Expand all files Collapse all files

111
litellm/proxy/utils.py
View file

@ -2447,117 +2447,6 @@ def _hash_token_if_needed(token: str) -> str:
return token
async def reset_budget(prisma_client: PrismaClient):
"""
Gets all the non-expired keys for a db, which need spend to be reset
Resets their spend
Updates db
"""
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"
)
### 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"
)
## Reset End-User (Customer) Spend and corresponding Budget duration
now = datetime.now(timezone.utc)
budgets_to_reset = await prisma_client.get_data(
table_name="budget", query_type="find_all", reset_at=now
)
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:
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",
data_list=budgets_to_reset,
table_name="budget",
)
endusers_to_reset = await prisma_client.get_data(
table_name="enduser",
query_type="find_all",
budget_id_list=budget_id_list_to_reset_enduser,
)
if endusers_to_reset is not None and len(endusers_to_reset) > 0:
for enduser in endusers_to_reset:
enduser.spend = 0.0
await prisma_client.update_data(
query_type="update_many",
data_list=endusers_to_reset,
table_name="enduser",
)
## Reset Team Budget
now = datetime.utcnow()
teams_to_reset = await prisma_client.get_data(
table_name="team",
query_type="find_all",
reset_at=now,
)
if teams_to_reset is not None and len(teams_to_reset) > 0:
team_reset_requests = []
for team in teams_to_reset:
duration_s = duration_in_seconds(duration=team.budget_duration)
reset_team_budget_request = ResetTeamBudgetRequest(
team_id=team.team_id,
spend=0.0,
budget_reset_at=now + timedelta(seconds=duration_s),
updated_at=now,
)
team_reset_requests.append(reset_team_budget_request)
await prisma_client.update_data(
query_type="update_many",
data_list=team_reset_requests,
table_name="team",
)
class ProxyUpdateSpend:
@staticmethod
async def update_end_user_spend(