From efe6d375e9adec23cc24f1934a41fc8e7e6fe298 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 18:40:03 -0700 Subject: [PATCH 01/20] add new SpendUpdateQueue --- litellm/proxy/db/spend_update_queue.py | 54 ++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 litellm/proxy/db/spend_update_queue.py diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py new file mode 100644 index 00000000000..778361be84e --- /dev/null +++ b/litellm/proxy/db/spend_update_queue.py @@ -0,0 +1,54 @@ +import asyncio +from typing import TYPE_CHECKING, Any, Dict, List + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient +else: + PrismaClient = Any + + +class SpendUpdateQueue: + """ + Handles buffering database `UPDATE` transactions in Redis before committing them to the database + + This is to prevent deadlocks and improve reliability + """ + + def __init__( + self, + ): + self.update_queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue() + + async def add_update(self, update: Dict[str, Any]) -> None: + """Enqueue an update. Each update might be a dict like {'entity_type': 'user', 'entity_id': '123', 'amount': 1.2}.""" + await self.update_queue.put(update) + + async def flush_all_updates_from_in_memory_queue(self) -> List[Dict[str, Any]]: + """Get all updates from the queue.""" + updates: List[Dict[str, Any]] = [] + while not self.update_queue.empty(): + updates.append(await self.update_queue.get()) + return updates + + async def flush_and_get_all_aggregated_updates_by_entity_type( + self, + ) -> Dict[str, Any]: + """Flush all updates from the queue and return all updates aggregated by entity type.""" + updates = await self.flush_all_updates_from_in_memory_queue() + return self.aggregate_updates_by_entity_type(updates) + + def aggregate_updates_by_entity_type( + self, updates: List[Dict[str, Any]] + ) -> Dict[str, Any]: + """Aggregate updates by entity type.""" + aggregated_updates = {} + for update in updates: + entity_type = update["entity_type"] + entity_id = update["entity_id"] + amount = update["amount"] + if entity_type not in aggregated_updates: + aggregated_updates[entity_type] = {} + if entity_id not in aggregated_updates[entity_type]: + aggregated_updates[entity_type][entity_id] = 0 + aggregated_updates[entity_type][entity_id] += amount + return aggregated_updates From bcd49204f6153629baa5a26cfd62d1d3ddba6294 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 18:40:52 -0700 Subject: [PATCH 02/20] use spend_update_queue for RedisUpdateBuffer --- litellm/proxy/db/redis_update_buffer.py | 42 ++++++++++--------------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index f98fc9300f8..7370b0ed08f 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -12,6 +12,7 @@ from litellm.caching import RedisCache from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_UPDATE_BUFFER_KEY from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DBSpendUpdateTransactions +from litellm.proxy.db.spend_update_queue import SpendUpdateQueue from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: @@ -32,6 +33,7 @@ class RedisUpdateBuffer: redis_cache: Optional[RedisCache] = None, ): self.redis_cache = redis_cache + self.spend_update_queue = SpendUpdateQueue() @staticmethod def _should_commit_spend_updates_to_redis() -> bool: @@ -54,7 +56,6 @@ class RedisUpdateBuffer: async def store_in_memory_spend_updates_in_redis( self, - prisma_client: PrismaClient, ): """ Stores the in-memory spend updates to Redis @@ -78,13 +79,21 @@ class RedisUpdateBuffer: "redis_cache is None, skipping store_in_memory_spend_updates_in_redis" ) return - db_spend_update_transactions: DBSpendUpdateTransactions = DBSpendUpdateTransactions( - 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, + + aggregated_updates = ( + await self.spend_update_queue.flush_and_get_all_aggregated_updates_by_entity_type() + ) + verbose_proxy_logger.debug("ALL AGGREGATED UPDATES: ", aggregated_updates) + + db_spend_update_transactions: DBSpendUpdateTransactions = ( + DBSpendUpdateTransactions( + user_list_transactions=aggregated_updates.get("user", {}), + end_user_list_transactions=aggregated_updates.get("end_user", {}), + key_list_transactions=aggregated_updates.get("key", {}), + team_list_transactions=aggregated_updates.get("team", {}), + team_member_list_transactions=aggregated_updates.get("team_member", {}), + org_list_transactions=aggregated_updates.get("org", {}), + ) ) # only store in redis if there are any updates to commit @@ -100,9 +109,6 @@ class RedisUpdateBuffer: values=list_of_transactions, ) - # clear the in-memory spend updates - RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client) - @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, @@ -116,20 +122,6 @@ class RedisUpdateBuffer: num_transactions += len(v) return num_transactions - @staticmethod - def _clear_all_in_memory_spend_updates( - prisma_client: PrismaClient, - ): - """ - Clears all in-memory spend updates - """ - prisma_client.user_list_transactions = {} - prisma_client.end_user_list_transactions = {} - prisma_client.key_list_transactions = {} - prisma_client.team_list_transactions = {} - prisma_client.team_member_list_transactions = {} - prisma_client.org_list_transactions = {} - @staticmethod def _remove_prefix_from_keys(data: Dict[str, Any], prefix: str) -> Dict[str, Any]: """ From 3e16a51ca6986095b89c0a3b967a8f28bdcc348b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:01:00 -0700 Subject: [PATCH 03/20] fix update_database helper on db_spend_update_writer --- litellm/proxy/db/db_spend_update_writer.py | 181 ++++++------------ litellm/proxy/db/redis_update_buffer.py | 8 +- litellm/proxy/db/spend_update_queue.py | 4 + .../proxy/hooks/proxy_track_cost_callback.py | 6 +- 4 files changed, 73 insertions(+), 126 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index f46b03b57ad..6b5719acbb3 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -25,6 +25,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.db.pod_lock_manager import PodLockManager from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.spend_update_queue import SpendUpdateQueue if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -48,10 +49,11 @@ class DBSpendUpdateWriter: self.redis_cache = redis_cache self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_lock_manager = PodLockManager(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) + self.spend_update_queue = SpendUpdateQueue() - @staticmethod async def update_database( # LiteLLM management object fields + self, token: Optional[str], user_id: Optional[str], end_user_id: Optional[str], @@ -84,7 +86,7 @@ class DBSpendUpdateWriter: hashed_token = token asyncio.create_task( - DBSpendUpdateWriter._update_user_db( + self._update_user_db( response_cost=response_cost, user_id=user_id, prisma_client=prisma_client, @@ -94,14 +96,14 @@ class DBSpendUpdateWriter: ) ) asyncio.create_task( - DBSpendUpdateWriter._update_key_db( + self._update_key_db( response_cost=response_cost, hashed_token=hashed_token, prisma_client=prisma_client, ) ) asyncio.create_task( - DBSpendUpdateWriter._update_team_db( + self._update_team_db( response_cost=response_cost, team_id=team_id, user_id=user_id, @@ -109,7 +111,7 @@ class DBSpendUpdateWriter: ) ) asyncio.create_task( - DBSpendUpdateWriter._update_org_db( + self._update_org_db( response_cost=response_cost, org_id=org_id, prisma_client=prisma_client, @@ -135,56 +137,8 @@ class DBSpendUpdateWriter: f"Error updating Prisma database: {traceback.format_exc()}" ) - @staticmethod - async def _update_transaction_list( - response_cost: Optional[float], - entity_id: Optional[str], - transaction_list: dict, - entity_type: Litellm_EntityType, - debug_msg: Optional[str] = None, - prisma_client: Optional[PrismaClient] = None, - ) -> bool: - """ - Common helper method to update a transaction list for an entity - - Args: - response_cost: The cost to add - entity_id: The ID of the entity to update - transaction_list: The transaction list dictionary to update - entity_type: The type of entity (from EntityType enum) - debug_msg: Optional custom debug message - - Returns: - bool: True if update happened, False otherwise - """ - try: - if debug_msg: - verbose_proxy_logger.debug(debug_msg) - else: - verbose_proxy_logger.debug( - f"adding spend to {entity_type.value} db. Response cost: {response_cost}. {entity_type.value}_id: {entity_id}." - ) - if prisma_client is None: - return False - - if entity_id is None: - verbose_proxy_logger.debug( - f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}" - ) - return False - transaction_list[entity_id] = response_cost + transaction_list.get( - entity_id, 0 - ) - return True - - except Exception as e: - verbose_proxy_logger.info( - f"Update {entity_type.value.capitalize()} DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - @staticmethod async def _update_key_db( + self, response_cost: Optional[float], hashed_token: Optional[str], prisma_client: Optional[PrismaClient], @@ -193,13 +147,12 @@ class DBSpendUpdateWriter: if hashed_token is None or prisma_client is None: return - await DBSpendUpdateWriter._update_transaction_list( - response_cost=response_cost, - entity_id=hashed_token, - 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}.", - prisma_client=prisma_client, + await self.spend_update_queue.add_update( + update={ + "entity_type": Litellm_EntityType.KEY.value, + "entity_id": hashed_token, + "amount": response_cost, + } ) except Exception as e: verbose_proxy_logger.exception( @@ -207,8 +160,8 @@ class DBSpendUpdateWriter: ) raise e - @staticmethod async def _update_user_db( + self, response_cost: Optional[float], user_id: Optional[str], prisma_client: Optional[PrismaClient], @@ -234,21 +187,21 @@ class DBSpendUpdateWriter: for _id in user_ids: if _id is not None: - await DBSpendUpdateWriter._update_transaction_list( - response_cost=response_cost, - entity_id=_id, - transaction_list=prisma_client.user_list_transactions, - entity_type=Litellm_EntityType.USER, - prisma_client=prisma_client, + await self.spend_update_queue.add_update( + update={ + "entity_type": Litellm_EntityType.USER.value, + "entity_id": _id, + "amount": response_cost, + } ) if end_user_id is not None: - await DBSpendUpdateWriter._update_transaction_list( - response_cost=response_cost, - entity_id=end_user_id, - transaction_list=prisma_client.end_user_list_transactions, - entity_type=Litellm_EntityType.END_USER, - prisma_client=prisma_client, + await self.spend_update_queue.add_update( + update={ + "entity_type": Litellm_EntityType.END_USER.value, + "entity_id": end_user_id, + "amount": response_cost, + } ) except Exception as e: verbose_proxy_logger.info( @@ -256,8 +209,8 @@ class DBSpendUpdateWriter: + f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}" ) - @staticmethod async def _update_team_db( + self, response_cost: Optional[float], team_id: Optional[str], user_id: Optional[str], @@ -270,12 +223,12 @@ class DBSpendUpdateWriter: ) return - await DBSpendUpdateWriter._update_transaction_list( - response_cost=response_cost, - entity_id=team_id, - transaction_list=prisma_client.team_list_transactions, - entity_type=Litellm_EntityType.TEAM, - prisma_client=prisma_client, + await self.spend_update_queue.add_update( + update={ + "entity_type": Litellm_EntityType.TEAM.value, + "entity_id": team_id, + "amount": response_cost, + } ) try: @@ -283,12 +236,12 @@ class DBSpendUpdateWriter: if user_id is not None: # key is "team_id::::user_id::" team_member_key = f"team_id::{team_id}::user_id::{user_id}" - await DBSpendUpdateWriter._update_transaction_list( - response_cost=response_cost, - entity_id=team_member_key, - transaction_list=prisma_client.team_member_list_transactions, - entity_type=Litellm_EntityType.TEAM_MEMBER, - prisma_client=prisma_client, + await self.spend_update_queue.add_update( + update={ + "entity_type": Litellm_EntityType.TEAM_MEMBER.value, + "entity_id": team_member_key, + "amount": response_cost, + } ) except Exception: pass @@ -298,8 +251,8 @@ class DBSpendUpdateWriter: ) raise e - @staticmethod async def _update_org_db( + self, response_cost: Optional[float], org_id: Optional[str], prisma_client: Optional[PrismaClient], @@ -311,12 +264,12 @@ class DBSpendUpdateWriter: ) return - await DBSpendUpdateWriter._update_transaction_list( - response_cost=response_cost, - entity_id=org_id, - transaction_list=prisma_client.org_list_transactions, - entity_type=Litellm_EntityType.ORGANIZATION, - prisma_client=prisma_client, + await self.spend_update_queue.add_update( + update={ + "entity_type": Litellm_EntityType.ORGANIZATION.value, + "entity_id": org_id, + "amount": response_cost, + } ) except Exception as e: verbose_proxy_logger.info( @@ -435,7 +388,7 @@ class DBSpendUpdateWriter: - Only 1 pod will commit to db at a time (based on if it can acquire the lock over writing to DB) """ await self.redis_update_buffer.store_in_memory_spend_updates_in_redis( - prisma_client=prisma_client, + spend_update_queue=self.spend_update_queue, ) # Only commit from redis to db if this pod is the leader @@ -447,7 +400,7 @@ class DBSpendUpdateWriter: await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() ) if db_spend_update_transactions is not None: - await DBSpendUpdateWriter._commit_spend_updates_to_db( + await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, @@ -471,23 +424,26 @@ class DBSpendUpdateWriter: Note: This flow causes Deadlocks in production (1K RPS+). Use self._commit_spend_updates_to_db_with_redis() instead if you expect 1K+ RPS. """ - db_spend_update_transactions = DBSpendUpdateTransactions( - 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, + aggregated_updates = ( + await self.spend_update_queue.flush_and_get_all_aggregated_updates_by_entity_type() ) - await DBSpendUpdateWriter._commit_spend_updates_to_db( + db_spend_update_transactions = DBSpendUpdateTransactions( + user_list_transactions=aggregated_updates.get("user", {}), + end_user_list_transactions=aggregated_updates.get("end_user", {}), + key_list_transactions=aggregated_updates.get("key", {}), + team_list_transactions=aggregated_updates.get("team", {}), + team_member_list_transactions=aggregated_updates.get("team_member", {}), + org_list_transactions=aggregated_updates.get("organization", {}), + ) + await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, db_spend_update_transactions=db_spend_update_transactions, ) - @staticmethod async def _commit_spend_updates_to_db( # noqa: PLR0915 + self, prisma_client: PrismaClient, n_retry_times: int, proxy_logging_obj: ProxyLogging, @@ -526,9 +482,6 @@ class DBSpendUpdateWriter: where={"user_id": user_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.user_list_transactions = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. break except DB_CONNECTION_ERROR_TYPES as e: if ( @@ -583,9 +536,6 @@ class DBSpendUpdateWriter: where={"token": token}, data={"spend": {"increment": response_cost}}, ) - prisma_client.key_list_transactions = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. break except DB_CONNECTION_ERROR_TYPES as e: if ( @@ -632,9 +582,6 @@ class DBSpendUpdateWriter: where={"team_id": team_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.team_list_transactions = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. break except DB_CONNECTION_ERROR_TYPES as e: if ( @@ -684,9 +631,6 @@ class DBSpendUpdateWriter: where={"team_id": team_id, "user_id": user_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.team_member_list_transactions = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. break except DB_CONNECTION_ERROR_TYPES as e: if ( @@ -725,9 +669,6 @@ class DBSpendUpdateWriter: where={"organization_id": org_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.org_list_transactions = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. break except DB_CONNECTION_ERROR_TYPES as e: if ( diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 7370b0ed08f..0dfaa72a16f 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -33,7 +33,6 @@ class RedisUpdateBuffer: redis_cache: Optional[RedisCache] = None, ): self.redis_cache = redis_cache - self.spend_update_queue = SpendUpdateQueue() @staticmethod def _should_commit_spend_updates_to_redis() -> bool: @@ -56,6 +55,7 @@ class RedisUpdateBuffer: async def store_in_memory_spend_updates_in_redis( self, + spend_update_queue: SpendUpdateQueue, ): """ Stores the in-memory spend updates to Redis @@ -81,9 +81,9 @@ class RedisUpdateBuffer: return aggregated_updates = ( - await self.spend_update_queue.flush_and_get_all_aggregated_updates_by_entity_type() + await spend_update_queue.flush_and_get_all_aggregated_updates_by_entity_type() ) - verbose_proxy_logger.debug("ALL AGGREGATED UPDATES: ", aggregated_updates) + verbose_proxy_logger.debug("ALL AGGREGATED UPDATES: %s", aggregated_updates) db_spend_update_transactions: DBSpendUpdateTransactions = ( DBSpendUpdateTransactions( @@ -92,7 +92,7 @@ class RedisUpdateBuffer: key_list_transactions=aggregated_updates.get("key", {}), team_list_transactions=aggregated_updates.get("team", {}), team_member_list_transactions=aggregated_updates.get("team_member", {}), - org_list_transactions=aggregated_updates.get("org", {}), + org_list_transactions=aggregated_updates.get("organization", {}), ) ) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 778361be84e..2462e6a547c 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -1,6 +1,8 @@ import asyncio from typing import TYPE_CHECKING, Any, Dict, List +from litellm._logging import verbose_proxy_logger + if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient else: @@ -21,6 +23,7 @@ class SpendUpdateQueue: async def add_update(self, update: Dict[str, Any]) -> None: """Enqueue an update. Each update might be a dict like {'entity_type': 'user', 'entity_id': '123', 'amount': 1.2}.""" + verbose_proxy_logger.debug("Adding update to queue: %s", update) await self.update_queue.put(update) async def flush_all_updates_from_in_memory_queue(self) -> List[Dict[str, Any]]: @@ -35,6 +38,7 @@ class SpendUpdateQueue: ) -> Dict[str, Any]: """Flush all updates from the queue and return all updates aggregated by entity type.""" updates = await self.flush_all_updates_from_in_memory_queue() + verbose_proxy_logger.debug("Aggregating updates by entity type: %s", updates) return self.aggregate_updates_by_entity_type(updates) def aggregate_updates_by_entity_type( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 39c1eeace98..dc0c27eb3ec 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -37,6 +37,8 @@ class _ProxyDBLogger(CustomLogger): if _ProxyDBLogger._should_track_errors_in_db() is False: return + from litellm.proxy.proxy_server import proxy_logging_obj + _metadata = dict( StandardLoggingUserAPIKeyMetadata( user_api_key_hash=user_api_key_dict.api_key, @@ -66,7 +68,7 @@ class _ProxyDBLogger(CustomLogger): request_data.get("proxy_server_request") or {} ) request_data["litellm_params"]["metadata"] = existing_metadata - await DBSpendUpdateWriter.update_database( + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, response_cost=0.0, user_id=user_api_key_dict.user_id, @@ -136,7 +138,7 @@ class _ProxyDBLogger(CustomLogger): end_user_id=end_user_id, ): ## UPDATE DATABASE - await DBSpendUpdateWriter.update_database( + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, From 95e674d1e98729e8c010046cda707dc9da645c50 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:01:17 -0700 Subject: [PATCH 04/20] test spend accuracy --- litellm/proxy/proxy_config.yaml | 24 +++++++-------- .../local_test_spend_accuracy_tests.py | 30 +++++++++---------- 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index c0107d7ea6f..fe8d73d26aa 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,15 +1,15 @@ model_list: - - model_name: gpt-4o + - model_name: fake-openai-endpoint litellm_params: - model: openai/gpt-4o - api_key: sk-xxxxxxx + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ -mcp_servers: - { - "zapier_mcp": { - "url": "https://actions.zapier.com/mcp/sk-akxxxxx/sse" - }, - "fetch": { - "url": "http://localhost:8000/sse" - } - } +general_settings: + use_redis_transaction_buffer: true + +litellm_settings: + cache: True + cache_params: + type: redis + supported_call_types: [] \ No newline at end of file diff --git a/tests/otel_tests/local_test_spend_accuracy_tests.py b/tests/otel_tests/local_test_spend_accuracy_tests.py index 6d756219c72..bbb12c58a1c 100644 --- a/tests/otel_tests/local_test_spend_accuracy_tests.py +++ b/tests/otel_tests/local_test_spend_accuracy_tests.py @@ -52,7 +52,7 @@ Additional Test Scenarios: async def create_organization(session, organization_alias: str): """Helper function to create a new organization""" - url = "http://0.0.0.0:4002/organization/new" + url = "http://0.0.0.0:4000/organization/new" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} data = {"organization_alias": organization_alias} async with session.post(url, headers=headers, json=data) as response: @@ -61,7 +61,7 @@ async def create_organization(session, organization_alias: str): async def create_team(session, org_id: str): """Helper function to create a new team under an organization""" - url = "http://0.0.0.0:4002/team/new" + url = "http://0.0.0.0:4000/team/new" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} data = {"organization_id": org_id, "team_alias": f"test-team-{uuid.uuid4()}"} async with session.post(url, headers=headers, json=data) as response: @@ -70,7 +70,7 @@ async def create_team(session, org_id: str): async def create_user(session, org_id: str): """Helper function to create a new user""" - url = "http://0.0.0.0:4002/user/new" + url = "http://0.0.0.0:4000/user/new" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} data = {"user_name": f"test-user-{uuid.uuid4()}"} async with session.post(url, headers=headers, json=data) as response: @@ -79,7 +79,7 @@ async def create_user(session, org_id: str): async def generate_key(session, user_id: str, team_id: str): """Helper function to generate a key for a specific user and team""" - url = "http://0.0.0.0:4002/key/generate" + url = "http://0.0.0.0:4000/key/generate" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} data = {"user_id": user_id, "team_id": team_id} async with session.post(url, headers=headers, json=data) as response: @@ -91,7 +91,7 @@ async def chat_completion(session, key: str): from openai import AsyncOpenAI import uuid - client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4002/v1") + client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000/v1") response = await client.chat.completions.create( model="fake-openai-endpoint", @@ -102,7 +102,7 @@ async def chat_completion(session, key: str): async def get_spend_info(session, entity_type: str, entity_id: str): """Helper function to get spend information for an entity""" - url = f"http://0.0.0.0:4002/{entity_type}/info" + url = f"http://0.0.0.0:4000/{entity_type}/info" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} if entity_type == "key": data = {"key": entity_id} @@ -266,14 +266,14 @@ async def test_long_term_spend_accuracy_with_bursts(): abs(key_info["info"]["spend"] - expected_spend) < TOLERANCE ), f"Key spend {key_info['info']['spend']} does not match expected {expected_spend}" - assert ( - abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE - ), f"User spend {user_info['user_info']['spend']} does not match expected {expected_spend}" + # assert ( + # abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE + # ), f"User spend {user_info['user_info']['spend']} does not match expected {expected_spend}" - assert ( - abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE - ), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}" + # assert ( + # abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE + # ), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}" - assert ( - abs(org_info["spend"] - expected_spend) < TOLERANCE - ), f"Organization spend {org_info['spend']} does not match expected {expected_spend}" + # assert ( + # abs(org_info["spend"] - expected_spend) < TOLERANCE + # ), f"Organization spend {org_info['spend']} does not match expected {expected_spend}" From 73d6af55723975711fca28f5d64d4ed9cd9e8dc9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:03:02 -0700 Subject: [PATCH 05/20] fix docstring --- litellm/proxy/db/spend_update_queue.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 2462e6a547c..2a1e336877c 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -11,9 +11,7 @@ else: class SpendUpdateQueue: """ - Handles buffering database `UPDATE` transactions in Redis before committing them to the database - - This is to prevent deadlocks and improve reliability + In memory buffer for spend updates that should be committed to the database """ def __init__( From 49595121703279ce924a0f88efb02124a0c20a4b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:05:58 -0700 Subject: [PATCH 06/20] fix update_end_user_spend --- litellm/proxy/db/db_spend_update_writer.py | 1 + litellm/proxy/utils.py | 17 +++++------------ 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 6b5719acbb3..5899ab3416a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -514,6 +514,7 @@ class DBSpendUpdateWriter: n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, + end_user_list_transactions=end_user_list_transactions, ) ### UPDATE KEY TABLE ### key_list_transactions = db_spend_update_transactions["key_list_transactions"] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 900b26f3f1a..76591d9fca7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1100,12 +1100,6 @@ def jsonify_object(data: dict) -> dict: 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 = [] daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {} @@ -2422,7 +2416,10 @@ def _hash_token_if_needed(token: str) -> str: class ProxyUpdateSpend: @staticmethod 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): start_time = time.time() @@ -2434,7 +2431,7 @@ class ProxyUpdateSpend: for ( end_user_id, 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: pass batcher.litellm_endusertable.upsert( @@ -2461,10 +2458,6 @@ class ProxyUpdateSpend: _raise_failed_update_spend_exception( 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 async def update_spend_logs( From a753fc9d9f8c6c833b8afd34f4e934f94957c5fe Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:17:13 -0700 Subject: [PATCH 07/20] test_long_term_spend_accuracy_with_bursts --- .../local_test_spend_accuracy_tests.py | 18 +++++++++--------- tests/proxy_unit_tests/test_update_spend.py | 6 ------ 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/tests/otel_tests/local_test_spend_accuracy_tests.py b/tests/otel_tests/local_test_spend_accuracy_tests.py index bbb12c58a1c..d60bc985a6c 100644 --- a/tests/otel_tests/local_test_spend_accuracy_tests.py +++ b/tests/otel_tests/local_test_spend_accuracy_tests.py @@ -266,14 +266,14 @@ async def test_long_term_spend_accuracy_with_bursts(): abs(key_info["info"]["spend"] - expected_spend) < TOLERANCE ), f"Key spend {key_info['info']['spend']} does not match expected {expected_spend}" - # assert ( - # abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE - # ), f"User spend {user_info['user_info']['spend']} does not match expected {expected_spend}" + assert ( + abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE + ), f"User spend {user_info['user_info']['spend']} does not match expected {expected_spend}" - # assert ( - # abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE - # ), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}" + assert ( + abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE + ), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}" - # assert ( - # abs(org_info["spend"] - expected_spend) < TOLERANCE - # ), f"Organization spend {org_info['spend']} does not match expected {expected_spend}" + assert ( + abs(org_info["spend"] - expected_spend) < TOLERANCE + ), f"Organization spend {org_info['spend']} does not match expected {expected_spend}" diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index 641768a7d2e..1fb2479792b 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -27,12 +27,6 @@ class MockPrismaClient: # Initialize transaction lists 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.daily_user_spend_transactions = {} def jsonify_object(self, obj): From 71e772dd4a5a6547c4ca97374d1ba3898e4660e9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:28:17 -0700 Subject: [PATCH 08/20] use typed data structure for queue --- litellm/proxy/_types.py | 6 ++ litellm/proxy/db/db_spend_update_writer.py | 73 +++++++++----------- litellm/proxy/db/redis_update_buffer.py | 17 ++--- litellm/proxy/db/spend_update_queue.py | 78 ++++++++++++++++------ 4 files changed, 101 insertions(+), 73 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7f13717e299..16d302aa9a6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2749,3 +2749,9 @@ class DBSpendUpdateTransactions(TypedDict): team_list_transactions: Optional[Dict[str, float]] team_member_list_transactions: Optional[Dict[str, float]] org_list_transactions: Optional[Dict[str, float]] + + +class SpendUpdateQueueItem(TypedDict, total=False): + entity_type: Litellm_EntityType + entity_id: str + response_cost: Optional[float] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 5899ab3416a..5bf255feae2 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload, + SpendUpdateQueueItem, ) from litellm.proxy.db.pod_lock_manager import PodLockManager from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer @@ -148,11 +149,11 @@ class DBSpendUpdateWriter: return await self.spend_update_queue.add_update( - update={ - "entity_type": Litellm_EntityType.KEY.value, - "entity_id": hashed_token, - "amount": response_cost, - } + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.KEY, + entity_id=hashed_token, + response_cost=response_cost, + ) ) except Exception as e: verbose_proxy_logger.exception( @@ -188,20 +189,20 @@ class DBSpendUpdateWriter: for _id in user_ids: if _id is not None: await self.spend_update_queue.add_update( - update={ - "entity_type": Litellm_EntityType.USER.value, - "entity_id": _id, - "amount": response_cost, - } + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, + entity_id=_id, + response_cost=response_cost, + ) ) if end_user_id is not None: await self.spend_update_queue.add_update( - update={ - "entity_type": Litellm_EntityType.END_USER.value, - "entity_id": end_user_id, - "amount": response_cost, - } + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.END_USER, + entity_id=end_user_id, + response_cost=response_cost, + ) ) except Exception as e: verbose_proxy_logger.info( @@ -224,11 +225,11 @@ class DBSpendUpdateWriter: return await self.spend_update_queue.add_update( - update={ - "entity_type": Litellm_EntityType.TEAM.value, - "entity_id": team_id, - "amount": response_cost, - } + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + response_cost=response_cost, + ) ) try: @@ -237,11 +238,11 @@ class DBSpendUpdateWriter: # key is "team_id::::user_id::" team_member_key = f"team_id::{team_id}::user_id::{user_id}" await self.spend_update_queue.add_update( - update={ - "entity_type": Litellm_EntityType.TEAM_MEMBER.value, - "entity_id": team_member_key, - "amount": response_cost, - } + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.TEAM_MEMBER, + entity_id=team_member_key, + response_cost=response_cost, + ) ) except Exception: pass @@ -265,11 +266,11 @@ class DBSpendUpdateWriter: return await self.spend_update_queue.add_update( - update={ - "entity_type": Litellm_EntityType.ORGANIZATION.value, - "entity_id": org_id, - "amount": response_cost, - } + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.ORGANIZATION, + entity_id=org_id, + response_cost=response_cost, + ) ) except Exception as e: verbose_proxy_logger.info( @@ -424,16 +425,8 @@ class DBSpendUpdateWriter: Note: This flow causes Deadlocks in production (1K RPS+). Use self._commit_spend_updates_to_db_with_redis() instead if you expect 1K+ RPS. """ - aggregated_updates = ( - await self.spend_update_queue.flush_and_get_all_aggregated_updates_by_entity_type() - ) - db_spend_update_transactions = DBSpendUpdateTransactions( - user_list_transactions=aggregated_updates.get("user", {}), - end_user_list_transactions=aggregated_updates.get("end_user", {}), - key_list_transactions=aggregated_updates.get("key", {}), - team_list_transactions=aggregated_updates.get("team", {}), - team_member_list_transactions=aggregated_updates.get("team_member", {}), - org_list_transactions=aggregated_updates.get("organization", {}), + db_spend_update_transactions = ( + await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 0dfaa72a16f..1a3fd3d42d1 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -80,20 +80,11 @@ class RedisUpdateBuffer: ) return - aggregated_updates = ( - await spend_update_queue.flush_and_get_all_aggregated_updates_by_entity_type() + db_spend_update_transactions = ( + await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() ) - verbose_proxy_logger.debug("ALL AGGREGATED UPDATES: %s", aggregated_updates) - - db_spend_update_transactions: DBSpendUpdateTransactions = ( - DBSpendUpdateTransactions( - user_list_transactions=aggregated_updates.get("user", {}), - end_user_list_transactions=aggregated_updates.get("end_user", {}), - key_list_transactions=aggregated_updates.get("key", {}), - team_list_transactions=aggregated_updates.get("team", {}), - team_member_list_transactions=aggregated_updates.get("team_member", {}), - org_list_transactions=aggregated_updates.get("organization", {}), - ) + verbose_proxy_logger.debug( + "ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions ) # only store in redis if there are any updates to commit diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 2a1e336877c..2d9792ac854 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -2,6 +2,11 @@ import asyncio from typing import TYPE_CHECKING, Any, Dict, List from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + DBSpendUpdateTransactions, + Litellm_EntityType, + SpendUpdateQueueItem, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -17,40 +22,73 @@ class SpendUpdateQueue: def __init__( self, ): - self.update_queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue() + self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue() - async def add_update(self, update: Dict[str, Any]) -> None: + async def add_update(self, update: SpendUpdateQueueItem) -> None: """Enqueue an update. Each update might be a dict like {'entity_type': 'user', 'entity_id': '123', 'amount': 1.2}.""" verbose_proxy_logger.debug("Adding update to queue: %s", update) await self.update_queue.put(update) - async def flush_all_updates_from_in_memory_queue(self) -> List[Dict[str, Any]]: + async def flush_all_updates_from_in_memory_queue( + self, + ) -> List[SpendUpdateQueueItem]: """Get all updates from the queue.""" - updates: List[Dict[str, Any]] = [] + updates: List[SpendUpdateQueueItem] = [] while not self.update_queue.empty(): updates.append(await self.update_queue.get()) return updates - async def flush_and_get_all_aggregated_updates_by_entity_type( + async def flush_and_get_aggregated_db_spend_update_transactions( self, - ) -> Dict[str, Any]: + ) -> DBSpendUpdateTransactions: """Flush all updates from the queue and return all updates aggregated by entity type.""" updates = await self.flush_all_updates_from_in_memory_queue() verbose_proxy_logger.debug("Aggregating updates by entity type: %s", updates) - return self.aggregate_updates_by_entity_type(updates) + return self.get_aggregated_db_spend_update_transactions(updates) - def aggregate_updates_by_entity_type( - self, updates: List[Dict[str, Any]] - ) -> Dict[str, Any]: + def get_aggregated_db_spend_update_transactions( + self, updates: List[SpendUpdateQueueItem] + ) -> DBSpendUpdateTransactions: """Aggregate updates by entity type.""" - aggregated_updates = {} + # Initialize all transaction lists as empty dicts + db_spend_update_transactions = DBSpendUpdateTransactions( + user_list_transactions={}, + end_user_list_transactions={}, + key_list_transactions={}, + team_list_transactions={}, + team_member_list_transactions={}, + org_list_transactions={}, + ) + + # Map entity types to their corresponding transaction dictionary keys + entity_type_to_dict_key = { + Litellm_EntityType.USER: "user_list_transactions", + Litellm_EntityType.END_USER: "end_user_list_transactions", + Litellm_EntityType.KEY: "key_list_transactions", + Litellm_EntityType.TEAM: "team_list_transactions", + Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", + Litellm_EntityType.ORGANIZATION: "org_list_transactions", + } + for update in updates: - entity_type = update["entity_type"] - entity_id = update["entity_id"] - amount = update["amount"] - if entity_type not in aggregated_updates: - aggregated_updates[entity_type] = {} - if entity_id not in aggregated_updates[entity_type]: - aggregated_updates[entity_type][entity_id] = 0 - aggregated_updates[entity_type][entity_id] += amount - return aggregated_updates + entity_type = update.get("entity_type") + entity_id = update.get("entity_id") + response_cost = update.get("response_cost") + + if entity_type is None or entity_id is None or response_cost is None: + raise ValueError("Invalid update: %s", update) + + dict_key = entity_type_to_dict_key.get(entity_type) + if dict_key is None: + continue # Skip unknown entity types + + transactions_dict = db_spend_update_transactions[dict_key] + if transactions_dict is None: + transactions_dict = {} + db_spend_update_transactions[dict_key] = transactions_dict + + transactions_dict[entity_id] = ( + transactions_dict.get(entity_id, 0) + response_cost + ) + + return db_spend_update_transactions From 811f488ca39f83bb56b02bac05992feebb52aae7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:31:53 -0700 Subject: [PATCH 09/20] get_aggregated_db_spend_update_transactions --- litellm/proxy/db/spend_update_queue.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 2d9792ac854..d502c2cd848 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -87,8 +87,9 @@ class SpendUpdateQueue: transactions_dict = {} db_spend_update_transactions[dict_key] = transactions_dict - transactions_dict[entity_id] = ( - transactions_dict.get(entity_id, 0) + response_cost - ) + if entity_id not in transactions_dict: + transactions_dict[entity_id] = 0 + + transactions_dict[entity_id] += response_cost return db_spend_update_transactions From aa8261af89e04f75da92ab4628d73028364eee90 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:33:10 -0700 Subject: [PATCH 10/20] test fixes --- litellm/proxy/hooks/proxy_track_cost_callback.py | 1 - tests/proxy_unit_tests/test_unit_test_proxy_hooks.py | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index dc0c27eb3ec..9ae1eb4c340 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -13,7 +13,6 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import log_db_metrics -from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( StandardLoggingPayload, diff --git a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py b/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py index 129be6d7546..46863889d26 100644 --- a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py +++ b/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py @@ -24,9 +24,10 @@ async def test_disable_spend_logs(): "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + db_spend_update_writer = DBSpendUpdateWriter() # Call update_database with disable_spend_logs=True - await DBSpendUpdateWriter.update_database( + await db_spend_update_writer.update_database( token="fake-token", response_cost=0.1, user_id="user123", From 271b8b95bc8d12cbd378591c35d6136624164b17 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:35:07 -0700 Subject: [PATCH 11/20] test spend accuracy --- .../test_spend_accuracy_tests.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{otel_tests/local_test_spend_accuracy_tests.py => spend_tracking_tests/test_spend_accuracy_tests.py} (100%) diff --git a/tests/otel_tests/local_test_spend_accuracy_tests.py b/tests/spend_tracking_tests/test_spend_accuracy_tests.py similarity index 100% rename from tests/otel_tests/local_test_spend_accuracy_tests.py rename to tests/spend_tracking_tests/test_spend_accuracy_tests.py From 6d4a6a84029acd7b2b2365b4e18bfff846c1a91f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:42:00 -0700 Subject: [PATCH 12/20] add spend tracking config.yaml --- .circleci/config.yml | 99 ++++++++++++++++++- .../spend_tracking_config.yaml | 15 +++ 2 files changed, 113 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/example_config_yaml/spend_tracking_config.yaml diff --git a/.circleci/config.yml b/.circleci/config.yml index e1488a9083b..d2a6aafef7f 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1450,7 +1450,7 @@ jobs: command: | pwd ls - python -m pytest -s -vv tests/*.py -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests --ignore=tests/load_tests --ignore=tests/llm_translation --ignore=tests/llm_responses_api_testing --ignore=tests/mcp_tests --ignore=tests/image_gen_tests --ignore=tests/pass_through_unit_tests + python -m pytest -s -vv tests/*.py -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/spend_tracking_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests --ignore=tests/load_tests --ignore=tests/llm_translation --ignore=tests/llm_responses_api_testing --ignore=tests/mcp_tests --ignore=tests/image_gen_tests --ignore=tests/pass_through_unit_tests no_output_timeout: 120m # Store test results @@ -1743,6 +1743,96 @@ jobs: # Store test results - store_test_results: path: test-results + proxy_spend_accuracy_tests: + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge + working_directory: ~/project + steps: + - checkout + - setup_google_dns + - run: + name: Install Docker CLI (In case it's not already installed) + command: | + sudo apt-get update + sudo apt-get install -y docker-ce docker-ce-cli containerd.io + - run: + name: Install Python 3.9 + command: | + curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh + bash miniconda.sh -b -p $HOME/miniconda + export PATH="$HOME/miniconda/bin:$PATH" + conda init bash + source ~/.bashrc + conda create -n myenv python=3.9 -y + conda activate myenv + python --version + - run: + name: Install Dependencies + command: | + pip install "pytest==7.3.1" + pip install "pytest-asyncio==0.21.1" + pip install aiohttp + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + - run: + name: Build Docker image + command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + - run: + name: Run Docker container + # intentionally give bad redis credentials here + # the OTEL test - should get this as a trace + command: | + docker run -d \ + -p 4000:4000 \ + -e DATABASE_URL=$PROXY_DATABASE_URL \ + -e REDIS_HOST=$REDIS_HOST \ + -e REDIS_PASSWORD=$REDIS_PASSWORD \ + -e REDIS_PORT=$REDIS_PORT \ + -e LITELLM_MASTER_KEY="sk-1234" \ + -e OPENAI_API_KEY=$OPENAI_API_KEY \ + -e LITELLM_LICENSE=$LITELLM_LICENSE \ + -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \ + -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \ + -e USE_DDTRACE=True \ + -e DD_API_KEY=$DD_API_KEY \ + -e DD_SITE=$DD_SITE \ + -e AWS_REGION_NAME=$AWS_REGION_NAME \ + --name my-app \ + -v $(pwd)/litellm/proxy/example_config_yaml/spend_tracking_config.yaml:/app/config.yaml \ + my-app:latest \ + --config /app/config.yaml \ + --port 4000 \ + --detailed_debug \ + - run: + name: Install curl and dockerize + command: | + sudo apt-get update + sudo apt-get install -y curl + sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz + sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz + sudo rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Start outputting logs + command: docker logs -f my-app + background: true + - run: + name: Wait for app to be ready + command: dockerize -wait http://localhost:4000 -timeout 5m + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/spend_tracking_tests -x --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: + 120m + # Clean up first container + - run: + name: Stop and remove first container + command: | + docker stop my-app + docker rm my-app proxy_multi_instance_tests: machine: @@ -2553,6 +2643,12 @@ workflows: only: - main - /litellm_.*/ + - proxy_spend_accuracy_tests: + filters: + branches: + only: + - main + - /litellm_.*/ - proxy_multi_instance_tests: filters: branches: @@ -2714,6 +2810,7 @@ workflows: - installing_litellm_on_python - installing_litellm_on_python_3_13 - proxy_logging_guardrails_model_info_tests + - proxy_spend_accuracy_tests - proxy_multi_instance_tests - proxy_store_model_in_db_tests - proxy_build_from_pip_tests diff --git a/litellm/proxy/example_config_yaml/spend_tracking_config.yaml b/litellm/proxy/example_config_yaml/spend_tracking_config.yaml new file mode 100644 index 00000000000..fe8d73d26aa --- /dev/null +++ b/litellm/proxy/example_config_yaml/spend_tracking_config.yaml @@ -0,0 +1,15 @@ +model_list: + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + +general_settings: + use_redis_transaction_buffer: true + +litellm_settings: + cache: True + cache_params: + type: redis + supported_call_types: [] \ No newline at end of file From 5fa5c1154eeb24d07e88ec8b034c61e5695af752 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 19:42:20 -0700 Subject: [PATCH 13/20] add SpendUpdateQueue --- litellm/proxy/db/spend_update_queue.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index d502c2cd848..e695f428413 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING, Any, List from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( From 923ac2303beedafea7a5f16e088eda586e8ff098 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 20:55:13 -0700 Subject: [PATCH 14/20] test_end_user_transactions_reset --- tests/proxy_unit_tests/test_proxy_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index d613118fc89..31c9711b0eb 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1509,13 +1509,16 @@ from litellm.proxy.utils import ProxyUpdateSpend async def test_end_user_transactions_reset(): # Setup mock_client = MagicMock() - mock_client.end_user_list_transactions = {"1": 10.0} # Bad log + end_user_list_transactions = {"1": 10.0} # Bad log mock_client.db.tx = AsyncMock(side_effect=Exception("DB Error")) # Call function - should raise error with pytest.raises(Exception): await ProxyUpdateSpend.update_end_user_spend( - n_retry_times=0, prisma_client=mock_client, proxy_logging_obj=MagicMock() + n_retry_times=0, + prisma_client=mock_client, + proxy_logging_obj=MagicMock(), + end_user_list_transactions=end_user_list_transactions, ) # Verify cleanup happened From 9951b356da8ad3e274902507d0ee1b62244871be Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 21:09:29 -0700 Subject: [PATCH 15/20] test_long_term_spend_accuracy_with_bursts --- tests/spend_tracking_tests/test_spend_accuracy_tests.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/spend_tracking_tests/test_spend_accuracy_tests.py b/tests/spend_tracking_tests/test_spend_accuracy_tests.py index d60bc985a6c..93228c2d06a 100644 --- a/tests/spend_tracking_tests/test_spend_accuracy_tests.py +++ b/tests/spend_tracking_tests/test_spend_accuracy_tests.py @@ -156,7 +156,7 @@ async def test_basic_spend_accuracy(): response = await chat_completion(session, key) print("response: ", response) - # wait 10 seconds for spend to be updated + # wait 15 seconds for spend to be updated await asyncio.sleep(15) # Get spend information for each entity @@ -235,7 +235,7 @@ async def test_long_term_spend_accuracy_with_bursts(): print(f"Burst 1 - Request {i+1}/{BURST_1_REQUESTS} completed") # Wait for spend to be updated - await asyncio.sleep(8) + await asyncio.sleep(15) # Check intermediate spend intermediate_key_info = await get_spend_info(session, "key", key) @@ -248,7 +248,7 @@ async def test_long_term_spend_accuracy_with_bursts(): print(f"Burst 2 - Request {i+1}/{BURST_2_REQUESTS} completed") # Wait for spend to be updated - await asyncio.sleep(8) + await asyncio.sleep(15) # Get final spend information for each entity key_info = await get_spend_info(session, "key", key) From f7ddc583f0371c95c2582d924ac6c1b0745b4c35 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 21:15:21 -0700 Subject: [PATCH 16/20] linting fix get_aggregated_db_spend_update_transactions --- litellm/proxy/db/spend_update_queue.py | 31 ++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index e695f428413..29118073efa 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Any, List +from typing import TYPE_CHECKING, Any, Dict, List from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -82,7 +82,34 @@ class SpendUpdateQueue: if dict_key is None: continue # Skip unknown entity types - transactions_dict = db_spend_update_transactions[dict_key] + # Type-safe access using if/elif statements + if dict_key == "user_list_transactions": + transactions_dict = db_spend_update_transactions[ + "user_list_transactions" + ] + elif dict_key == "end_user_list_transactions": + transactions_dict = db_spend_update_transactions[ + "end_user_list_transactions" + ] + elif dict_key == "key_list_transactions": + transactions_dict = db_spend_update_transactions[ + "key_list_transactions" + ] + elif dict_key == "team_list_transactions": + transactions_dict = db_spend_update_transactions[ + "team_list_transactions" + ] + elif dict_key == "team_member_list_transactions": + transactions_dict = db_spend_update_transactions[ + "team_member_list_transactions" + ] + elif dict_key == "org_list_transactions": + transactions_dict = db_spend_update_transactions[ + "org_list_transactions" + ] + else: + continue + if transactions_dict is None: transactions_dict = {} db_spend_update_transactions[dict_key] = transactions_dict From 115946d402450e800d848636a6aa6159479dc724 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 21:25:24 -0700 Subject: [PATCH 17/20] unit testing for SpendUpdateQueue --- litellm/proxy/db/spend_update_queue.py | 20 ++- .../proxy/db/test_spend_update_queue.py | 152 ++++++++++++++++++ 2 files changed, 166 insertions(+), 6 deletions(-) create mode 100644 tests/litellm/proxy/db/test_spend_update_queue.py diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 29118073efa..fbd96c6c2df 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING, Any, List from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -72,14 +72,22 @@ class SpendUpdateQueue: for update in updates: entity_type = update.get("entity_type") - entity_id = update.get("entity_id") - response_cost = update.get("response_cost") + entity_id = update.get("entity_id") or "" + response_cost = update.get("response_cost") or 0 - if entity_type is None or entity_id is None or response_cost is None: - raise ValueError("Invalid update: %s", update) + if entity_type is None: + verbose_proxy_logger.debug( + "Skipping update spend for update: %s, because entity_type is None", + update, + ) + continue dict_key = entity_type_to_dict_key.get(entity_type) if dict_key is None: + verbose_proxy_logger.debug( + "Skipping update spend for update: %s, because entity_type is not in entity_type_to_dict_key", + update, + ) continue # Skip unknown entity types # Type-safe access using if/elif statements @@ -117,6 +125,6 @@ class SpendUpdateQueue: if entity_id not in transactions_dict: transactions_dict[entity_id] = 0 - transactions_dict[entity_id] += response_cost + transactions_dict[entity_id] += response_cost or 0 return db_spend_update_transactions diff --git a/tests/litellm/proxy/db/test_spend_update_queue.py b/tests/litellm/proxy/db/test_spend_update_queue.py new file mode 100644 index 00000000000..89d494a0701 --- /dev/null +++ b/tests/litellm/proxy/db/test_spend_update_queue.py @@ -0,0 +1,152 @@ +import asyncio +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem +from litellm.proxy.db.spend_update_queue import SpendUpdateQueue + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +@pytest.fixture +def spend_queue(): + return SpendUpdateQueue() + + +@pytest.mark.asyncio +async def test_add_update(spend_queue): + # Test adding a single update + update: SpendUpdateQueueItem = { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user123", + "response_cost": 0.5, + } + await spend_queue.add_update(update) + + # Verify update was added by checking queue size + assert spend_queue.update_queue.qsize() == 1 + + +@pytest.mark.asyncio +async def test_missing_response_cost(spend_queue): + # Test with missing response_cost - should default to 0 + update: SpendUpdateQueueItem = { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user123", + } + + await spend_queue.add_update(update) + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + # Should have created entry with 0 cost + assert aggregated["user_list_transactions"]["user123"] == 0 + + +@pytest.mark.asyncio +async def test_missing_entity_id(spend_queue): + # Test with missing entity_id - should default to empty string + update: SpendUpdateQueueItem = { + "entity_type": Litellm_EntityType.USER, + "response_cost": 1.0, + } + + await spend_queue.add_update(update) + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + # Should use empty string as key + assert aggregated["user_list_transactions"][""] == 1.0 + + +@pytest.mark.asyncio +async def test_none_values(spend_queue): + # Test with None values + update: SpendUpdateQueueItem = { + "entity_type": Litellm_EntityType.USER, + "entity_id": None, # type: ignore + "response_cost": None, + } + + await spend_queue.add_update(update) + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + # Should handle None values gracefully + assert aggregated["user_list_transactions"][""] == 0 + + +@pytest.mark.asyncio +async def test_multiple_updates_with_missing_fields(spend_queue): + # Test multiple updates with various missing fields + updates: list[SpendUpdateQueueItem] = [ + { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user123", + "response_cost": 0.5, + }, + { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user123", # missing response_cost + }, + { + "entity_type": Litellm_EntityType.USER, # missing entity_id + "response_cost": 1.5, + }, + ] + + for update in updates: + await spend_queue.add_update(update) + + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + # Verify aggregation + assert ( + aggregated["user_list_transactions"]["user123"] == 0.5 + ) # only the first update with valid cost + assert ( + aggregated["user_list_transactions"][""] == 1.5 + ) # update with missing entity_id + + +@pytest.mark.asyncio +async def test_unknown_entity_type(spend_queue): + # Test with unknown entity type + update: SpendUpdateQueueItem = { + "entity_type": "UNKNOWN_TYPE", # type: ignore + "entity_id": "123", + "response_cost": 0.5, + } + + await spend_queue.add_update(update) + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + # Should ignore unknown entity type + assert all(len(transactions) == 0 for transactions in aggregated.values()) + + +@pytest.mark.asyncio +async def test_missing_entity_type(spend_queue): + # Test with missing entity type + update: SpendUpdateQueueItem = {"entity_id": "123", "response_cost": 0.5} + + await spend_queue.add_update(update) + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + # Should ignore updates without entity type + assert all(len(transactions) == 0 for transactions in aggregated.values()) From d0a7e44a6ee232c8d2a75c5096304b61105c449f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 21:33:05 -0700 Subject: [PATCH 18/20] fix linting --- litellm/proxy/db/spend_update_queue.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index fbd96c6c2df..28e05246fad 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -120,7 +120,9 @@ class SpendUpdateQueue: if transactions_dict is None: transactions_dict = {} - db_spend_update_transactions[dict_key] = transactions_dict + + # type ignore: dict_key is guaranteed to be one of "one of ("user_list_transactions", "end_user_list_transactions", "key_list_transactions", "team_list_transactions", "team_member_list_transactions", "org_list_transactions")" + db_spend_update_transactions[dict_key] = transactions_dict # type: ignore if entity_id not in transactions_dict: transactions_dict[entity_id] = 0 From 7a2442d6c052dd1f82a56d70087f8400ea58d2e2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:12:29 -0700 Subject: [PATCH 19/20] test_batch_update_spend --- tests/local_testing/test_update_spend.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_update_spend.py b/tests/local_testing/test_update_spend.py index fffa3062d79..cc2c94af270 100644 --- a/tests/local_testing/test_update_spend.py +++ b/tests/local_testing/test_update_spend.py @@ -62,6 +62,8 @@ from litellm.proxy._types import ( KeyRequest, NewUserRequest, UpdateKeyRequest, + SpendUpdateQueueItem, + Litellm_EntityType, ) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) @@ -93,7 +95,13 @@ def prisma_client(): @pytest.mark.asyncio async def test_batch_update_spend(prisma_client): - prisma_client.user_list_transactions["test-litellm-user-5"] = 23 + await proxy_logging_obj.db_spend_update_writer.spend_update_queue.add_update( + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, + entity_id="test-litellm-user-5", + response_cost=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() From 55763ae276ba1413ce70d7e76057ca9e73e76d10 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:13:25 -0700 Subject: [PATCH 20/20] test_end_user_transactions_reset --- tests/proxy_unit_tests/test_proxy_utils.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 31c9711b0eb..a8c0ebe2ffa 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1521,12 +1521,6 @@ async def test_end_user_transactions_reset(): end_user_list_transactions=end_user_list_transactions, ) - # Verify cleanup happened - assert ( - mock_client.end_user_list_transactions == {} - ), "Transactions list should be empty after error" - - @pytest.mark.asyncio async def test_spend_logs_cleanup_after_error(): # Setup test data