mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(spend): add exclude_team_models param to /global/spend/report
Team BYOK deployments share their public model name with admin-managed models, so billing pulled from /global/spend/report cannot distinguish spend a team already paid its provider for, and that spend gets billed twice. Zeroing the deployment cost is not a workaround because spend drives member-level budget enforcement. When exclude_team_models=true, every SQL branch of the report anti-joins LiteLLM_SpendLogs.model_id against LiteLLM_ProxyModelTable rows whose model_info carries a team_id. Ownership via model_info.team_id (the same signal Router._is_team_specific_model reads) survives deployment renames and never matches admin models, unlike the mangled model_name prefix. With the flag absent or false every query is byte-identical to before. Known limitation: spend from deployments deleted since the request was logged cannot be classified and remains included.
This commit is contained in:
parent
2227bd5c2c
commit
e118ebec55
4 changed files with 222 additions and 14 deletions
|
|
@ -29,6 +29,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
get_spend_by_team,
|
||||
get_spend_by_team_and_customer,
|
||||
team_model_exclusion_clause,
|
||||
)
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
from litellm.repositories.table_repositories import SpendLogsRepository
|
||||
|
|
@ -991,6 +992,13 @@ async def get_global_spend_report(
|
|||
default=None,
|
||||
description="View spend for a specific customer_id. Example customer_id='1234. Can be used in conjunction with team_id as well.",
|
||||
),
|
||||
exclude_team_models: bool = fastapi.Query(
|
||||
default=False,
|
||||
description=(
|
||||
"When True, exclude spend from team-owned (BYOK) model deployments. Spend from "
|
||||
"deployments that have since been deleted cannot be classified and remains included"
|
||||
),
|
||||
),
|
||||
):
|
||||
"""
|
||||
Get Daily Spend per Team, based on specific startTime and endTime. Per team, view usage by each key, model
|
||||
|
|
@ -1041,11 +1049,14 @@ async def get_global_spend_report(
|
|||
if premium_user is not True:
|
||||
verbose_proxy_logger.debug("accessing /spend/report but not a premium user")
|
||||
raise ValueError("/spend/report endpoint " + CommonProxyErrors.not_premium_user.value)
|
||||
|
||||
exclusion_sql = team_model_exclusion_clause(exclude_team_models)
|
||||
|
||||
if api_key is not None:
|
||||
verbose_proxy_logger.debug("Getting /spend for api_key: [set=%s]", api_key is not None)
|
||||
if api_key.startswith("sk-"):
|
||||
api_key = hash_token(token=api_key)
|
||||
sql_query = """
|
||||
sql_query = f"""
|
||||
WITH SpendByModelApiKey AS (
|
||||
SELECT
|
||||
sl.api_key,
|
||||
|
|
@ -1058,7 +1069,7 @@ async def get_global_spend_report(
|
|||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl.api_key = $3
|
||||
AND sl.api_key = $3{exclusion_sql}
|
||||
GROUP BY
|
||||
sl.api_key,
|
||||
sl.model
|
||||
|
|
@ -1088,7 +1099,7 @@ async def get_global_spend_report(
|
|||
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 = """
|
||||
sql_query = f"""
|
||||
WITH SpendByModelApiKey AS (
|
||||
SELECT
|
||||
sl.api_key,
|
||||
|
|
@ -1101,7 +1112,7 @@ async def get_global_spend_report(
|
|||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl.user = $3
|
||||
AND sl.user = $3{exclusion_sql}
|
||||
GROUP BY
|
||||
sl.api_key,
|
||||
sl.model
|
||||
|
|
@ -1131,13 +1142,24 @@ async def get_global_spend_report(
|
|||
return db_response
|
||||
elif team_id is not None and customer_id is not None:
|
||||
return await get_spend_by_team_and_customer(
|
||||
start_date_obj, end_date_obj, team_id, customer_id, prisma_client
|
||||
start_date_obj,
|
||||
end_date_obj,
|
||||
team_id,
|
||||
customer_id,
|
||||
prisma_client,
|
||||
exclude_team_models=exclude_team_models,
|
||||
)
|
||||
if group_by == "team":
|
||||
return await get_spend_by_team(start_date_obj, end_date_obj, team_id, prisma_client)
|
||||
return await get_spend_by_team(
|
||||
start_date_obj,
|
||||
end_date_obj,
|
||||
team_id,
|
||||
prisma_client,
|
||||
exclude_team_models=exclude_team_models,
|
||||
)
|
||||
|
||||
elif group_by == "customer":
|
||||
sql_query = """
|
||||
sql_query = f"""
|
||||
|
||||
WITH SpendByModelApiKey AS (
|
||||
SELECT
|
||||
|
|
@ -1151,7 +1173,7 @@ async def get_global_spend_report(
|
|||
"LiteLLM_SpendLogs" sl
|
||||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC'){exclusion_sql}
|
||||
GROUP BY
|
||||
date_trunc('day', sl."startTime"),
|
||||
customer,
|
||||
|
|
@ -1195,7 +1217,7 @@ async def get_global_spend_report(
|
|||
|
||||
return db_response
|
||||
elif group_by == "api_key":
|
||||
sql_query = """
|
||||
sql_query = f"""
|
||||
WITH SpendByModelApiKey AS (
|
||||
SELECT
|
||||
sl.api_key,
|
||||
|
|
@ -1207,7 +1229,7 @@ async def get_global_spend_report(
|
|||
"LiteLLM_SpendLogs" sl
|
||||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC'){exclusion_sql}
|
||||
GROUP BY
|
||||
sl.api_key,
|
||||
sl.model
|
||||
|
|
|
|||
|
|
@ -510,13 +510,35 @@ def _ensure_datetime_utc(timestamp: datetime) -> datetime:
|
|||
return timestamp
|
||||
|
||||
|
||||
def team_model_exclusion_clause(exclude_team_models: bool) -> str:
|
||||
"""
|
||||
SQL predicate that drops spend rows produced by team-owned (BYOK) deployments.
|
||||
|
||||
Ownership is `model_info.team_id` (what `Router._is_team_specific_model` reads), so the
|
||||
exclusion survives a rename of the deployment's mangled `model_name`. `->>` yields NULL
|
||||
on an absent key or a non-object `model_info`, which fails open to "included".
|
||||
|
||||
Only a literal True opts in, so every existing query stays byte-identical; in-process
|
||||
callers that omit the argument leave FastAPI's `Query` default object in place, and that
|
||||
object is truthy.
|
||||
"""
|
||||
if exclude_team_models is not True:
|
||||
return ""
|
||||
return (
|
||||
'\n AND NOT EXISTS (SELECT 1 FROM "LiteLLM_ProxyModelTable" pm '
|
||||
"WHERE pm.model_id = sl.model_id AND pm.model_info ->> 'team_id' IS NOT NULL)"
|
||||
)
|
||||
|
||||
|
||||
async def get_spend_by_team(
|
||||
start_date: dt,
|
||||
end_date: dt,
|
||||
team_id: Optional[str],
|
||||
prisma_client: PrismaClient,
|
||||
exclude_team_models: bool = False,
|
||||
):
|
||||
sql_query = """
|
||||
exclusion_sql = team_model_exclusion_clause(exclude_team_models)
|
||||
sql_query = f"""
|
||||
WITH SpendByModelApiKey AS (
|
||||
SELECT
|
||||
date_trunc('day', sl."startTime") AS group_by_day,
|
||||
|
|
@ -534,7 +556,7 @@ async def get_spend_by_team(
|
|||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND ($3::text IS NULL OR sl.team_id = $3)
|
||||
AND ($3::text IS NULL OR sl.team_id = $3){exclusion_sql}
|
||||
GROUP BY
|
||||
date_trunc('day', sl."startTime"),
|
||||
tt.team_alias,
|
||||
|
|
@ -584,8 +606,10 @@ async def get_spend_by_team_and_customer(
|
|||
team_id: str,
|
||||
customer_id: str,
|
||||
prisma_client: PrismaClient,
|
||||
exclude_team_models: bool = False,
|
||||
):
|
||||
sql_query = """
|
||||
exclusion_sql = team_model_exclusion_clause(exclude_team_models)
|
||||
sql_query = f"""
|
||||
WITH SpendByModelApiKey AS (
|
||||
SELECT
|
||||
date_trunc('day', sl."startTime") AS group_by_day,
|
||||
|
|
@ -605,7 +629,7 @@ async def get_spend_by_team_and_customer(
|
|||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl.team_id = $3
|
||||
AND sl.end_user = $4
|
||||
AND sl.end_user = $4{exclusion_sql}
|
||||
GROUP BY
|
||||
date_trunc('day', sl."startTime"),
|
||||
tt.team_alias,
|
||||
|
|
|
|||
|
|
@ -502,3 +502,163 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch):
|
|||
params = mock_prisma.db.query_raw.call_args[0][1:]
|
||||
assert "team_x" in params, "team_id must be forwarded into the DB query params"
|
||||
assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}"
|
||||
|
||||
|
||||
_OMITTED = object()
|
||||
|
||||
_EXCLUSION_PREDICATE = (
|
||||
'\n AND NOT EXISTS (SELECT 1 FROM "LiteLLM_ProxyModelTable" pm '
|
||||
"WHERE pm.model_id = sl.model_id AND pm.model_info ->> 'team_id' IS NOT NULL)"
|
||||
)
|
||||
|
||||
_REPORT_BRANCHES = {
|
||||
"api_key_filter": {"api_key": "sk-1234"},
|
||||
"internal_user_id_filter": {"internal_user_id": "user-1"},
|
||||
"team_and_customer_filter": {"team_id": "team-1", "customer_id": "cust-1"},
|
||||
"group_by_team": {"group_by": "team"},
|
||||
"group_by_customer": {"group_by": "customer"},
|
||||
"group_by_api_key": {"group_by": "api_key"},
|
||||
}
|
||||
|
||||
|
||||
async def _capture_report_sql(monkeypatch, branch_kwargs, exclude_team_models):
|
||||
"""
|
||||
Run GET /global/spend/report for one branch and return the SQL it executed.
|
||||
|
||||
group_by is always passed explicitly because the priority chain (api_key ->
|
||||
internal_user_id -> team+customer -> group_by) reads the raw arguments; `_OMITTED`
|
||||
calls the endpoint the way every caller predating the flag does, with no keyword.
|
||||
"""
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
get_global_spend_report,
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
await get_global_spend_report(
|
||||
**{
|
||||
"start_date": "2026-07-01",
|
||||
"end_date": "2026-07-03",
|
||||
"group_by": "team",
|
||||
"api_key": None,
|
||||
"internal_user_id": None,
|
||||
"team_id": None,
|
||||
"customer_id": None,
|
||||
**branch_kwargs,
|
||||
**({} if exclude_team_models is _OMITTED else {"exclude_team_models": exclude_team_models}),
|
||||
}
|
||||
)
|
||||
|
||||
assert mock_prisma.db.query_raw.called, "query_raw should have been called"
|
||||
return mock_prisma.db.query_raw.call_args[0][0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("branch", sorted(_REPORT_BRANCHES))
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_spend_report_excludes_team_models_when_opted_in(monkeypatch, branch):
|
||||
"""
|
||||
Every branch must honor the flag, not just the default group_by=team one. Non-team rows
|
||||
survive the anti-join either with no "LiteLLM_ProxyModelTable" match (config models log a
|
||||
router-generated hash id) or with a match carrying no model_info.team_id.
|
||||
`_EXCLUSION_PREDICATE` pins the clause's exact leading whitespace and is spelled out
|
||||
rather than imported, so a mutated production clause fails here.
|
||||
"""
|
||||
sql = await _capture_report_sql(monkeypatch, _REPORT_BRANCHES[branch], exclude_team_models=True)
|
||||
|
||||
assert _EXCLUSION_PREDICATE in sql, f"{branch} must exclude team deployments. SQL was:\n{sql}"
|
||||
assert sql.index(_EXCLUSION_PREDICATE) < sql.index("GROUP BY"), (
|
||||
f"the exclusion must sit in the WHERE clause of the SpendLogs CTE, before its GROUP BY. SQL was:\n{sql}"
|
||||
)
|
||||
assert "NOT EXISTS" in sql and "IN (SELECT" not in sql, (
|
||||
f"an anti-join (NOT EXISTS) keeps rows with an empty model_id; a subquery IN would drop them. SQL was:\n{sql}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("branch", sorted(_REPORT_BRANCHES))
|
||||
@pytest.mark.parametrize("exclude_team_models", [False, _OMITTED], ids=["explicit_false", "omitted"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_spend_report_keeps_team_models_by_default(monkeypatch, branch, exclude_team_models):
|
||||
"""
|
||||
Opt-in only: with the flag off (or absent) the executed SQL must be exactly what it
|
||||
was before the flag existed, so existing /global/spend/report consumers see no change.
|
||||
"""
|
||||
sql = await _capture_report_sql(monkeypatch, _REPORT_BRANCHES[branch], exclude_team_models=exclude_team_models)
|
||||
|
||||
assert "LiteLLM_ProxyModelTable" not in sql, f"{branch} must not touch the model table by default. SQL was:\n{sql}"
|
||||
assert "NOT EXISTS" not in sql, f"{branch} must not filter deployments by default. SQL was:\n{sql}"
|
||||
|
||||
opted_in = await _capture_report_sql(monkeypatch, _REPORT_BRANCHES[branch], exclude_team_models=True)
|
||||
assert opted_in.replace(_EXCLUSION_PREDICATE, "") == sql, (
|
||||
f"the flag must only add the exclusion predicate, nothing else. SQL was:\n{sql}"
|
||||
)
|
||||
|
||||
|
||||
def test_global_spend_report_binds_exclude_team_models_from_query_string(monkeypatch):
|
||||
"""
|
||||
Calling the endpoint function directly cannot catch a binding regression, so drive the
|
||||
real route once: `?exclude_team_models=true` has to reach the query builder as a bool.
|
||||
"""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="admin",
|
||||
)
|
||||
try:
|
||||
client = TestClient(app)
|
||||
params = {"start_date": "2026-07-01", "end_date": "2026-07-03", "group_by": "team"}
|
||||
|
||||
response = client.get("/global/spend/report", params={**params, "exclude_team_models": "true"})
|
||||
assert response.status_code == 200, response.text
|
||||
assert _EXCLUSION_PREDICATE in mock_prisma.db.query_raw.call_args[0][0]
|
||||
|
||||
response = client.get("/global/spend/report", params=params)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _EXCLUSION_PREDICATE not in mock_prisma.db.query_raw.call_args[0][0]
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_by_team_helpers_default_to_including_team_models():
|
||||
"""
|
||||
Both helpers are public and called positionally elsewhere; the new keyword must
|
||||
default to off and only add the exclusion when a caller asks for it.
|
||||
"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_query_raw = AsyncMock(return_value=[])
|
||||
mock_prisma.db.query_raw = mock_query_raw
|
||||
|
||||
start_date = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc)
|
||||
end_date = datetime.datetime(2024, 1, 31, tzinfo=timezone.utc)
|
||||
|
||||
await get_spend_by_team(start_date, end_date, "team-1", mock_prisma)
|
||||
assert "LiteLLM_ProxyModelTable" not in mock_query_raw.call_args[0][0]
|
||||
|
||||
await get_spend_by_team_and_customer(start_date, end_date, "team-1", "cust-1", mock_prisma)
|
||||
assert "LiteLLM_ProxyModelTable" not in mock_query_raw.call_args[0][0]
|
||||
|
||||
await get_spend_by_team(start_date, end_date, "team-1", mock_prisma, exclude_team_models=True)
|
||||
assert _EXCLUSION_PREDICATE in mock_query_raw.call_args[0][0]
|
||||
|
||||
await get_spend_by_team_and_customer(
|
||||
start_date, end_date, "team-1", "cust-1", mock_prisma, exclude_team_models=True
|
||||
)
|
||||
assert _EXCLUSION_PREDICATE in mock_query_raw.call_args[0][0]
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -41045,6 +41045,8 @@ export interface operations {
|
|||
team_id?: string | null;
|
||||
/** @description View spend for a specific customer_id. Example customer_id='1234. Can be used in conjunction with team_id as well. */
|
||||
customer_id?: string | null;
|
||||
/** @description When True, exclude spend from team-owned (BYOK) model deployments. Spend from deployments that have since been deleted cannot be classified and remains included */
|
||||
exclude_team_models?: boolean;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue