From 55225a09cf9379ad81246bdf197a3ec2f4cc9394 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 11 Feb 2026 12:48:33 -0800 Subject: [PATCH] Allow sorting on /spend/logs/ui --- .../spend_management_endpoints.py | 34 +++- .../test_spend_management_endpoints.py | 192 +++++++++++++++++- 2 files changed, 222 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c05eda85c81..a853e764cdd 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1687,6 +1687,14 @@ async def ui_view_spend_logs( # noqa: PLR0915 error_message: Optional[str] = fastapi.Query( default=None, description="Filter logs by error message (partial string match)" ), + sort_by: Optional[str] = fastapi.Query( + default=None, + description="Sort logs by field: spend, total_tokens, startTime, or endTime (default: startTime)", + ), + sort_order: Optional[str] = fastapi.Query( + default="desc", + description="Sort order: asc or desc", + ), ): """ View spend logs with pagination support. @@ -1718,6 +1726,23 @@ async def ui_view_spend_logs( # noqa: PLR0915 code=status.HTTP_400_BAD_REQUEST, ) + # Validate sort_by and sort_order + valid_sort_fields = {"spend", "total_tokens", "startTime", "endTime"} + if sort_by is not None and sort_by not in valid_sort_fields: + raise ProxyException( + message=f"Invalid sort_by: {sort_by}. Must be one of: {', '.join(sorted(valid_sort_fields))}", + type="bad_request", + param="sort_by", + code=status.HTTP_400_BAD_REQUEST, + ) + if sort_order is not None and sort_order.lower() not in {"asc", "desc"}: + raise ProxyException( + message=f"Invalid sort_order: {sort_order}. Must be one of: asc, desc", + type="bad_request", + param="sort_order", + code=status.HTTP_400_BAD_REQUEST, + ) + try: is_v2 = "/spend/logs/v2" in request.url.path formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] @@ -1830,6 +1855,11 @@ async def ui_view_spend_logs( # noqa: PLR0915 # Calculate skip value for pagination skip = (page - 1) * page_size + # Build order clause from sort_by and sort_order + order_column = sort_by if sort_by else "startTime" + order_direction = (sort_order or "desc").lower() + order_clause = {order_column: order_direction} + # Get total count of records total_records = await prisma_client.db.litellm_spendlogs.count( where=where_conditions, @@ -1838,9 +1868,7 @@ async def ui_view_spend_logs( # noqa: PLR0915 # Get paginated data data = await prisma_client.db.litellm_spendlogs.find_many( where=where_conditions, - order={ - "startTime": "desc", - }, + order=order_clause, skip=skip, take=page_size, ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 08205cd2d9d..0d4a711812f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -12,7 +12,7 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import litellm import litellm.proxy.proxy_server as ps @@ -347,6 +347,196 @@ async def test_ui_view_spend_logs_with_user_id(client, monkeypatch): assert data["data"][0]["user"] == "test_user_1" +# Mock spend logs with distinct values for sorting tests. +# req_a: spend=0.10, tokens=500, start/end earliest +# req_b: spend=0.05, tokens=200, start/end 2nd +# req_c: spend=0.20, tokens=50, start/end latest +# req_d: spend=0.01, tokens=100, start/end 3rd +_SORT_TEST_LOGS = [ + { + "request_id": "req_a", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.10, + "total_tokens": 500, + "startTime": "2025-01-01T00:00:00+00:00", + "endTime": "2025-01-01T00:01:00+00:00", + "model": "gpt-3.5-turbo", + }, + { + "request_id": "req_b", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.05, + "total_tokens": 200, + "startTime": "2025-01-01T00:00:01+00:00", + "endTime": "2025-01-01T00:01:01+00:00", + "model": "gpt-3.5-turbo", + }, + { + "request_id": "req_c", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.20, + "total_tokens": 50, + "startTime": "2025-01-01T00:00:03+00:00", + "endTime": "2025-01-01T00:01:03+00:00", + "model": "gpt-3.5-turbo", + }, + { + "request_id": "req_d", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.01, + "total_tokens": 100, + "startTime": "2025-01-01T00:00:02+00:00", + "endTime": "2025-01-01T00:01:02+00:00", + "model": "gpt-3.5-turbo", + }, +] + + +def _sort_logs(logs, order_clause): + """Sort logs by the given Prisma-style order clause, e.g. {'spend': 'asc'}.""" + if not order_clause: + return list(logs) + key, direction = next(iter(order_clause.items())) + reverse = direction.lower() == "desc" + return sorted(logs, key=lambda x: x.get(key, 0), reverse=reverse) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sort_by,sort_order,expected_request_ids", + [ + # spend: 0.01(d) < 0.05(b) < 0.10(a) < 0.20(c) + ("spend", "asc", ["req_d", "req_b", "req_a", "req_c"]), + ("spend", "desc", ["req_c", "req_a", "req_b", "req_d"]), + # total_tokens: 50(c) < 100(d) < 200(b) < 500(a) + ("total_tokens", "asc", ["req_c", "req_d", "req_b", "req_a"]), + ("total_tokens", "desc", ["req_a", "req_b", "req_d", "req_c"]), + # startTime: 00:00:00(a) < 00:00:01(b) < 00:00:02(d) < 00:00:03(c) + ("startTime", "asc", ["req_a", "req_b", "req_d", "req_c"]), + ("startTime", "desc", ["req_c", "req_d", "req_b", "req_a"]), + # endTime: same ordering as startTime + ("endTime", "asc", ["req_a", "req_b", "req_d", "req_c"]), + ("endTime", "desc", ["req_c", "req_d", "req_b", "req_a"]), + # default when sort_by not provided: startTime desc + (None, "desc", ["req_c", "req_d", "req_b", "req_a"]), + ], +) +async def test_ui_view_spend_logs_sort_by_and_sort_order( + client, monkeypatch, sort_by, sort_order, expected_request_ids +): + """Test that spend logs are returned in the correct order for each sort_by/sort_order.""" + base_logs = list(_SORT_TEST_LOGS) + + async def mock_find_many(*args, **kwargs): + order = kwargs.get("order", {}) + return _sort_logs(base_logs, order) + + async def mock_count(*args, **kwargs): + return len(base_logs) + + class MockPrismaClient: + def __init__(self): + self.db = MagicMock() + self.db.litellm_spendlogs = MagicMock() + self.db.litellm_spendlogs.find_many = AsyncMock(side_effect=mock_find_many) + self.db.litellm_spendlogs.count = AsyncMock(side_effect=mock_count) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + start_date = "2024-12-25 00:00:00" + end_date = "2025-01-02 23:59:59" + + params = { + "start_date": start_date, + "end_date": end_date, + } + if sort_by is not None: + params["sort_by"] = sort_by + if sort_order is not None: + params["sort_order"] = sort_order + + response = client.get( + "/spend/logs/ui", + params=params, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200, response.text + data = response.json() + assert "data" in data + + actual_ids = [log["request_id"] for log in data["data"]] + assert actual_ids == expected_request_ids, ( + f"Expected order {expected_request_ids}, got {actual_ids} " + f"(sort_by={sort_by}, sort_order={sort_order})" + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sort_by,sort_order", + [ + ("invalid", "asc"), + ("spend", "invalid"), + ], +) +async def test_ui_view_spend_logs_sort_validation_errors( + client, monkeypatch, sort_by, sort_order +): + """Test that invalid sort_by and sort_order return 400.""" + async def mock_count(*args, **kwargs): + return 0 + + class MockPrismaClient: + def __init__(self): + self.db = MagicMock() + self.db.litellm_spendlogs = MagicMock() + self.db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + self.db.litellm_spendlogs.count = AsyncMock(side_effect=mock_count) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + start_date = "2024-12-25 00:00:00" + end_date = "2025-01-02 23:59:59" + + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "sort_by": sort_by, + "sort_order": sort_order, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 400 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_team_id(client, monkeypatch): # Mock data for the test