diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 0183fa433b..7184262ee5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -201,7 +201,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=hashed_token, - transaction_list=prisma_client.key_list_transactons, + transaction_list=prisma_client.key_list_transactions, entity_type=Litellm_EntityType.KEY, debug_msg=f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}.", ) @@ -241,7 +241,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=_id, - transaction_list=prisma_client.user_list_transactons, + transaction_list=prisma_client.user_list_transactions, entity_type=Litellm_EntityType.USER, ) @@ -249,7 +249,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=end_user_id, - transaction_list=prisma_client.end_user_list_transactons, + transaction_list=prisma_client.end_user_list_transactions, entity_type=Litellm_EntityType.END_USER, ) except Exception as e: @@ -275,7 +275,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=team_id, - transaction_list=prisma_client.team_list_transactons, + transaction_list=prisma_client.team_list_transactions, entity_type=Litellm_EntityType.TEAM, ) @@ -287,7 +287,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=team_member_key, - transaction_list=prisma_client.team_member_list_transactons, + transaction_list=prisma_client.team_member_list_transactions, entity_type=Litellm_EntityType.TEAM_MEMBER, ) except Exception: @@ -314,7 +314,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=org_id, - transaction_list=prisma_client.org_list_transactons, + transaction_list=prisma_client.org_list_transactions, entity_type=Litellm_EntityType.ORGANIZATION, ) except Exception as e: @@ -424,12 +424,12 @@ class DBSpendUpdateWriter: await self.pod_leader_manager.release_lock() else: db_spend_update_transactions = DBSpendUpdateTransactions( - user_list_transactons=prisma_client.user_list_transactons, - end_user_list_transactons=prisma_client.end_user_list_transactons, - key_list_transactons=prisma_client.key_list_transactons, - team_list_transactons=prisma_client.team_list_transactons, - team_member_list_transactons=prisma_client.team_member_list_transactons, - org_list_transactons=prisma_client.org_list_transactons, + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, ) await DBSpendUpdateWriter._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -455,8 +455,8 @@ class DBSpendUpdateWriter: ) ### UPDATE USER TABLE ### - user_list_transactons = db_spend_update_transactions["user_list_transactons"] - if len(user_list_transactons.keys()) > 0: + user_list_transactions = db_spend_update_transactions["user_list_transactions"] + if len(user_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -467,12 +467,12 @@ class DBSpendUpdateWriter: for ( user_id, response_cost, - ) in user_list_transactons.items(): + ) in user_list_transactions.items(): batcher.litellm_usertable.update_many( where={"user_id": user_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.user_list_transactons = ( + prisma_client.user_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -495,24 +495,24 @@ class DBSpendUpdateWriter: ### UPDATE END-USER TABLE ### verbose_proxy_logger.debug( "End-User Spend transactions: {}".format( - len(prisma_client.end_user_list_transactons.keys()) + len(prisma_client.end_user_list_transactions.keys()) ) ) - end_user_list_transactons = db_spend_update_transactions[ - "end_user_list_transactons" + end_user_list_transactions = db_spend_update_transactions[ + "end_user_list_transactions" ] - if len(end_user_list_transactons.keys()) > 0: + if len(end_user_list_transactions.keys()) > 0: await ProxyUpdateSpend.update_end_user_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) ### UPDATE KEY TABLE ### - key_list_transactons = db_spend_update_transactions["key_list_transactons"] + key_list_transactions = db_spend_update_transactions["key_list_transactions"] verbose_proxy_logger.debug( - "KEY Spend transactions: {}".format(len(key_list_transactons.keys())) + "KEY Spend transactions: {}".format(len(key_list_transactions.keys())) ) - if len(key_list_transactons.keys()) > 0: + if len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -523,12 +523,12 @@ class DBSpendUpdateWriter: for ( token, response_cost, - ) in key_list_transactons.items(): + ) in key_list_transactions.items(): batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, data={"spend": {"increment": response_cost}}, ) - prisma_client.key_list_transactons = ( + prisma_client.key_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -551,11 +551,11 @@ class DBSpendUpdateWriter: ### UPDATE TEAM TABLE ### verbose_proxy_logger.debug( "Team Spend transactions: {}".format( - len(prisma_client.team_list_transactons.keys()) + len(prisma_client.team_list_transactions.keys()) ) ) - team_list_transactons = db_spend_update_transactions["team_list_transactons"] - if len(team_list_transactons.keys()) > 0: + team_list_transactions = db_spend_update_transactions["team_list_transactions"] + if len(team_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -566,7 +566,7 @@ class DBSpendUpdateWriter: for ( team_id, response_cost, - ) in team_list_transactons.items(): + ) in team_list_transactions.items(): verbose_proxy_logger.debug( "Updating spend for team id={} by {}".format( team_id, response_cost @@ -576,7 +576,7 @@ class DBSpendUpdateWriter: where={"team_id": team_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.team_list_transactons = ( + prisma_client.team_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -597,10 +597,10 @@ class DBSpendUpdateWriter: ) ### UPDATE TEAM Membership TABLE with spend ### - team_member_list_transactons = db_spend_update_transactions[ - "team_member_list_transactons" + team_member_list_transactions = db_spend_update_transactions[ + "team_member_list_transactions" ] - if len(team_member_list_transactons.keys()) > 0: + if len(team_member_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -611,7 +611,7 @@ class DBSpendUpdateWriter: for ( key, response_cost, - ) in team_member_list_transactons.items(): + ) in team_member_list_transactions.items(): # key is "team_id::::user_id::" team_id = key.split("::")[1] user_id = key.split("::")[3] @@ -620,7 +620,7 @@ class DBSpendUpdateWriter: where={"team_id": team_id, "user_id": user_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.team_member_list_transactons = ( + prisma_client.team_member_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -641,8 +641,8 @@ class DBSpendUpdateWriter: ) ### UPDATE ORG TABLE ### - org_list_transactons = db_spend_update_transactions["org_list_transactons"] - if len(org_list_transactons.keys()) > 0: + org_list_transactions = db_spend_update_transactions["org_list_transactions"] + if len(org_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -653,12 +653,12 @@ class DBSpendUpdateWriter: for ( org_id, response_cost, - ) in org_list_transactons.items(): + ) in org_list_transactions.items(): batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"organization_id": org_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.org_list_transactons = ( + prisma_client.org_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index ea7f14d321..22afc56483 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -17,12 +17,12 @@ else: class DBSpendUpdateTransactions(TypedDict): - user_list_transactons: Dict[str, float] - end_user_list_transactons: Dict[str, float] - key_list_transactons: Dict[str, float] - team_list_transactons: Dict[str, float] - team_member_list_transactons: Dict[str, float] - org_list_transactons: Dict[str, float] + user_list_transactions: Dict[str, float] + end_user_list_transactions: Dict[str, float] + key_list_transactions: Dict[str, float] + team_list_transactions: Dict[str, float] + team_member_list_transactions: Dict[str, float] + org_list_transactions: Dict[str, float] class RedisUpdateBuffer: @@ -78,12 +78,12 @@ class RedisUpdateBuffer: """ IN_MEMORY_UPDATE_TRANSACTIONS: DBSpendUpdateTransactions = ( DBSpendUpdateTransactions( - user_list_transactons=prisma_client.user_list_transactons, - end_user_list_transactons=prisma_client.end_user_list_transactons, - key_list_transactons=prisma_client.key_list_transactons, - team_list_transactons=prisma_client.team_list_transactons, - team_member_list_transactons=prisma_client.team_member_list_transactons, - org_list_transactons=prisma_client.org_list_transactons, + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, ) ) for key, _transactions in IN_MEMORY_UPDATE_TRANSACTIONS.items(): @@ -120,21 +120,23 @@ class RedisUpdateBuffer: if self.redis_cache is None: return None expected_keys = [ - "user_list_transactons", - "end_user_list_transactons", - "key_list_transactons", - "team_list_transactons", - "team_member_list_transactons", - "org_list_transactons", + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", ] result = await self.redis_cache.async_batch_get_cache(expected_keys) if result is None: return None return DBSpendUpdateTransactions( - user_list_transactons=result.get("user_list_transactons", {}), - end_user_list_transactons=result.get("end_user_list_transactons", {}), - key_list_transactons=result.get("key_list_transactons", {}), - team_list_transactons=result.get("team_list_transactons", {}), - team_member_list_transactons=result.get("team_member_list_transactons", {}), - org_list_transactons=result.get("org_list_transactons", {}), + user_list_transactions=result.get("user_list_transactions", {}), + end_user_list_transactions=result.get("end_user_list_transactions", {}), + key_list_transactions=result.get("key_list_transactions", {}), + team_list_transactions=result.get("team_list_transactions", {}), + team_member_list_transactions=result.get( + "team_member_list_transactions", {} + ), + org_list_transactions=result.get("org_list_transactions", {}), ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 040c6c14ef..d0d4ea9b4e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1101,12 +1101,12 @@ def jsonify_object(data: dict) -> dict: class PrismaClient: - user_list_transactons: dict = {} - end_user_list_transactons: dict = {} - key_list_transactons: dict = {} - team_list_transactons: dict = {} - team_member_list_transactons: dict = {} # key is ["team_id" + "user_id"] - org_list_transactons: dict = {} + 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 = [] daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {} @@ -2433,7 +2433,7 @@ class ProxyUpdateSpend: for ( end_user_id, response_cost, - ) in prisma_client.end_user_list_transactons.items(): + ) in prisma_client.end_user_list_transactions.items(): if litellm.max_end_user_budget is not None: pass batcher.litellm_endusertable.upsert( @@ -2461,7 +2461,7 @@ class ProxyUpdateSpend: e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) finally: - prisma_client.end_user_list_transactons = ( + prisma_client.end_user_list_transactions = ( {} ) # reset the end user list transactions - prevent bad data from causing issues diff --git a/tests/local_testing/test_update_spend.py b/tests/local_testing/test_update_spend.py index 6aeae851ab..fffa3062d7 100644 --- a/tests/local_testing/test_update_spend.py +++ b/tests/local_testing/test_update_spend.py @@ -93,7 +93,7 @@ def prisma_client(): @pytest.mark.asyncio async def test_batch_update_spend(prisma_client): - prisma_client.user_list_transactons["test-litellm-user-5"] = 23 + prisma_client.user_list_transactions["test-litellm-user-5"] = 23 setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index dccf0d1842..d613118fc8 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1509,7 +1509,7 @@ from litellm.proxy.utils import ProxyUpdateSpend async def test_end_user_transactions_reset(): # Setup mock_client = MagicMock() - mock_client.end_user_list_transactons = {"1": 10.0} # Bad log + mock_client.end_user_list_transactions = {"1": 10.0} # Bad log mock_client.db.tx = AsyncMock(side_effect=Exception("DB Error")) # Call function - should raise error @@ -1520,7 +1520,7 @@ async def test_end_user_transactions_reset(): # Verify cleanup happened assert ( - mock_client.end_user_list_transactons == {} + mock_client.end_user_list_transactions == {} ), "Transactions list should be empty after error" diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index 6efc68a077..355cb8a403 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -22,12 +22,12 @@ class MockPrismaClient: def __init__(self): self.db = MagicMock() self.spend_log_transactions = [] - self.user_list_transactons = {} - self.end_user_list_transactons = {} - self.key_list_transactons = {} - self.team_list_transactons = {} - self.team_member_list_transactons = {} - self.org_list_transactons = {} + self.user_list_transactions = {} + self.end_user_list_transactions = {} + self.key_list_transactions = {} + self.team_list_transactions = {} + self.team_member_list_transactions = {} + self.org_list_transactions = {} def jsonify_object(self, obj): return obj