""" Handles buffering database `UPDATE` transactions in Redis before committing them to the database This is to prevent deadlocks and improve reliability """ import asyncio import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import (DailyAgentSpendTransaction, DailyEndUserSpendTransaction, DailyOrganizationSpendTransaction, DailyTagSpendTransaction, DailyTeamSpendTransaction, DailyUserSpendTransaction, DBSpendUpdateTransactions) from litellm.proxy.db.db_transaction_queue.base_update_queue import \ service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \ DailySpendUpdateQueue from litellm.proxy.db.db_transaction_queue.spend_update_queue import \ SpendUpdateQueue from litellm.secret_managers.main import str_to_bool from litellm.types.caching import (RedisPipelineLpopOperation, RedisPipelineRpushOperation) from litellm.types.services import ServiceTypes if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient else: PrismaClient = Any class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database This is to prevent deadlocks and improve reliability """ def __init__( self, redis_cache: Optional[RedisCache] = None, ): self.redis_cache = redis_cache @staticmethod def _should_commit_spend_updates_to_redis() -> bool: """ Checks if the Pod should commit spend updates to Redis This setting enables buffering database transactions in Redis to improve reliability and reduce database contention """ from litellm.proxy.proxy_server import general_settings _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( general_settings.get("use_redis_transaction_buffer", False) ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) if _use_redis_transaction_buffer is None: return False return _use_redis_transaction_buffer async def _store_transactions_in_redis( self, transactions: Any, redis_key: str, service_type: ServiceTypes, ) -> None: """ Helper method to store transactions in Redis and emit an event Args: transactions: The transactions to store redis_key: The Redis key to store under service_type: The service type for event emission """ if transactions is None or len(transactions) == 0: return list_of_transactions = [safe_dumps(transactions)] if self.redis_cache is None: return try: current_redis_buffer_size = await self.redis_cache.async_rpush( key=redis_key, values=list_of_transactions, ) verbose_proxy_logger.debug( "Spend tracking - pushed spend updates to Redis buffer. " "redis_key=%s, buffer_size=%s", redis_key, current_redis_buffer_size, ) await self._emit_new_item_added_to_redis_buffer_event( queue_size=current_redis_buffer_size, service=service_type, ) except Exception as e: verbose_proxy_logger.error( "Spend tracking - failed to push spend updates to Redis (redis_key=%s). " "Error: %s", redis_key, str(e), ) async def store_in_memory_spend_updates_in_redis( self, spend_update_queue: SpendUpdateQueue, daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, daily_tag_spend_update_queue: DailySpendUpdateQueue, ): """ Stores the in-memory spend updates to Redis Stores the following in memory data structures in Redis: - SpendUpdateQueue - Key, User, Team, TeamMember, Org, EndUser Spend updates - DailySpendUpdateQueue - Daily Spend updates Aggregate view For SpendUpdateQueue: Each transaction is a dict stored as following: - key is the entity id - value is the spend amount ``` Redis List: key_list_transactions: [ "0929880201": 1.2, "0929880202": 0.01, "0929880203": 0.001, ] ``` For DailySpendUpdateQueue: Each transaction is a Dict[str, DailyUserSpendTransaction] stored as following: - key is the daily_transaction_key - value is the DailyUserSpendTransaction ``` Redis List: daily_spend_update_transactions: [ { "user_keyhash_1_model_1": { "spend": 1.2, "prompt_tokens": 1000, "completion_tokens": 1000, "api_requests": 1000, "successful_requests": 1000, }, } ] ``` """ if self.redis_cache is None: verbose_proxy_logger.debug( "redis_cache is None, skipping store_in_memory_spend_updates_in_redis" ) return # Get all transactions db_spend_update_transactions = ( await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() ) daily_spend_update_transactions = ( await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) daily_team_spend_update_transactions = ( await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) daily_org_spend_update_transactions = ( await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) daily_end_user_spend_update_transactions = ( await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) daily_agent_spend_update_transactions = ( await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) daily_tag_spend_update_transactions = ( await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) verbose_proxy_logger.debug( "ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions ) verbose_proxy_logger.debug( "ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions ) # Build a list of rpush operations, skipping empty/None transaction sets _queue_configs: List[Tuple[Any, str, ServiceTypes]] = [ (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE), (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE), (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE), (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE), ] rpush_list: List[RedisPipelineRpushOperation] = [] service_types: List[ServiceTypes] = [] for transactions, redis_key, service_type in _queue_configs: if transactions is None or len(transactions) == 0: continue rpush_list.append( RedisPipelineRpushOperation( key=redis_key, values=[safe_dumps(transactions)], ) ) service_types.append(service_type) if len(rpush_list) == 0: return result_lengths = await self.redis_cache.async_rpush_pipeline( rpush_list=rpush_list, ) # Emit gauge events for each queue for i, queue_size in enumerate(result_lengths): if i < len(service_types): await self._emit_new_item_added_to_redis_buffer_event( queue_size=queue_size, service=service_types[i], ) @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, ) -> int: """ Gets the number of transactions to store in Redis """ num_transactions = 0 for v in db_spend_update_transactions.values(): if isinstance(v, dict): num_transactions += len(v) return num_transactions @staticmethod def _remove_prefix_from_keys(data: Dict[str, Any], prefix: str) -> Dict[str, Any]: """ Removes the specified prefix from the keys of a dictionary. """ return {key.replace(prefix, "", 1): value for key, value in data.items()} async def get_all_update_transactions_from_redis_buffer( self, ) -> Optional[DBSpendUpdateTransactions]: """ Gets all the update transactions from Redis On Redis we store a list of transactions as a JSON string eg. [ DBSpendUpdateTransactions( user_list_transactions={ "user_id_1": 1.2, "user_id_2": 0.01, }, end_user_list_transactions={}, key_list_transactions={ "0929880201": 1.2, "0929880202": 0.01, }, team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, ), DBSpendUpdateTransactions( user_list_transactions={ "user_id_3": 1.2, "user_id_4": 0.01, }, end_user_list_transactions={}, key_list_transactions={ "key_id_1": 1.2, "key_id_2": 0.01, }, team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, ] """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None verbose_proxy_logger.info( "Spend tracking - popped %d spend update batches from Redis buffer (key=%s). " "These items are now removed from Redis and must be committed to DB.", len(list_of_transactions) if isinstance(list_of_transactions, list) else 1, REDIS_UPDATE_BUFFER_KEY, ) # Parse the list of transactions from JSON strings parsed_transactions = self._parse_list_of_transactions(list_of_transactions) # If there are no transactions, return None if len(parsed_transactions) == 0: return None # Combine all transactions into a single transaction combined_transaction = self._combine_list_of_transactions(parsed_transactions) return combined_transaction async def get_all_transactions_from_redis_buffer_pipeline( self, ) -> Tuple[ Optional[DBSpendUpdateTransactions], Optional[Dict[str, DailyUserSpendTransaction]], Optional[Dict[str, DailyTeamSpendTransaction]], Optional[Dict[str, DailyOrganizationSpendTransaction]], Optional[Dict[str, DailyEndUserSpendTransaction]], Optional[Dict[str, DailyAgentSpendTransaction]], Optional[Dict[str, DailyTagSpendTransaction]], ]: """ Drains all 7 Redis buffer queues in a single pipeline round-trip. Returns a 7-tuple of parsed results in this order: 0: DBSpendUpdateTransactions 1: daily user spend 2: daily team spend 3: daily org spend 4: daily end-user spend 5: daily agent spend 6: daily tag spend """ if self.redis_cache is None: return None, None, None, None, None, None, None lpop_list: List[RedisPipelineLpopOperation] = [ RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation(key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation(key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation(key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation(key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation(key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), ] raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) # Pad with None if pipeline returned fewer results than expected while len(raw_results) < 7: raw_results.append(None) # Slot 0: DBSpendUpdateTransactions db_spend: Optional[DBSpendUpdateTransactions] = None if raw_results[0] is not None: parsed = self._parse_list_of_transactions(raw_results[0]) if len(parsed) > 0: db_spend = self._combine_list_of_transactions(parsed) # Slots 1-6: daily spend categories daily_results: List[Optional[Dict[str, Any]]] = [] for slot in range(1, 7): if raw_results[slot] is None: daily_results.append(None) else: list_of_daily = [json.loads(t) for t in raw_results[slot]] # type: ignore aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily ) daily_results.append(aggregated) return ( db_spend, cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]), cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]), cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]), cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]), cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]), cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]), ) async def get_all_daily_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyUserSpendTransaction]]: """ Gets all the daily spend update transactions from Redis """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None list_of_daily_spend_update_transactions = [ json.loads(transaction) for transaction in list_of_transactions ] return cast( Dict[str, DailyUserSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), ) async def get_all_daily_team_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyTeamSpendTransaction]]: """ Gets all the daily team spend update transactions from Redis """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None list_of_daily_spend_update_transactions = [ json.loads(transaction) for transaction in list_of_transactions ] return cast( Dict[str, DailyTeamSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), ) async def get_all_daily_org_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyOrganizationSpendTransaction]]: """ Gets all the daily organization spend update transactions from Redis """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None list_of_daily_spend_update_transactions = [ json.loads(transaction) for transaction in list_of_transactions ] return cast( Dict[str, DailyOrganizationSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), ) async def get_all_daily_end_user_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyEndUserSpendTransaction]]: """ Gets all the daily end-user spend update transactions from Redis """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None list_of_daily_spend_update_transactions = [ json.loads(transaction) for transaction in list_of_transactions ] return cast( Dict[str, DailyEndUserSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), ) async def get_all_daily_agent_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyAgentSpendTransaction]]: """ Gets all the daily agent spend update transactions from Redis """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None list_of_daily_spend_update_transactions = [ json.loads(transaction) for transaction in list_of_transactions ] return cast( Dict[str, DailyAgentSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), ) async def get_all_daily_tag_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyTagSpendTransaction]]: """ Gets all the daily tag spend update transactions from Redis """ if self.redis_cache is None: return None list_of_transactions = await self.redis_cache.async_lpop( key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) if list_of_transactions is None: return None list_of_daily_spend_update_transactions = [ json.loads(transaction) for transaction in list_of_transactions ] return cast( Dict[str, DailyTagSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), ) @staticmethod def _parse_list_of_transactions( list_of_transactions: Union[Any, List[Any]], ) -> List[DBSpendUpdateTransactions]: """ Parses the list of transactions from Redis """ if isinstance(list_of_transactions, list): return [json.loads(transaction) for transaction in list_of_transactions] else: return [json.loads(list_of_transactions)] @staticmethod def _combine_list_of_transactions( list_of_transactions: List[DBSpendUpdateTransactions], ) -> DBSpendUpdateTransactions: """ Combines the list of transactions into a single DBSpendUpdateTransactions object """ # Initialize a new combined transaction object with empty dictionaries combined_transaction = DBSpendUpdateTransactions( user_list_transactions={}, end_user_list_transactions={}, key_list_transactions={}, team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, ) # Define the transaction fields to process transaction_fields = [ "user_list_transactions", "end_user_list_transactions", "key_list_transactions", "team_list_transactions", "team_member_list_transactions", "org_list_transactions", "tag_list_transactions", "agent_list_transactions", ] # Loop through each transaction and combine the values for transaction in list_of_transactions: # Process each field type for field in transaction_fields: if transaction.get(field): for entity_id, amount in transaction[field].items(): # type: ignore combined_transaction[field][entity_id] = ( # type: ignore combined_transaction[field].get(entity_id, 0) + amount # type: ignore ) return combined_transaction async def _emit_new_item_added_to_redis_buffer_event( self, service: ServiceTypes, queue_size: int, ): asyncio.create_task( service_logger_obj.async_service_success_hook( service=service, duration=0, call_type="_emit_new_item_added_to_queue_event", event_metadata={ "gauge_labels": service, "gauge_value": queue_size, }, ) )