scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 43

perf(proxy): make /spend/logs/ui session listing use an index and run its queries concurrently

Browse source
This commit is contained in:
mubashir1osmani 2026-09-11 10:46:32 -04:00
parent 9a715df212
commit 46ef016f2f
6 changed files with 63 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm-proxy-extras/litellm_proxy_extras/migrations/20260911000000_spend_logs_session_window_index/migration.sql Normal file
View file

@ -0,0 +1,14 @@
-- CreateIndex (CONCURRENTLY)
--
-- Covering index for the session-grouped /spend/logs/ui listing. The page,
-- count, and representative queries all GROUP BY
-- COALESCE(NULLIF(session_id, ''), request_id), api_key over a startTime
-- window; with only the plain startTime index they aggregate full heap rows
-- (hundreds of KB each from messages/response), which is the 40s+ page load.
-- This index makes those window scans index-only.
--
-- CREATE INDEX CONCURRENTLY cannot run inside a transaction, so this
-- migration must stay a single statement (see
-- 20260415120000_health_check_latest_per_model_index for the full
-- disclaimer).
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_session_window_idx" ON "LiteLLM_SpendLogs"("startTime", session_id, request_id, api_key, status);

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -665,6 +665,7 @@ model LiteLLM_SpendLogs {
@@index([startTime, request_id])
@@index([end_user])
@@index([session_id])
@@index([startTime, session_id, request_id, api_key, status], map: "LiteLLM_SpendLogs_session_window_idx")
}
model LiteLLM_BudgetWindowSpend {

1
litellm/proxy/schema.prisma
View file

@ -665,6 +665,7 @@ model LiteLLM_SpendLogs {
@@index([startTime, request_id])
@@index([end_user])
@@index([session_id])
@@index([startTime, session_id, request_id, api_key, status], map: "LiteLLM_SpendLogs_session_window_idx")
}
model LiteLLM_BudgetWindowSpend {

65
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -1,8 +1,9 @@
#### SPEND MANAGEMENT #####
import asyncio
import collections
import json
import os
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Mapping, Sequence
from datetime import date, datetime, timedelta, timezone
from itertools import groupby
from types import MappingProxyType
@ -2821,13 +2822,6 @@ async def ui_view_spend_logs(
LIMIT ${p}
) AS bounded_matches
"""
count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_raw_or_none(
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
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
sql_query: Final = (
f"""
SELECT * FROM (
@ -2850,17 +2844,20 @@ async def ui_view_spend_logs(
LIMIT ${p} OFFSET ${p + 1}
"""
)
sql_params.extend([page_size, skip])
data: Final = await prisma_client.db.query_raw(sql_query, *sql_params)
count_task: Final[Awaitable[Sequence[_SpendLogsCountRow] | None]] = _query_raw_or_none(
prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
)
data_task: Final = prisma_client.db.query_raw(sql_query, *sql_params, page_size, skip)
count_rows, data = await asyncio.gather(count_task, data_task)
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
_hydrate_spend_log_metadata(data)
# Calculate total pages
total_pages: Final = (total_records + page_size - 1) // page_size
verbose_proxy_logger.debug("data= %s", json.dumps(data, indent=4, default=str))
return await _build_ui_spend_logs_response(
prisma_client,
data,
@ -2899,15 +2896,25 @@ async def _fetch_session_representatives(
next_param_index: int,
session_keys: Sequence[tuple[str, str]],
) -> list[dict[str, object]]: # mutable-ok: _build_ui_spend_logs_response writes session counts onto each row
"""Fetch the newest non-MCP row of each ``(session_key, api_key)`` session, in ``session_keys`` order."""
"""Fetch the newest non-MCP row of each ``(session_key, api_key)`` session, in ``session_keys`` order.
The SQL matches candidates on the plain ``session_id`` / ``request_id`` /
``api_key`` columns so the ``(startTime, session_id, request_id, api_key)``
index applies; a row-tuple ``IN`` on the COALESCE group key forces a full
window scan. That predicate admits a superset of the requested pairs (the
cross product of session ids and api keys, plus rows whose request_id
collides with a requested key), so the exact ``(session_key, api_key)``
pairing is restored by the ``rep_by_key`` lookup below.
"""
rep_query: Final = f"""
SELECT * FROM (
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
{_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
AND ({_SESSION_GROUP_KEY_SQL}) IN (
SELECT * FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[])
AND (
(session_id = ANY(${next_param_index}::text[]) AND api_key = ANY(${next_param_index + 1}::text[]))
OR request_id = ANY(${next_param_index}::text[])
)
ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
) AS session_representatives
@ -2975,18 +2982,6 @@ async def _ui_session_grouped_spend_logs(
ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
LIMIT ${limit_index}
"""
page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw(
prisma_client, page_query, *sql_params, *cursor_params, page_size + 1
)
has_more: Final = len(page_rows) > page_size
visible_rows: Final = page_rows[:page_size]
next_cursor: Final = (
f"{visible_rows[-1]['last_activity']}|{visible_rows[-1]['api_key']}|{visible_rows[-1]['session_key']}"
if has_more and visible_rows
else None
)
count_query: Final = f"""
SELECT COUNT(*) AS total_count
FROM (
@ -2997,9 +2992,21 @@ async def _ui_session_grouped_spend_logs(
LIMIT ${next_param_index}
) AS bounded_sessions
"""
count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw(
page_rows_task: Final[Awaitable[Sequence[_SessionPageRow]]] = _query_raw(
prisma_client, page_query, *sql_params, *cursor_params, page_size + 1
)
count_rows_task: Final[Awaitable[Sequence[_SpendLogsCountRow]]] = _query_raw(
prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
)
page_rows, count_rows = await asyncio.gather(page_rows_task, count_rows_task)
has_more: Final = len(page_rows) > page_size
visible_rows: Final = page_rows[:page_size]
next_cursor: Final = (
f"{visible_rows[-1]['last_activity']}|{visible_rows[-1]['api_key']}|{visible_rows[-1]['session_key']}"
if has_more and visible_rows
else None
)
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

1
schema.prisma
View file

@ -665,6 +665,7 @@ model LiteLLM_SpendLogs {
@@index([startTime, request_id])
@@index([end_user])
@@index([session_id])
@@index([startTime, session_id, request_id, api_key, status], map: "LiteLLM_SpendLogs_session_window_idx")
}
model LiteLLM_BudgetWindowSpend {

10
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -6640,6 +6640,10 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
reps = [
_session_representative_row("req-solo", None),
_session_representative_row("req-1", "sess-1"),
# The plain-column predicate admits the cross product of session ids and
# api keys, so the endpoint must drop rows whose (session_key, api_key)
# pair was never requested.
_session_representative_row("req-cross-product", "sess-unrequested"),
]
mock_prisma = _session_grouped_mock_prisma(page_rows, 3, reps)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
@ -6689,6 +6693,12 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
f"ORDER BY {SESSION_GROUP_KEY_SQL}, call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC"
in rep_call[0]
), "the session representative must prefer the newest non-MCP call"
assert "session_id = ANY(" in rep_call[0] and "OR request_id = ANY(" in rep_call[0], (
"the representative lookup must filter on the plain session_id/request_id columns so the "
"(startTime, session_id, request_id, api_key, status) index applies; a row-tuple IN on the "
"COALESCE group key forces a full window scan (the 40s+ logs page)"
)
assert f"({SESSION_GROUP_KEY_SQL}) IN" not in rep_call[0]
assert rep_call[-2] == ["sess-1", "req-solo"]
assert rep_call[-1] == ["hashed-key", "hashed-key"]
finally: