diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 220a0d5ddbe..ebe7d1e955c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2646,3 +2646,15 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase): models: Optional[List[str]] = Field( default=None, description="Default list of models that new users can access" ) + + +class DailyUserSpendTransaction(TypedDict): + user_id: str + date: str + api_key: str + model: str + model_group: Optional[str] + custom_llm_provider: Optional[str] + prompt_tokens: int + completion_tokens: int + spend: float diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 568253735c1..eefc0bdd06e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -558,6 +558,7 @@ async def proxy_startup_event(app: FastAPI): proxy_batch_write_at=proxy_batch_write_at, proxy_logging_obj=proxy_logging_obj, ) + ## [Optional] Initialize dd tracer ProxyStartupEvent._init_dd_tracer() @@ -914,6 +915,8 @@ def _set_spend_logs_payload( prisma_client.spend_log_transactions.append(payload) elif prisma_client is not None: prisma_client.spend_log_transactions.append(payload) + + prisma_client.add_spend_log_transaction_to_daily_user_transaction(payload.copy()) return prisma_client diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index e453e74b46f..9269e890142 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -313,3 +313,26 @@ model LiteLLM_AuditLog { before_value Json? // value of the row updated_values Json? // value of the row after change } + + +// Track daily user spend metrics per model and key +model LiteLLM_DailyUserSpend { + id String @id @default(uuid()) + user_id String + date String + api_key String // Hashed API Token + model String // The specific model used + model_group String? // public model_name / model_group + custom_llm_provider String? // The LLM provider (e.g., "openai", "anthropic") + prompt_tokens Int @default(0) + completion_tokens Int @default(0) + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([user_id, date, api_key, model, custom_llm_provider]) + @@index([date]) + @@index([user_id]) + @@index([api_key]) + @@index([model]) +} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0e7ae455960..7f1ac814a8f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -10,13 +10,15 @@ import traceback from datetime import datetime, timedelta from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText -from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union, overload +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, overload from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, CommonProxyErrors, + DailyUserSpendTransaction, ProxyErrorTypes, ProxyException, + SpendLogsPayload, ) from litellm.types.guardrails import GuardrailEventHooks @@ -1103,6 +1105,7 @@ class PrismaClient: team_member_list_transactons: dict = {} # key is ["team_id" + "user_id"] org_list_transactons: dict = {} spend_log_transactions: List = [] + daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {} def __init__( self, @@ -1140,6 +1143,61 @@ class PrismaClient: ) # Client to connect to Prisma db verbose_proxy_logger.debug("Success - Created Prisma Client") + def add_spend_log_transaction_to_daily_user_transaction( + self, payload: Union[dict, SpendLogsPayload] + ): + """ + Add a spend log transaction to the daily user transaction list + + Key = @@unique([user_id, date, api_key, model, custom_llm_provider]) ) + + If key exists, update the transaction with the new spend and usage + """ + expected_keys = ["user", "startTime", "api_key", "model", "custom_llm_provider"] + if not all(key in payload for key in expected_keys): + verbose_proxy_logger.debug( + f"Missing expected keys: {expected_keys}, in payload, skipping from daily_user_spend_transactions" + ) + return + + if isinstance(payload["startTime"], datetime): + start_time = payload["startTime"].isoformat() + date = start_time.split("T")[0] + elif isinstance(payload["startTime"], str): + date = payload["startTime"].split("T")[0] + else: + verbose_proxy_logger.debug( + f"Invalid start time: {payload['startTime']}, skipping from daily_user_spend_transactions" + ) + return + try: + daily_transaction_key = f"{payload['user']}_{date}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}" + if daily_transaction_key in self.daily_user_spend_transactions: + daily_transaction = self.daily_user_spend_transactions[ + daily_transaction_key + ] + daily_transaction["spend"] += payload["spend"] + daily_transaction["prompt_tokens"] += payload["prompt_tokens"] + daily_transaction["completion_tokens"] += payload["completion_tokens"] + else: + daily_transaction = DailyUserSpendTransaction( + user_id=payload["user"], + date=date, + api_key=payload["api_key"], + model=payload["model"], + model_group=payload["model_group"], + custom_llm_provider=payload["custom_llm_provider"], + prompt_tokens=payload["prompt_tokens"], + completion_tokens=payload["completion_tokens"], + spend=payload["spend"], + ) + + self.daily_user_spend_transactions[daily_transaction_key] = ( + daily_transaction + ) + except Exception as e: + raise e + def hash_token(self, token: str): # Hash the string using SHA-256 hashed_token = hashlib.sha256(token.encode()).hexdigest() @@ -2476,6 +2534,116 @@ class ProxyUpdateSpend: e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) + @staticmethod + async def update_daily_user_spend( + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + ): + """ + Batch job to update LiteLLM_DailyUserSpend table using in-memory daily_spend_transactions + """ + BATCH_SIZE = ( + 100 # Number of aggregated records to update in each database operation + ) + start_time = time.time() + + try: + for i in range(n_retry_times + 1): + try: + # Get transactions to process + transactions_to_process = dict( + list(prisma_client.daily_user_spend_transactions.items())[ + :BATCH_SIZE + ] + ) + + if len(transactions_to_process) == 0: + verbose_proxy_logger.debug( + "No new transactions to process for daily spend update" + ) + break + + # Update DailyUserSpend table in batches + async with prisma_client.db.batch_() as batcher: + for _, transaction in transactions_to_process.items(): + user_id = transaction.get("user_id") + if not user_id: # Skip if no user_id + continue + + batcher.litellm_dailyuserspend.upsert( + where={ + "user_id_date_api_key_model_custom_llm_provider": { + "user_id": user_id, + "date": transaction["date"], + "api_key": transaction["api_key"], + "model": transaction["model"], + "custom_llm_provider": transaction.get( + "custom_llm_provider" + ), + } + }, + data={ + "create": { + "user_id": user_id, + "date": transaction["date"], + "api_key": transaction["api_key"], + "model": transaction["model"], + "model_group": transaction.get("model_group"), + "custom_llm_provider": transaction.get( + "custom_llm_provider" + ), + "prompt_tokens": transaction["prompt_tokens"], + "completion_tokens": transaction[ + "completion_tokens" + ], + "spend": transaction["spend"], + }, + "update": { + "prompt_tokens": { + "increment": transaction["prompt_tokens"] + }, + "completion_tokens": { + "increment": transaction[ + "completion_tokens" + ] + }, + "spend": {"increment": transaction["spend"]}, + }, + }, + ) + + verbose_proxy_logger.info( + f"Processed {len(transactions_to_process)} daily spend transactions in {time.time() - start_time:.2f}s" + ) + + # Remove processed transactions + for key in transactions_to_process.keys(): + prisma_client.daily_user_spend_transactions.pop(key, None) + + verbose_proxy_logger.debug( + f"Processed {len(transactions_to_process)} daily spend transactions in {time.time() - start_time:.2f}s" + ) + break + + except DB_CONNECTION_ERROR_TYPES as e: + if i >= n_retry_times: + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + await asyncio.sleep(2**i) # Exponential backoff + + except Exception as e: + # Remove processed transactions even if there was an error + if "transactions_to_process" in locals(): + for key in transactions_to_process.keys(): # type: ignore + prisma_client.daily_user_spend_transactions.pop(key, None) + _raise_failed_update_spend_exception( + e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + ) + @staticmethod def disable_spend_updates() -> bool: """ @@ -2716,6 +2884,20 @@ async def update_spend( # noqa: PLR0915 db_writer_client=db_writer_client, ) + ### UPDATE DAILY USER SPEND ### + verbose_proxy_logger.debug( + "Daily User Spend transactions: {}".format( + len(prisma_client.daily_user_spend_transactions) + ) + ) + + if len(prisma_client.daily_user_spend_transactions) > 0: + await ProxyUpdateSpend.update_daily_user_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + def _raise_failed_update_spend_exception( e: Exception, start_time: float, proxy_logging_obj: ProxyLogging @@ -2911,3 +3093,50 @@ def _premium_user_check(): "error": f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}" }, ) + + +async def _update_daily_spend_batch(prisma_client, spend_aggregates): + """Helper function to update daily spend in batches""" + async with prisma_client.db.batch_() as batcher: + for ( + user_id, + date, + api_key, + model, + model_group, + provider, + ), metrics in spend_aggregates.items(): + if not user_id: # Skip if no user_id + continue + + batcher.litellm_dailyuserspend.upsert( + where={ + "user_id_date_api_key_model_custom_llm_provider": { + "user_id": user_id, + "date": date, + "api_key": api_key, + "model": model, + "custom_llm_provider": provider, + } + }, + data={ + "create": { + "user_id": user_id, + "date": date, + "api_key": api_key, + "model": model, + "model_group": model_group, + "custom_llm_provider": provider, + "prompt_tokens": metrics["prompt_tokens"], + "completion_tokens": metrics["completion_tokens"], + "spend": metrics["spend"], + }, + "update": { + "prompt_tokens": {"increment": metrics["prompt_tokens"]}, + "completion_tokens": { + "increment": metrics["completion_tokens"] + }, + "spend": {"increment": metrics["spend"]}, + }, + }, + ) diff --git a/schema.prisma b/schema.prisma index e453e74b46f..3312b26354d 100644 --- a/schema.prisma +++ b/schema.prisma @@ -313,3 +313,25 @@ model LiteLLM_AuditLog { before_value Json? // value of the row updated_values Json? // value of the row after change } + +// Track daily user spend metrics per model and key +model LiteLLM_DailyUserSpend { + id String @id @default(uuid()) + user_id String + date String + api_key String + model String + model_group String? + custom_llm_provider String? + prompt_tokens Int @default(0) + completion_tokens Int @default(0) + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([user_id, date, api_key, model, custom_llm_provider]) + @@index([date]) + @@index([user_id]) + @@index([api_key]) + @@index([model]) +}