From 4bcd60e72b60414ff4f5b7ba08b1ef4eb3f60452 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 14:51:05 -0700 Subject: [PATCH] perf(proxy): total a short grouped log page from the page itself A cursorless page that comes back without its lookahead row is the end of the list, so the total is offset + len(page) and the bounded grouped COUNT over the whole spend-log table is skipped. First pages on small deployments and every offset last page now cost one query less. Moves the count into _count_grouped_sessions and reworks the query-optimization test that asserted the count always runs second onto a full page, where it does. Claude-Session: https://claude.ai/code/session_01ESi9JwaXDww1vP3Qsrr4Mz --- .../spend_management_endpoints.py | 50 +++++++++++++------ .../test_spend_management_endpoints.py | 40 ++++++++++++++- .../test_spend_query_optimization.py | 14 +++--- 3 files changed, 79 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 2105d682d9a..582622cfb50 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2925,6 +2925,32 @@ async def _fetch_session_representatives( return [rep_by_key[key] for key in session_keys if key in rep_by_key] # mutable-ok: rows are enriched in place +async def _count_grouped_sessions( + prisma_client: "PrismaClient", + where_clause: str, + sql_params: Sequence[object], + next_param_index: int, +) -> tuple[int, bool]: + """Count the sessions matching the filter, returning ``(total, total_is_capped)`` bounded by the count cap.""" + count_query: Final = f""" + SELECT COUNT(*) AS total_count + FROM ( + SELECT 1 + FROM "LiteLLM_SpendLogs" + WHERE {where_clause} + GROUP BY {_SESSION_GROUP_KEY_SQL} + LIMIT ${next_param_index} + ) AS bounded_sessions + """ + count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( + prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + ) + raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 + return ( + (SPEND_LOGS_PAGINATION_COUNT_CAP, True) if raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP else (raw_total, False) + ) + + async def _ui_session_grouped_spend_logs( prisma_client: "PrismaClient", sql_conditions: Sequence[str], @@ -2953,7 +2979,9 @@ async def _ui_session_grouped_spend_logs( by its newest non-MCP row, enriched by ``_build_ui_spend_logs_response`` exactly like the flat listing, and the response carries ``next_session_cursor`` / ``has_more`` while ``total`` counts sessions - (capped like the flat total). + (capped like the flat total). A page that runs out of sessions is itself + the end of the list, so its ``total`` is ``offset + len(page)`` and the + grouped count query is skipped. """ where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE" cmp_op: Final = "<" if sort_desc else ">" @@ -2998,22 +3026,12 @@ async def _ui_session_grouped_spend_logs( else None ) - count_query: Final = f""" - SELECT COUNT(*) AS total_count - FROM ( - SELECT 1 - FROM "LiteLLM_SpendLogs" - WHERE {where_clause} - GROUP BY {_SESSION_GROUP_KEY_SQL} - LIMIT ${next_param_index} - ) AS bounded_sessions - """ - count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( - prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + page_ends_the_list: Final = cursor is None and page_limit > 0 and not has_more + total_records, total_is_capped = ( + (offset + len(page_rows), False) + if page_ends_the_list + else await _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index) ) - raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 - total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP - total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total session_keys: Final = tuple((row["session_key"], row["api_key"]) for row in visible_rows) data: Final[list[dict[str, object]]] = ( # mutable-ok: _build_ui_spend_logs_response writes onto each row 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 5920a984239..c709ac77a0b 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 @@ -6629,12 +6629,12 @@ def _session_page_row(session_key, last_activity): return {"session_key": session_key, "api_key": "hashed-key", "last_activity": last_activity} -def _session_grouped_paginating_prisma(sessions): +def _session_grouped_paginating_prisma(sessions, counted_total=None): """Mock prisma serving the grouped page query out of ``sessions``, honoring the LIMIT and OFFSET it asks for.""" async def mock_query_raw(sql_query, *params): if "COUNT(*) AS total_count" in sql_query: - return [{"total_count": min(len(sessions), params[-1])}] + return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}] if "DISTINCT ON" in sql_query: return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] if "COALESCE(SUM(spend)" in sql_query: @@ -6803,6 +6803,42 @@ async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor( app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_short_page_totals_itself(client, monkeypatch): + """A page that runs out of sessions is the end of the list, so the total comes from it and nothing is counted.""" + sessions = tuple((f"sess-{index:02d}", f"2026-08-29 10:{59 - index:02d}:00") for index in range(10)) + mock_prisma = _session_grouped_paginating_prisma(sessions, counted_total=999) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + 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, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": 1, + "page_size": 25, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["total"] == 10, "the count query's 999 would have won if it had been asked" + assert data["total_is_capped"] is False + assert data["total_pages"] == 1 + assert len(data["data"]) == 10 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_page_past_count_cap_is_empty(client, monkeypatch): """The last page inside the capped total still lists sessions; the page after it is empty and costs no query.""" 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 a7de3f1d8d6..54e5a6d5385 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 @@ -528,8 +528,8 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" session_rows = [ - {"session_key": "req-1", "api_key": "k", "last_activity": "2026-02-16 10:00:00"}, - {"session_key": "req-2", "api_key": "k", "last_activity": "2026-02-16 09:00:00"}, + {"session_key": f"req-{index}", "api_key": "k", "last_activity": f"2026-02-16 10:{59 - index:02d}:00"} + for index in range(51) ] representative_rows = [ {"request_id": "req-1", "api_key": "k", "metadata": "{}", "session_id": None}, @@ -538,7 +538,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): async def mock_query_raw(sql_query, *params): if "COUNT(*) AS total_count" in sql_query: - return [{"total_count": 12}] + return [{"total_count": 60}] if "DISTINCT ON" in sql_query: return representative_rows return session_rows @@ -590,11 +590,11 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): assert "COUNT(*) OVER ()" not in rep_sql assert [row["request_id"] for row in response["data"]] == ["req-1", "req-2"] - assert response["total"] == 12 + assert response["total"] == 60 assert response["total_is_capped"] is False - assert response["total_pages"] == 1 - assert response["has_more"] is False - assert response["next_session_cursor"] is None + assert response["total_pages"] == 2 + assert response["has_more"] is True + assert response["next_session_cursor"] == "2026-02-16 10:10:00|k|req-49" @pytest.mark.asyncio