feat - viewing spend report per customer / team

This commit is contained in:
Ishaan Jaff 2024-06-10 12:58:24 -07:00
parent 6daeacb817
commit f00fa12e69

View file

@ -7814,6 +7814,10 @@ async def get_global_spend_report(
default=None, default=None,
description="Time till which to view spend", description="Time till which to view spend",
), ),
group_by: Optional[Literal["team", "customer"]] = fastapi.Query(
default="team",
description="Group spend by internal team or customer",
),
): ):
""" """
Get Daily Spend per Team, based on specific startTime and endTime. Per team, view usage by each key, model Get Daily Spend per Team, based on specific startTime and endTime. Per team, view usage by each key, model
@ -7860,69 +7864,130 @@ async def get_global_spend_report(
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
) )
# first get data from spend logs -> SpendByModelApiKey if group_by == "team":
# then read data from "SpendByModelApiKey" to format the response obj # first get data from spend logs -> SpendByModelApiKey
sql_query = """ # then read data from "SpendByModelApiKey" to format the response obj
sql_query = """
WITH SpendByModelApiKey AS ( WITH SpendByModelApiKey AS (
SELECT SELECT
date_trunc('day', sl."startTime") AS group_by_day, date_trunc('day', sl."startTime") AS group_by_day,
COALESCE(tt.team_alias, 'Unassigned Team') AS team_name, COALESCE(tt.team_alias, 'Unassigned Team') AS team_name,
sl.model, sl.model,
sl.api_key, sl.api_key,
SUM(sl.spend) AS model_api_spend, SUM(sl.spend) AS model_api_spend,
SUM(sl.total_tokens) AS model_api_tokens SUM(sl.total_tokens) AS model_api_tokens
FROM FROM
"LiteLLM_SpendLogs" sl "LiteLLM_SpendLogs" sl
LEFT JOIN LEFT JOIN
"LiteLLM_TeamTable" tt "LiteLLM_TeamTable" tt
ON ON
sl.team_id = tt.team_id sl.team_id = tt.team_id
WHERE WHERE
sl."startTime" BETWEEN $1::date AND $2::date sl."startTime" BETWEEN $1::date AND $2::date
GROUP BY GROUP BY
date_trunc('day', sl."startTime"), date_trunc('day', sl."startTime"),
tt.team_alias, tt.team_alias,
sl.model, sl.model,
sl.api_key sl.api_key
) )
SELECT
group_by_day,
jsonb_agg(jsonb_build_object(
'team_name', team_name,
'total_spend', total_spend,
'metadata', metadata
)) AS teams
FROM (
SELECT
group_by_day,
team_name,
SUM(model_api_spend) AS total_spend,
jsonb_agg(jsonb_build_object(
'model', model,
'api_key', api_key,
'spend', model_api_spend,
'total_tokens', model_api_tokens
)) AS metadata
FROM
SpendByModelApiKey
GROUP BY
group_by_day,
team_name
) AS aggregated
GROUP BY
group_by_day
ORDER BY
group_by_day;
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
if db_response is None:
return []
return db_response
elif group_by == "customer":
sql_query = """
WITH SpendByModelApiKey AS (
SELECT
date_trunc('day', sl."startTime") AS group_by_day,
sl.end_user AS customer,
sl.model,
sl.api_key,
SUM(sl.spend) AS model_api_spend,
SUM(sl.total_tokens) AS model_api_tokens
FROM
"LiteLLM_SpendLogs" sl
WHERE
sl."startTime" BETWEEN $1::date AND $2::date
GROUP BY
date_trunc('day', sl."startTime"),
customer,
sl.model,
sl.api_key
)
SELECT SELECT
group_by_day, group_by_day,
jsonb_agg(jsonb_build_object( jsonb_agg(jsonb_build_object(
'team_name', team_name, 'customer', customer,
'total_spend', total_spend, 'total_spend', total_spend,
'metadata', metadata 'metadata', metadata
)) AS teams )) AS customers
FROM ( FROM
SELECT (
group_by_day, SELECT
team_name, group_by_day,
SUM(model_api_spend) AS total_spend, customer,
jsonb_agg(jsonb_build_object( SUM(model_api_spend) AS total_spend,
'model', model, jsonb_agg(jsonb_build_object(
'api_key', api_key, 'model', model,
'spend', model_api_spend, 'api_key', api_key,
'total_tokens', model_api_tokens 'spend', model_api_spend,
)) AS metadata 'total_tokens', model_api_tokens
FROM )) AS metadata
SpendByModelApiKey FROM
GROUP BY SpendByModelApiKey
group_by_day, GROUP BY
team_name group_by_day,
) AS aggregated customer
) AS aggregated
GROUP BY GROUP BY
group_by_day group_by_day
ORDER BY ORDER BY
group_by_day; group_by_day;
""" """
db_response = await prisma_client.db.query_raw( db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj sql_query, start_date_obj, end_date_obj
) )
if db_response is None: if db_response is None:
return [] return []
return db_response return db_response
except Exception as e: except Exception as e:
raise HTTPException( raise HTTPException(