diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5482661c3de..20cd62bf81b 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -552,6 +552,20 @@ async def spend_user_fn( ) +def _parse_utc_day(value: str | None, param_name: str) -> datetime | None: + """Parse an optional YYYY-MM-DD query param to a UTC datetime, so the SQL + day window never depends on the database session timezone.""" + if value is None: + return None + try: + return datetime.strptime(value, "%Y-%m-%d").replace(tzinfo=timezone.utc) + except ValueError: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": f"{param_name} must be a YYYY-MM-DD date, got {value!r}"}, + ) + + @router.get( "/spend/tags", tags=["Budget & Spend Tracking"], @@ -561,6 +575,8 @@ async def spend_user_fn( }, ) async def view_spend_tags( + log_team_lookup: LogTeamLookupDependency, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], start_date: str | None = fastapi.Query( default=None, description="Time from which to start viewing key spend", @@ -573,6 +589,11 @@ async def view_spend_tags( """ LiteLLM Enterprise - View Spend Per Request Tag + - Admin callers (PROXY_ADMIN / PROXY_ADMIN_VIEW_ONLY) see tag aggregates + across every spend log. + - All other callers see aggregates over their own spend logs and those of + teams they may view, exactly like ``/spend/logs/ui``. + Example Request: ``` curl -X GET "http://0.0.0.0:8000/spend/tags" \ @@ -594,16 +615,19 @@ async def view_spend_tags( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - # run the following SQL query on prisma - """ - SELECT - jsonb_array_elements_text(request_tags) AS individual_request_tag, - COUNT(*) AS log_count, - SUM(spend) AS total_spend - FROM "LiteLLM_SpendLogs" - GROUP BY individual_request_tag; - """ - response: Final = await get_spend_by_tags(start_date=start_date, end_date=end_date, prisma_client=prisma_client) + read_scope: Final = ( + AllRows() + if is_admin_view_safe(user_api_key_dict=user_api_key_dict) + else await _spend_log_read_scope(user_api_key_dict, log_team_lookup) + if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict) + else OwnedRows(user_api_key_dict.user_id) + ) + response: Final = await get_spend_by_tags( + prisma_client=prisma_client, + read_scope=read_scope, + start_date=_parse_utc_day(start_date, "start_date"), + end_date=_parse_utc_day(end_date, "end_date"), + ) return response except Exception as e: @@ -4377,17 +4401,41 @@ async def provider_budgets() -> ProviderBudgetResponse: raise handle_exception_on_proxy(e) -async def get_spend_by_tags(prisma_client: PrismaClient, start_date=None, end_date=None): +async def get_spend_by_tags( + prisma_client: PrismaClient, + read_scope: ReadScope, + start_date: datetime | None = None, + end_date: datetime | None = None, +): + scope_clause, scope_params = read_scope_sql(read_scope, 1) + start_clause: Final = ( + (f"\"startTime\" >= (${len(scope_params) + 1}::timestamptz AT TIME ZONE 'UTC')",) if start_date else () + ) + end_clause: Final = ( + ( + f"\"startTime\" < ((${len(scope_params) + len(start_clause) + 1}::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')", + ) + if end_date + else () + ) + clauses: Final = tuple(clause for clause in (scope_clause, *start_clause, *end_clause) if clause) + where_sql: Final = f" WHERE {' AND '.join(clauses)}" if clauses else "" + params: Final = ( + *scope_params, + *((start_date,) if start_date else ()), + *((end_date,) if end_date else ()), + ) response: Final[Sequence[Mapping[str, object]]] = await _query_raw( prisma_client, - """ + f""" SELECT jsonb_array_elements_text(request_tags) AS individual_request_tag, COUNT(*) AS log_count, SUM(spend) AS total_spend - FROM "LiteLLM_SpendLogs" + FROM "LiteLLM_SpendLogs"{where_sql} GROUP BY individual_request_tag; """, + *params, ) return response diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 25180038e2a..f6b36c184ec 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1,4 +1,3 @@ -from typing import Final import asyncio import collections import datetime @@ -7,6 +6,7 @@ import json import re import sqlite3 from datetime import timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import prisma @@ -278,8 +278,8 @@ from litellm.proxy._types import ( from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.management.teams import authz as team_access from litellm.proxy.proxy_server import app -from litellm.proxy.utils import PrismaClient from litellm.proxy.spend_tracking import spend_management_endpoints +from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.utils import BudgetConfig @@ -4120,7 +4120,7 @@ async def test_view_spend_tags(client, monkeypatch): ] # Mock the get_spend_by_tags function - async def mock_get_spend_by_tags(prisma_client, start_date=None, end_date=None): + async def mock_get_spend_by_tags(prisma_client, read_scope, start_date=None, end_date=None): return mock_response monkeypatch.setattr( @@ -4168,6 +4168,116 @@ async def test_view_spend_tags(client, monkeypatch): app.dependency_overrides.pop(ps.user_api_key_auth, None) +def _tag_query_prisma_client(assert_query_raw): + class MockDB: + async def query_raw(self, sql_query, *params): + return assert_query_raw(sql_query, params) + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + + return MockPrismaClient() + + +@pytest.mark.asyncio +async def test_view_spend_tags_scopes_internal_user_to_own_and_team_logs(client, monkeypatch): + def assert_query_raw(sql_query, params): + assert '("user" = $1 OR team_id = ANY($2::text[]))' in sql_query + assert params[0] == "user-1" + assert tuple(params[1]) == ("team-9",) + return [{"individual_request_tag": "own-tag", "log_count": 1, "total_spend": 0.01}] + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _tag_query_prisma_client(assert_query_raw)) + + async def permitted_teams(*args, **kwargs): + return ["team-9"] + + monkeypatch.setitem(app.dependency_overrides, get_log_team_lookup, lambda: permitted_teams) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1" + ) + + try: + response = client.get("/spend/tags", headers={"Authorization": "Bearer sk-test"}) + assert response.status_code == 200 + assert [row["individual_request_tag"] for row in response.json()] == ["own-tag"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_view_spend_tags_admin_sees_unscoped_aggregates_with_date_filter(client, monkeypatch): + def assert_query_raw(sql_query, params): + assert '"user"' not in sql_query + assert "team_id" not in sql_query + assert '"startTime" >= ($1' in sql_query + assert "$2::timestamptz + INTERVAL '1 day'" in sql_query + assert params == ( + datetime.datetime(2024, 1, 1, tzinfo=timezone.utc), + datetime.datetime(2024, 1, 31, tzinfo=timezone.utc), + ), "dates must be bound as UTC datetimes so the window is session-timezone independent" + return [{"individual_request_tag": "any-tag", "log_count": 2, "total_spend": 0.05}] + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _tag_query_prisma_client(assert_query_raw)) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + response = client.get( + "/spend/tags", + params={"start_date": "2024-01-01", "end_date": "2024-01-31"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + assert len(response.json()) == 1 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_view_spend_tags_rejects_malformed_dates_with_400(client, monkeypatch): + def assert_query_raw(sql_query, params): + pytest.fail("malformed dates must never reach the database") + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _tag_query_prisma_client(assert_query_raw)) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + response = client.get( + "/spend/tags", + params={"start_date": "01/31/2024"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 400 + assert "start_date" in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_view_spend_tags_key_without_user_id_gets_no_rows(client, monkeypatch): + def assert_query_raw(sql_query, params): + assert "WHERE FALSE" in sql_query + assert params == () + return [] + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _tag_query_prisma_client(assert_query_raw)) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.TEAM, user_id=None + ) + + try: + response = client.get("/spend/tags", headers={"Authorization": "Bearer sk-test"}) + assert response.status_code == 200 + assert response.json() == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_view_spend_tags_no_database(client, monkeypatch): """Test /spend/tags endpoint when database is not connected""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c3dde3b4d47..73cd88d2e50 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15634,6 +15634,11 @@ export interface paths { * View Spend Tags * @description LiteLLM Enterprise - View Spend Per Request Tag * + * - Admin callers (PROXY_ADMIN / PROXY_ADMIN_VIEW_ONLY) see tag aggregates + * across every spend log. + * - All other callers see aggregates over their own spend logs and those of + * teams they may view, exactly like ``/spend/logs/ui``. + * * Example Request: * ``` * curl -X GET "http://0.0.0.0:8000/spend/tags" -H "Authorization: Bearer $LITELLM_MASTER_KEY"