fix update_end_user_spend

This commit is contained in:
Ishaan Jaff 2025-03-31 19:05:58 -07:00
parent 73d6af5572
commit 4959512170
2 changed files with 6 additions and 12 deletions

View file

@ -514,6 +514,7 @@ class DBSpendUpdateWriter:
n_retry_times=n_retry_times, n_retry_times=n_retry_times,
prisma_client=prisma_client, prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj, proxy_logging_obj=proxy_logging_obj,
end_user_list_transactions=end_user_list_transactions,
) )
### UPDATE KEY TABLE ### ### UPDATE KEY TABLE ###
key_list_transactions = db_spend_update_transactions["key_list_transactions"] key_list_transactions = db_spend_update_transactions["key_list_transactions"]

View file

@ -1100,12 +1100,6 @@ def jsonify_object(data: dict) -> dict:
class PrismaClient: class PrismaClient:
user_list_transactions: dict = {}
end_user_list_transactions: dict = {}
key_list_transactions: dict = {}
team_list_transactions: dict = {}
team_member_list_transactions: dict = {} # key is ["team_id" + "user_id"]
org_list_transactions: dict = {}
spend_log_transactions: List = [] spend_log_transactions: List = []
daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {} daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {}
@ -2422,7 +2416,10 @@ def _hash_token_if_needed(token: str) -> str:
class ProxyUpdateSpend: class ProxyUpdateSpend:
@staticmethod @staticmethod
async def update_end_user_spend( async def update_end_user_spend(
n_retry_times: int, prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
end_user_list_transactions: Dict[str, float],
): ):
for i in range(n_retry_times + 1): for i in range(n_retry_times + 1):
start_time = time.time() start_time = time.time()
@ -2434,7 +2431,7 @@ class ProxyUpdateSpend:
for ( for (
end_user_id, end_user_id,
response_cost, response_cost,
) in prisma_client.end_user_list_transactions.items(): ) in end_user_list_transactions.items():
if litellm.max_end_user_budget is not None: if litellm.max_end_user_budget is not None:
pass pass
batcher.litellm_endusertable.upsert( batcher.litellm_endusertable.upsert(
@ -2461,10 +2458,6 @@ class ProxyUpdateSpend:
_raise_failed_update_spend_exception( _raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
) )
finally:
prisma_client.end_user_list_transactions = (
{}
) # reset the end user list transactions - prevent bad data from causing issues
@staticmethod @staticmethod
async def update_spend_logs( async def update_spend_logs(