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

(Bug Fix + Better Observability) - BudgetResetJob: (#8562)

Browse source 
* use class ResetBudgetJob

* refactor reset budget job

* update reset_budget job

* refactor reset budget job

* fix LiteLLM_UserTable

* refactor reset budget job

* add telemetry for reset budget job

* dd - log service success/failure on DD

* add detailed reset budget reset info on DD

* initialize_scheduled_background_jobs

* refactor reset budget job

* trigger service failure hook when fails to reset a budget for team, key, user

* fix resetBudgetJob

* unit testing for ResetBudgetJob

* test_duration_in_seconds_basic

* testing for triggering service logging

* fix logs on test teams fail

* remove unused imports

* fix import duration in s

* duration_in_seconds
This commit is contained in:
Ishaan Jaff 2025-02-15 16:13:08 -08:00 • committed by GitHub
parent 5a6ccb0c48
commit 2753de1458
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1107 additions and 87 deletions
Download patch file Download diff file Expand all files Collapse all files

69
litellm/proxy/utils.py
View file

@ -13,7 +13,6 @@ from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union, overload
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
CommonProxyErrors,
@ -49,7 +48,6 @@ from litellm.proxy._types import (
CallInfo,
LiteLLM_VerificationTokenView,
Member,
ResetTeamBudgetRequest,
UserAPIKeyAuth,
)
from litellm.proxy.db.create_views import (
@ -2363,73 +2361,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 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(