diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0c525ee9466..7a4c1278d40 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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 diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a6105b6dff9..99b80e325c0 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index 19083486974..4b8903c973f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -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] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9a33a6bd758..ccc7b678330 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;