From f691ff7005d8d3263c69e705f3563eb28147392e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 6 Jul 2024 14:45:58 -0700 Subject: [PATCH] get spend per internal user / api_key --- .../spend_management_endpoints.py | 138 ++++++++++++------ 1 file changed, 94 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 6f887cc1bd..9e0ef877e8 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -825,6 +825,10 @@ async def get_global_spend_report( default=None, description="View spend for a specific api_key. Example api_key='sk-1234", ), + internal_user_id: Optional[str] = fastapi.Query( + default=None, + description="View spend for a specific internal_user_id. Example internal_user_id='1234", + ), ): """ Get Daily Spend per Team, based on specific startTime and endTime. Per team, view usage by each key, model @@ -877,6 +881,96 @@ async def get_global_spend_report( raise ValueError( "/spend/report endpoint " + CommonProxyErrors.not_premium_user.value ) + if api_key is not None: + verbose_proxy_logger.debug("Getting /spend for api_key: %s", api_key) + if api_key.startswith("sk-"): + api_key = hash_token(token=api_key) + sql_query = """ + WITH SpendByModelApiKey AS ( + SELECT + sl.api_key, + sl.model, + SUM(sl.spend) AS model_cost, + SUM(sl.prompt_tokens) AS model_input_tokens, + SUM(sl.completion_tokens) AS model_output_tokens + FROM + "LiteLLM_SpendLogs" sl + WHERE + sl."startTime" BETWEEN $1::date AND $2::date AND sl.api_key = $3 + GROUP BY + sl.api_key, + sl.model + ) + SELECT + api_key, + SUM(model_cost) AS total_cost, + SUM(model_input_tokens) AS total_input_tokens, + SUM(model_output_tokens) AS total_output_tokens, + jsonb_agg(jsonb_build_object( + 'model', model, + 'total_cost', model_cost, + 'total_input_tokens', model_input_tokens, + 'total_output_tokens', model_output_tokens + )) AS model_details + FROM + SpendByModelApiKey + GROUP BY + api_key + ORDER BY + total_cost DESC; + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date_obj, end_date_obj, api_key + ) + if db_response is None: + return [] + + return db_response + elif internal_user_id is not None: + verbose_proxy_logger.debug( + "Getting /spend for internal_user_id: %s", internal_user_id + ) + sql_query = """ + WITH SpendByModelApiKey AS ( + SELECT + sl.api_key, + sl.model, + SUM(sl.spend) AS model_cost, + SUM(sl.prompt_tokens) AS model_input_tokens, + SUM(sl.completion_tokens) AS model_output_tokens + FROM + "LiteLLM_SpendLogs" sl + WHERE + sl."startTime" BETWEEN $1::date AND $2::date AND sl.user = $3 + GROUP BY + sl.api_key, + sl.model + ) + SELECT + api_key, + SUM(model_cost) AS total_cost, + SUM(model_input_tokens) AS total_input_tokens, + SUM(model_output_tokens) AS total_output_tokens, + jsonb_agg(jsonb_build_object( + 'model', model, + 'total_cost', model_cost, + 'total_input_tokens', model_input_tokens, + 'total_output_tokens', model_output_tokens + )) AS model_details + FROM + SpendByModelApiKey + GROUP BY + api_key + ORDER BY + total_cost DESC; + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date_obj, end_date_obj, internal_user_id + ) + if db_response is None: + return [] + + return db_response if group_by == "team": # first get data from spend logs -> SpendByModelApiKey @@ -1044,50 +1138,6 @@ async def get_global_spend_report( return [] return db_response - elif api_key is not None: - if api_key.startswith("sk-"): - api_key = hash_token(token=api_key) - sql_query = """ - WITH SpendByModelApiKey AS ( - SELECT - sl.api_key, - sl.model, - SUM(sl.spend) AS model_cost, - SUM(sl.prompt_tokens) AS model_input_tokens, - SUM(sl.completion_tokens) AS model_output_tokens - FROM - "LiteLLM_SpendLogs" sl - WHERE - sl."startTime" BETWEEN $1::date AND $2::date AND sl.api_key = $3 - GROUP BY - sl.api_key, - sl.model - ) - SELECT - api_key, - SUM(model_cost) AS total_cost, - SUM(model_input_tokens) AS total_input_tokens, - SUM(model_output_tokens) AS total_output_tokens, - jsonb_agg(jsonb_build_object( - 'model', model, - 'total_cost', model_cost, - 'total_input_tokens', model_input_tokens, - 'total_output_tokens', model_output_tokens - )) AS model_details - FROM - SpendByModelApiKey - GROUP BY - api_key - ORDER BY - total_cost DESC; - """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj, api_key - ) - if db_response is None: - return [] - - return db_response except Exception as e: raise HTTPException(