scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

fix(spend): read only the newest rows when paging sessions (LIT-8513) (#48886)

Browse source
This commit is contained in:
shrey 2026-11-03 10:51:09 -08:00 • committed by GitHub
parent add7dea9aa
commit 0aa7079f04
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1117 additions and 239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

221
litellm/proxy/spend_tracking/session_page.py Normal file
View file

@ -0,0 +1,221 @@
"""Which sessions a logs page shows and how many match, read from the newest rows instead of every row. Reads grow
until one settles both, so the result is what grouping every row would give."""
import math
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, TypeAlias
from typing_extensions import ReadOnly, TypedDict
MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)"
SESSION_REPRESENTATIVE_ORDER_SQL: Final = (
f"(call_type = {AGENT_CALL_TYPE_SQL}) DESC, "
f'CASE WHEN call_type = {AGENT_CALL_TYPE_SQL} THEN "endTime" END DESC NULLS LAST, '
f'call_type IN {MCP_CALL_TYPES_SQL}, "startTime" DESC, request_id'
)
class _ReadRow(TypedDict):
session_key: ReadOnly[str | None]
api_key: ReadOnly[str | None]
last_activity: ReadOnly[str | None]
settled: ReadOnly[bool | None]
rows_read: ReadOnly[int]
session_count: ReadOnly[int]
class _SupportsQueryRaw(Protocol):
async def query_raw(self, query: str, *args: object) -> Sequence[_ReadRow]: ...
@dataclass(frozen=True, slots=True)
class Session:
key: str
api_key: str
last_activity: str
@dataclass(frozen=True, slots=True)
class SessionPage:
sessions: tuple[Session, ...]
has_more: bool
next_cursor: str | None
total: int
total_is_capped: bool
@dataclass(frozen=True, slots=True)
class _Request:
cursor: tuple[str, str, str] | None
offset: int
page_limit: int
sort_desc: bool
status: str | None
cap: int
def first_read_size(self) -> int | None:
"""Every row when oldest first, or when filtered by outcome, which any of a session's rows can decide. Otherwise
the fewest that could settle a page: a session past the cap and an older row."""
return self.cap + 2 if self.sort_desc and self.status is None else None
@dataclass(frozen=True, slots=True)
class _Read:
candidates: tuple[Session, ...]
last_settled: bool
rows_read: int
session_count: int
saw_every_row: bool
def settles(self, request: _Request) -> bool:
page_settled: Final = (
request.page_limit <= 0
or self.saw_every_row
or (len(self.candidates) > request.page_limit and self.last_settled)
)
return page_settled and (self.saw_every_row or self.session_count > request.cap)
def next_size(self, request: _Request) -> int:
rows_per_session: Final = self.rows_read / self.session_count
return max(2 * self.rows_read, math.ceil(1.25 * (request.cap + 1) * rows_per_session))
_Reader: TypeAlias = Callable[[int | None], Awaitable[_Read]]
def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] | None:
if session_cursor is None or session_cursor.count("|") < 2:
return None
last_activity, _, rest = session_cursor.partition("|")
api_key, _, session_key = rest.partition("|")
if not last_activity or not session_key:
return None
return (last_activity, session_key, api_key)
async def read_session_page(
db: _SupportsQueryRaw,
*,
where_clause: str,
params: Sequence[object],
status: str | None,
page: int,
page_size: int,
sort_desc: bool,
session_cursor: str | None,
cap: int,
) -> SessionPage:
"""One page of sessions after ``session_cursor``, or after ``(page - 1) * page_size`` sessions, keeping only those
whose outcome is ``status`` when given. ``where_clause`` reads ``params`` as ``$1`` onward."""
cursor: Final = _parse_session_cursor(session_cursor)
offset: Final = (page - 1) * page_size if cursor is None else 0
request: Final = _Request(cursor, offset, min(page_size, cap - offset), sort_desc, status, cap)
return await _session_page(_sql_reader(db, where_clause, params, request), request)
async def _session_page(read: _Reader, request: _Request) -> SessionPage:
last = await read(request.first_read_size()) # rebind-ok: each read replaces the last until one settles the page
while not last.settles(request):
last = await read(last.next_size(request))
has_more: Final = len(last.candidates) > request.page_limit
sessions: Final = last.candidates[: request.page_limit]
final: Final = sessions[-1] if has_more and sessions else None
return SessionPage(
sessions=sessions,
has_more=has_more,
next_cursor=f"{final.last_activity}|{final.api_key}|{final.key}" if final else None,
total=min(last.session_count, request.cap),
total_is_capped=last.session_count > request.cap,
)
def _sessions_sql(request: _Request, status_index: int) -> str:
"""Each session's last activity, from its rows in ``scanned``, keeping only sessions whose outcome is the status."""
if request.status is None:
return (
'SELECT session_key, api_key, MAX("startTime") AS last_activity FROM scanned GROUP BY session_key, api_key'
)
return f"""
SELECT session_key, api_key, last_activity
FROM (
SELECT DISTINCT ON (session_key, api_key)
session_key,
api_key,
status,
MAX("startTime") OVER (PARTITION BY session_key, api_key) AS last_activity
FROM scanned
ORDER BY session_key, api_key, {SESSION_REPRESENTATIVE_ORDER_SQL}
) AS outcomes
WHERE COALESCE(status, 'success') = ${status_index}
"""
def _sql_reader(db: _SupportsQueryRaw, where_clause: str, params: Sequence[object], request: _Request) -> _Reader:
status_index: Final = len(params) + 1
status_args: Final = () if request.status is None else (request.status,)
cursor_index: Final = status_index + len(status_args)
after_cursor: Final = (
f"WHERE (last_activity, session_key, api_key) {'<' if request.sort_desc else '>'} "
f"(${cursor_index}::timestamp, ${cursor_index + 1}, ${cursor_index + 2})"
if request.cursor
else ""
)
outcome_columns: Final = ', status, call_type, "endTime", request_id' if request.status is not None else ""
direction: Final = "DESC" if request.sort_desc else "ASC"
page_rows: Final = request.page_limit + 1 if request.page_limit > 0 else 0
skip: Final = f"OFFSET {request.offset}" if page_rows and request.offset else ""
async def read(size: int | None) -> _Read:
# LEFT JOIN keeps the summary row when no session is on the page.
rows: Final = await db.query_raw(
f"""
WITH scanned AS (
SELECT {SESSION_KEY_EXPR} AS session_key, api_key, "startTime"{outcome_columns}
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
{f'ORDER BY "startTime" DESC LIMIT {size}' if size is not None else ""}
),
sessions AS ({_sessions_sql(request, status_index)}),
summary AS (
SELECT COUNT(*) AS rows_read,
MIN("startTime") AS oldest,
(SELECT COUNT(*) FROM sessions) AS session_count
FROM scanned
)
SELECT page.session_key,
page.api_key,
page.last_activity::text AS last_activity,
page.last_activity > summary.oldest AS settled,
summary.rows_read,
summary.session_count
FROM summary
LEFT JOIN LATERAL (
SELECT session_key, api_key, last_activity
FROM sessions
{after_cursor}
ORDER BY last_activity {direction}, session_key {direction}, api_key {direction}
LIMIT {page_rows} {skip}
) AS page ON TRUE
ORDER BY page.last_activity {direction}, page.session_key {direction}, page.api_key {direction}
""",
*params,
*status_args,
*(request.cursor or ()),
)
rows_read: Final = int(rows[0]["rows_read"])
return _Read(
candidates=tuple(
Session(str(row["session_key"]), str(row["api_key"]), str(row["last_activity"]))
for row in rows
if row["session_key"] is not None
),
last_settled=bool(rows[-1]["settled"]),
rows_read=rows_read,
session_count=int(rows[0]["session_count"]),
saw_every_row=size is None or rows_read < size,
)
return read

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

@ -52,6 +52,13 @@ from litellm.proxy.auth.authorization_dependencies import (
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.openai_error_payload import error_status_code, openai_error_param, openai_error_type
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.spend_tracking.session_page import (
AGENT_CALL_TYPE_SQL,
MCP_CALL_TYPES_SQL,
SESSION_KEY_EXPR,
SESSION_REPRESENTATIVE_ORDER_SQL,
read_session_page,
)
from litellm.proxy.spend_tracking.spend_capture_rate import (
ProviderBillingCredentialMissing,
ProviderBillingRequestFailed,
@ -87,23 +94,15 @@ router: Final = APIRouter()
SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000
_DECODED_JSON: Final = TypeAdapter(object)
_SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)"
_SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key"
_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
_AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
_SESSION_REPRESENTATIVE_ORDER_SQL: Final = (
f"(call_type = {_AGENT_CALL_TYPE_SQL}) DESC, "
f'CASE WHEN call_type = {_AGENT_CALL_TYPE_SQL} THEN "endTime" END DESC NULLS LAST, '
f'call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC, request_id'
)
_SESSION_GROUP_KEY_SQL: Final = f"{SESSION_KEY_EXPR}, api_key"
_BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')"
_SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType(
{
"mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}",
"agent": f"call_type = {_AGENT_CALL_TYPE_SQL}",
"mcp": f"call_type IN {MCP_CALL_TYPES_SQL}",
"agent": f"call_type = {AGENT_CALL_TYPE_SQL}",
"batch": f"call_type IN {_BATCH_CALL_TYPES_SQL}",
"llm": (
f"(call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} "
f"(call_type NOT IN {MCP_CALL_TYPES_SQL} AND call_type != {AGENT_CALL_TYPE_SQL} "
f"AND call_type NOT IN {_BATCH_CALL_TYPES_SQL})"
),
}
@ -2948,23 +2947,6 @@ async def ui_view_spend_logs(
sql_params.append(json.dumps(used_client_oauth_token))
p += 1
if status_filter is not None and group_by_session is True and not is_search_lookup:
session_filter_conditions: Final = " AND ".join(sql_conditions) or "TRUE"
sql_conditions.append(
f"""({_SESSION_GROUP_KEY_SQL}) IN (
SELECT session_key, api_key FROM (
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
{_SESSION_KEY_EXPR} AS session_key, api_key, status
FROM "LiteLLM_SpendLogs"
WHERE {session_filter_conditions}
ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL}
) AS session_outcomes
WHERE COALESCE(status, 'success') = ${p}
)"""
)
sql_params.append(status_filter)
p += 1
if (
group_by_session is True
and not is_v2
@ -2977,12 +2959,30 @@ async def ui_view_spend_logs(
sql_conditions=sql_conditions,
sql_params=sql_params,
next_param_index=p,
status_filter=status_filter,
page=page,
page_size=page_size,
sort_desc=order_direction != "asc",
session_cursor=session_cursor,
)
if status_filter is not None and group_by_session is True and not is_search_lookup:
session_filter_conditions: Final = " AND ".join(sql_conditions) or "TRUE"
sql_conditions.append(
f"""({_SESSION_GROUP_KEY_SQL}) IN (
SELECT session_key, api_key FROM (
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
{SESSION_KEY_EXPR} AS session_key, api_key, status
FROM "LiteLLM_SpendLogs"
WHERE {session_filter_conditions}
ORDER BY {_SESSION_GROUP_KEY_SQL}, {SESSION_REPRESENTATIVE_ORDER_SQL}
) AS session_outcomes
WHERE COALESCE(status, 'success') = ${p}
)"""
)
sql_params.append(status_filter)
p += 1
# Build the ORDER BY expression. ttft_ms is computed from
# completionStartTime - startTime; non-streaming rows (where
# completionStartTime is null or equals endTime) yield NULL, so we
@ -3031,7 +3031,7 @@ async def ui_view_spend_logs(
{_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
WHERE {joined_conditions}
ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL}
ORDER BY {_SESSION_GROUP_KEY_SQL}, {SESSION_REPRESENTATIVE_ORDER_SQL}
) AS session_representatives
ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause}, request_id
LIMIT ${p} OFFSET ${p + 1}
@ -3080,22 +3080,6 @@ async def ui_view_spend_logs(
raise handle_exception_on_proxy(e)
class _SessionPageRow(TypedDict):
session_key: ReadOnly[str]
api_key: ReadOnly[str]
last_activity: ReadOnly[str]
def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] | None:
if session_cursor is None or session_cursor.count("|") < 2:
return None
last_activity, _, rest = session_cursor.partition("|")
api_key, _, session_key = rest.partition("|")
if not last_activity or not session_key:
return None
return (last_activity, session_key, api_key)
async def _fetch_session_representatives(
prisma_client: "PrismaClient",
where_clause: str,
@ -3104,16 +3088,18 @@ async def _fetch_session_representatives(
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 final agent outcome, or newest non-MCP row, of each ``(session_key, api_key)`` session, in ``session_keys`` order."""
# Lets Postgres use the session_id and request_id indexes.
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_id = ANY(${next_param_index}::text[]) OR request_id = ANY(${next_param_index}::text[]))
AND ({_SESSION_GROUP_KEY_SQL}) IN (
SELECT * FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[])
)
ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL}
ORDER BY {_SESSION_GROUP_KEY_SQL}, {SESSION_REPRESENTATIVE_ORDER_SQL}
) AS session_representatives
"""
rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place
@ -3129,117 +3115,32 @@ async def _fetch_session_representatives(
return [rep_by_key[key] for key in session_keys if key in rep_by_key]
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],
sql_params: Sequence[object],
next_param_index: int,
status_filter: str | None,
page: int,
page_size: int,
sort_desc: bool,
session_cursor: str | None,
) -> Mapping[str, object]:
"""
One row per session, keyset-paginated by session last activity.
Sessions are derived on the fly from ``LiteLLM_SpendLogs`` (no extra
table): rows sharing a ``session_id`` and ``api_key`` form a session, rows
without a session id are singletons keyed by ``request_id``. A page is the
next ``page_size`` sessions ordered by ``(MAX(startTime), session_key,
api_key)``, resumed from the ``session_cursor`` keyset
``'<last_activity>|<api_key>|<session_key>'`` instead of an OFFSET, so
page depth does not degrade the query plan. A request for ``page > 1``
without a cursor (the UI jumping straight to the last page, or back to a
page it never walked through) falls back to ``OFFSET (page - 1) *
page_size``, trimmed to the end of the ``SPEND_LOGS_PAGINATION_COUNT_CAP``
window the capped ``total`` promises, so a page never runs past that total
and one starting at or past it returns no rows without a query. Each session is represented
by its final agent outcome (or 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). A page that runs out of sessions while still
holding some is itself the end of the list, so its ``total`` is
``offset + len(page)`` and the grouped count query is skipped; a page that
starts past the end says nothing about the total, so that one is counted.
"""
"""One row per session, keyset-paginated by last activity. ``read_session_page`` picks and counts the sessions,
and each is shown by its final agent outcome or newest non-MCP row."""
where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE"
cmp_op: Final = "<" if sort_desc else ">"
direction: Final = "DESC" if sort_desc else "ASC"
cursor: Final = _parse_session_cursor(session_cursor)
having_clause: Final = (
f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} '
f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})"
if cursor
else ""
session_page: Final = await read_session_page(
prisma_client.db, # pyright: ignore[reportArgumentType] # PrismaWrapper forwards raw queries via __getattr__
where_clause=where_clause,
params=sql_params,
status=status_filter,
page=page,
page_size=page_size,
sort_desc=sort_desc,
session_cursor=session_cursor,
cap=SPEND_LOGS_PAGINATION_COUNT_CAP,
)
cursor_params: Final[tuple[object, ...]] = cursor if cursor else ()
limit_index: Final = next_param_index + len(cursor_params)
offset: Final = (page - 1) * page_size if cursor is None else 0
page_limit: Final = min(page_size, SPEND_LOGS_PAGINATION_COUNT_CAP - offset)
offset_params: Final[tuple[int, ...]] = (offset,) if offset and page_limit > 0 else ()
offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else ""
page_query: Final = f"""
SELECT {_SESSION_KEY_EXPR} AS session_key,
api_key,
MAX("startTime")::text AS last_activity
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
GROUP BY {_SESSION_GROUP_KEY_SQL}
{having_clause}
ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
LIMIT ${limit_index} {offset_clause}
"""
page_rows: Final[Sequence[_SessionPageRow]] = (
()
if page_limit <= 0
else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params)
)
has_more: Final = len(page_rows) > page_limit
visible_rows: Final = page_rows[:page_limit]
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
)
page_starts_inside_the_list: Final = offset == 0 or len(page_rows) > 0
page_ends_the_list: Final = cursor is None and page_limit > 0 and not has_more and page_starts_inside_the_list
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)
)
session_keys: Final = tuple((row["session_key"], row["api_key"]) for row in visible_rows)
session_keys: Final = tuple((session.key, session.api_key) for session in session_page.sessions)
data: Final[list[dict[str, object]]] = ( # mutable-ok: _build_ui_spend_logs_response writes onto each row
await _fetch_session_representatives(
prisma_client=prisma_client,
@ -3253,18 +3154,17 @@ async def _ui_session_grouped_spend_logs(
)
_hydrate_spend_log_metadata(data)
total_pages: Final = (total_records + page_size - 1) // page_size
response: Final[Mapping[str, object]] = await _build_ui_spend_logs_response(
prisma_client,
data,
total_records,
session_page.total,
page,
page_size,
total_pages,
(session_page.total + page_size - 1) // page_size,
enrich_session_counts=True,
total_is_capped=total_is_capped,
total_is_capped=session_page.total_is_capped,
)
return {**response, "next_session_cursor": next_cursor, "has_more": has_more}
return {**response, "next_session_cursor": session_page.next_cursor, "has_more": session_page.has_more}
class RequestResponsePayload(NamedTuple):
@ -4720,16 +4620,16 @@ async def _build_ui_spend_logs_response(
)
), 0)::bigint AS session_total_duration_ms,
COUNT(*) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
WHERE call_type IN {MCP_CALL_TYPES_SQL}
)::int AS mcp_tool_call_count,
COALESCE(SUM(spend) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
WHERE call_type IN {MCP_CALL_TYPES_SQL}
), 0)::double precision AS mcp_tool_call_spend,
COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count,
COUNT(*) FILTER (
WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
WHERE call_type NOT IN {MCP_CALL_TYPES_SQL} AND call_type != {AGENT_CALL_TYPE_SQL}
)::int AS session_llm_count,
COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count,
COUNT(*) FILTER (WHERE call_type = {AGENT_CALL_TYPE_SQL})::int AS session_agent_count,
COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens,
COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens,
COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens

455
tests/integration/spend/test_session_grouped_logs_pages.py Normal file
View file

@ -0,0 +1,455 @@
import itertools
import json
import os
import uuid
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import psycopg
import pytest
from prisma import Prisma
from psycopg import sql
from pydantic import BaseModel, ConfigDict, Field
from litellm.proxy.spend_tracking.session_page import read_session_page
from litellm.proxy.spend_tracking.spend_management_endpoints import _fetch_session_representatives
from tests.integration._support.client import Gateway
from tests.integration._support.database import write_rows
_URL: Final = "/spend/logs/ui"
_WINDOW_SECONDS: Final = 3600
@dataclass(frozen=True, slots=True)
class _Row:
request: str
session: str | None
api_key: str
second: int
status: str = "success"
model: str = "m-a"
_ROWS: Final = (
_Row("z1", "z", "key-1", -1),
_Row("a1", "a", "key-1", 10),
_Row("a2", "a", "key-1", 50),
_Row("b1", "b", "key-1", 20, model="m-b"),
_Row("b2", "b", "key-1", 60),
_Row("b3", "b", "key-1", 30),
_Row("c1", "c", "key-2", 60),
_Row("d1", "b", "key-2", 40, model="m-b"),
_Row("e1", None, "key-1", 55, model="m-b"),
_Row("f1", "", "key-1", 45),
_Row("g1", "g", "key-3", 5, status="failure"),
_Row("g2", "g", "key-3", 15, status="failure"),
*(_Row(f"h{second - 69}", "h", "key-1", second) for second in range(70, 76)),
_Row("i1", "i", "key-3", 35, status="failure"),
_Row("i2", "i", "key-3", _WINDOW_SECONDS + 1),
)
_NEWEST_FIRST: Final = (
("h", "key-1"),
("c", "key-2"),
("b", "key-1"),
("e1", "key-1"),
("a", "key-1"),
("f1", "key-1"),
("b", "key-2"),
("i", "key-3"),
("g", "key-3"),
)
class _LogRow(BaseModel):
model_config = ConfigDict(frozen=True)
request_id: str
session_id: str | None
api_key: str
session_total_count: int
class _Page(BaseModel):
model_config = ConfigDict(frozen=True)
data: tuple[_LogRow, ...]
total: int
total_is_capped: bool
has_more: bool
next_session_cursor: str | None
@dataclass(frozen=True, slots=True)
class _Seeded:
prefix: str
window: Mapping[str, str]
def session(self, row: _LogRow) -> tuple[str, str]:
key: Final = row.session_id or row.request_id
assert key.startswith(self.prefix), key
return key.removeprefix(self.prefix), row.api_key
@pytest.fixture
def seeded(gateway: Gateway) -> Iterator[_Seeded]:
prefix: Final = f"sess{uuid.uuid4().hex[:12]}-"
start: Final = datetime(1900, 1, 1) + timedelta(days=int(prefix[4:16], 16) % 200000)
rows: Final = [
{
"request_id": f"{prefix}{row.request}",
"session_id": None if row.session is None else f"{prefix}{row.session}" if row.session else "",
"api_key": row.api_key,
"start_time": (start + timedelta(seconds=row.second)).isoformat(),
"status": row.status,
"model": row.model,
}
for row in _ROWS
]
with gateway.scenario() as scenario:
scenario.cleanups.callback(
write_rows, 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (f"{prefix}%",)
)
write_rows(
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, "startTime", "endTime", model, '
"session_id, status) SELECT request_id, 'acompletion', api_key, start_time, start_time, model, session_id, "
"status FROM jsonb_to_recordset(%s::jsonb) AS seeded(request_id text, session_id text, api_key text, "
"start_time timestamp, status text, model text)",
(json.dumps(rows),),
)
yield _Seeded(
prefix,
{
"start_date": f"{start:%Y-%m-%d %H:%M:%S}",
"end_date": f"{start + timedelta(seconds=_WINDOW_SECONDS):%Y-%m-%d %H:%M:%S}",
"group_by_session": "true",
},
)
def _page(gateway: Gateway, seeded: _Seeded, params: Mapping[str, str]) -> _Page:
return _Page.model_validate(gateway.get(_URL, {**seeded.window, **params}))
def _cursor_walk(gateway: Gateway, seeded: _Seeded, params: Mapping[str, str], number: int = 1) -> tuple[_Page, ...]:
page: Final = _page(gateway, seeded, {**params, "page_size": "2", "page": str(number)})
if page.next_session_cursor is None or number > len(_NEWEST_FIRST):
return (page,)
following: Final = {**params, "session_cursor": page.next_session_cursor}
return (page, *_cursor_walk(gateway, seeded, following, number + 1))
def _rows(pages: tuple[_Page, ...]) -> Iterator[_LogRow]:
return itertools.chain.from_iterable(page.data for page in pages)
def _sessions(seeded: _Seeded, pages: tuple[_Page, ...]) -> tuple[tuple[str, str], ...]:
return tuple(seeded.session(row) for row in _rows(pages))
def test_cursor_pages_show_every_session_once_newest_first_with_its_newest_row(
gateway: Gateway, seeded: _Seeded
) -> None:
pages: Final = _cursor_walk(gateway, seeded, {})
assert _sessions(seeded, pages) == _NEWEST_FIRST
assert [(page.total, page.total_is_capped) for page in pages] == [(9, False)] * 5
assert [page.has_more for page in pages] == [True, True, True, True, False]
representatives: Final = tuple(
(row.request_id.removeprefix(seeded.prefix), row.session_total_count) for row in _rows(pages)
)
assert representatives == (
("h6", 6),
("c1", 1),
("b2", 3),
("e1", 1),
("a2", 2),
("f1", 1),
("d1", 1),
("i1", 2),
("g2", 2),
)
@pytest.mark.parametrize(
"sort_order, expected",
[("desc", _NEWEST_FIRST), ("asc", tuple(reversed(_NEWEST_FIRST)))],
ids=["newest-first", "oldest-first"],
)
def test_numbered_pages_show_the_same_sessions_as_the_cursor_walk(
gateway: Gateway, seeded: _Seeded, sort_order: str, expected: tuple[tuple[str, str], ...]
) -> None:
pages: Final = tuple(
_page(gateway, seeded, {"page_size": "2", "page": str(number), "sort_order": sort_order})
for number in range(1, 7)
)
assert _sessions(seeded, pages) == expected
assert [len(page.data) for page in pages] == [2, 2, 2, 2, 1, 0]
assert [page.has_more for page in pages] == [True, True, True, True, False, False]
assert {page.total for page in pages} == {len(expected)}
def test_failure_filter_keeps_only_sessions_whose_outcome_failed(gateway: Gateway, seeded: _Seeded) -> None:
pages: Final = _cursor_walk(gateway, seeded, {"status_filter": "failure"})
assert _sessions(seeded, pages) == (("i", "key-3"), ("g", "key-3"))
assert pages[0].total == 2
def test_model_filter_ranks_each_session_by_its_matching_rows_only(gateway: Gateway, seeded: _Seeded) -> None:
pages: Final = _cursor_walk(gateway, seeded, {"model": "m-b"})
assert _sessions(seeded, pages) == (("e1", "key-1"), ("b", "key-2"), ("b", "key-1"))
assert tuple(row.request_id.removeprefix(seeded.prefix) for row in _rows(pages)) == ("e1", "d1", "b1")
assert pages[0].total == 3
_COPIED_ROWS: Final = 20_000
_COPY_START: Final = datetime(2000, 1, 1, tzinfo=UTC)
_COPY_WINDOW: Final = (
"\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC') AND \"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')"
)
_COPY_WINDOW_PARAMS: Final = (_COPY_START, _COPY_START + timedelta(days=1))
def _generated_sessions(rows_per_session: int) -> tuple[sql.SQL, tuple[object, ...]]:
"""`_COPIED_ROWS` rows a second apart in sessions of `rows_per_session` rows, a one-row session keyed by request_id."""
return sql.SQL(
'INSERT INTO {table} (request_id, call_type, api_key, "startTime", "endTime", session_id) '
"SELECT 'r' || n, 'acompletion', 'key-1', moment, moment, "
"CASE WHEN %s > 1 THEN 's' || (n - 1) / %s END FROM generate_series(1, %s) n, "
"LATERAL (SELECT %s::timestamp + n * INTERVAL '1 second' AS moment) seconds"
), (rows_per_session, rows_per_session, _COPIED_ROWS, _COPY_START.replace(tzinfo=None))
def _sessions_tied_at_the_newest_moment() -> tuple[sql.SQL, tuple[object, ...]]:
"""`_generated_sessions(1)` plus thirty one-row sessions sharing a newer moment, inserted and keyed against session
order so a read cut inside the tie misses the sessions that rank first."""
start: Final = _COPY_START.replace(tzinfo=None)
return sql.SQL(
'INSERT INTO {table} (request_id, call_type, api_key, "startTime", "endTime", session_id) '
"SELECT request_id, 'acompletion', 'key-1', moment, moment, session_id FROM ("
"SELECT n AS ordinal, 'r' || n AS request_id, %s::timestamp + n * INTERVAL '1 second' AS moment, "
"NULL AS session_id FROM generate_series(1, %s) n UNION ALL "
"SELECT %s + 30 - n, 't' || lpad((29 - n)::text, 2, '0'), %s::timestamp + %s * INTERVAL '1 second', "
"'s' || lpad(n::text, 2, '0') FROM generate_series(0, 29) n) seeded ORDER BY ordinal"
), (start, _COPIED_ROWS, _COPIED_ROWS, start, _COPIED_ROWS + 1)
@dataclass(frozen=True, slots=True)
class _Call:
session: str
call_type: str
status: str | None
second: int
def _calls(calls: tuple[_Call, ...]) -> tuple[sql.SQL, tuple[object, ...]]:
"""One row per call, `second` seconds into the window, ending a second later."""
seeded: Final = tuple(
{
"request_id": f"{call.session}-{call.second}",
"call_type": call.call_type,
"session_id": call.session,
"status": call.status,
"start_time": (_COPY_START + timedelta(seconds=call.second)).replace(tzinfo=None).isoformat(),
}
for call in calls
)
return sql.SQL(
'INSERT INTO {table} (request_id, call_type, api_key, "startTime", "endTime", session_id, status) '
"SELECT request_id, call_type, 'key-1', start_time, start_time + INTERVAL '1 second', session_id, status "
"FROM jsonb_to_recordset(%s::jsonb) AS calls(request_id text, call_type text, session_id text, status text, "
"start_time timestamp)"
), (json.dumps(seeded),)
@asynccontextmanager
async def _spend_logs_copy(insert: sql.SQL, args: tuple[object, ...]) -> AsyncIterator[Prisma]:
"""A private schema holding a copy of the migrated `LiteLLM_SpendLogs` with its indexes, seeded by `insert`."""
schema: Final = f"integration_{uuid.uuid4().hex}"
url: Final = os.environ["DATABASE_URL"]
table: Final = sql.Identifier(schema, "LiteLLM_SpendLogs")
with psycopg.connect(url, autocommit=True) as setup:
setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
try:
setup.execute(sql.SQL('CREATE TABLE {} (LIKE "LiteLLM_SpendLogs" INCLUDING ALL)').format(table))
setup.execute(insert.format(table=table), args)
setup.execute(sql.SQL("ANALYZE {}").format(table))
parsed: Final = urlsplit(url)
query: Final = urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})
database: Final = Prisma(datasource={"url": urlunsplit(parsed._replace(query=query))})
await database.connect()
try:
yield database
finally:
await database.disconnect()
finally:
setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
class _PlanNode(BaseModel):
model_config = ConfigDict(frozen=True)
relation_name: str | None = Field(default=None, alias="Relation Name")
actual_rows: float = Field(alias="Actual Rows")
rows_removed_by_filter: float = Field(default=0, alias="Rows Removed by Filter")
actual_loops: int = Field(alias="Actual Loops")
plans: tuple["_PlanNode", ...] = Field(default=(), alias="Plans")
def spend_log_rows_scanned(self) -> int:
own: Final = (
(self.actual_rows + self.rows_removed_by_filter) * self.actual_loops
if self.relation_name == "LiteLLM_SpendLogs"
else 0
)
return round(own) + sum(child.spend_log_rows_scanned() for child in self.plans)
class _Explained(BaseModel):
model_config = ConfigDict(frozen=True)
plan: _PlanNode = Field(alias="Plan")
class _ExplainedReads:
"""Runs each read through Prisma, first under EXPLAIN ANALYZE to count the `LiteLLM_SpendLogs` rows it scans."""
def __init__(self, database: Prisma) -> None:
self.database: Final = database
self.rows_scanned: tuple[int, ...] = ()
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
explained: Final = await self.database.query_raw(f"EXPLAIN (ANALYZE, FORMAT JSON) {query}", *args)
plan: Final = _Explained.model_validate(explained[0]["QUERY PLAN"][0]).plan
self.rows_scanned = (*self.rows_scanned, plan.spend_log_rows_scanned())
return await self.database.query_raw(query, *args)
@pytest.mark.parametrize(
"rows_per_session, page, newest_session, rows_scanned_per_read, total",
[
pytest.param(1, 1, ("r", 20_000), (102,), (100, True), id="first-page"),
pytest.param(1, 3, ("r", 19_980), (102,), (100, True), id="third-page"),
pytest.param(10, 1, ("s", 1_999), (102, 1_171), (100, True), id="second-read-sized-by-the-first-reads-density"),
pytest.param(100, 1, ("s", 199), (102, 6_439, 12_878), (100, True), id="third-read-at-least-doubled"),
pytest.param(1_000, 1, ("s", 19), (102, 12_878, 20_000), (20, False), id="last-read-runs-out-of-rows"),
],
)
async def test_a_newest_first_page_scans_only_the_rows_its_reads_ask_for(
rows_per_session: int,
page: int,
newest_session: tuple[str, int],
rows_scanned_per_read: tuple[int, ...],
total: tuple[int, bool],
) -> None:
async with _spend_logs_copy(*_generated_sessions(rows_per_session)) as database:
reads: Final = _ExplainedReads(database)
session_page: Final = await read_session_page(
reads,
where_clause=_COPY_WINDOW,
params=_COPY_WINDOW_PARAMS,
page=page,
page_size=10,
status=None,
sort_desc=True,
session_cursor=None,
cap=100,
)
prefix, newest = newest_session
assert tuple(session.key for session in session_page.sessions) == tuple(
f"{prefix}{newest - n}" for n in range(10)
)
assert (session_page.total, session_page.total_is_capped) == total
assert reads.rows_scanned == rows_scanned_per_read
async def test_a_page_cut_inside_a_tie_reads_on_until_the_tie_is_whole() -> None:
async with _spend_logs_copy(*_sessions_tied_at_the_newest_moment()) as database:
reads: Final = _ExplainedReads(database)
session_page: Final = await read_session_page(
reads,
where_clause=_COPY_WINDOW,
params=_COPY_WINDOW_PARAMS,
page=1,
page_size=5,
status=None,
sort_desc=True,
session_cursor=None,
cap=10,
)
assert tuple(session.key for session in session_page.sessions) == ("s29", "s28", "s27", "s26", "s25")
assert reads.rows_scanned == (12, 24, 48)
@dataclass(frozen=True, slots=True)
class _ExplainedPrisma:
db: _ExplainedReads
async def test_a_pages_representatives_are_found_by_session_not_by_scanning_the_window() -> None:
async with _spend_logs_copy(*_generated_sessions(rows_per_session=10)) as database:
reads: Final = _ExplainedReads(database)
representatives: Final = await _fetch_session_representatives(
prisma_client=_ExplainedPrisma(reads), # pyright: ignore[reportArgumentType] # only .db is used
where_clause=_COPY_WINDOW,
sql_params=_COPY_WINDOW_PARAMS,
next_param_index=len(_COPY_WINDOW_PARAMS) + 1,
session_keys=(("s1999", "key-1"), ("s1998", "key-1")),
)
assert tuple(row["request_id"] for row in representatives) == ("r20000", "r19990")
assert reads.rows_scanned == (20,), "only the two sessions' rows, not the window's 20,000"
async def test_a_page_number_far_past_the_cap_shows_no_session() -> None:
async with _spend_logs_copy(*_generated_sessions(rows_per_session=1)) as database:
session_page: Final = await read_session_page(
_ExplainedReads(database),
where_clause=_COPY_WINDOW,
params=_COPY_WINDOW_PARAMS,
page=9_300_000_000_000_000,
page_size=1000,
status=None,
sort_desc=True,
session_cursor=None,
cap=100,
)
assert (session_page.sessions, session_page.total, session_page.total_is_capped) == ((), 100, True)
_OUTCOME_CALLS: Final = (
_Call("agent", "asend_message", "failure", 10),
_Call("agent", "acompletion", "success", 20),
_Call("mcp", "acompletion", "failure", 5),
_Call("mcp", "call_mcp_tool", "success", 30),
_Call("unknown", "acompletion", None, 40),
_Call("failed", "acompletion", "success", 1),
_Call("failed", "acompletion", "failure", 50),
_Call("recovered", "acompletion", "failure", 2),
_Call("recovered", "acompletion", "success", 45),
)
@pytest.mark.parametrize(
"status, expected",
[("failure", ("failed", "mcp", "agent")), ("success", ("recovered", "unknown"))],
)
async def test_a_status_filter_keeps_sessions_whose_representative_call_has_that_status(
status: str, expected: tuple[str, ...]
) -> None:
"""An agent message stands for its session, then the newest call that is not an MCP tool call, and a missing status
counts as success. The page stays ordered by each session's newest call."""
async with _spend_logs_copy(*_calls(_OUTCOME_CALLS)) as database:
session_page: Final = await read_session_page(
_ExplainedReads(database),
where_clause=_COPY_WINDOW,
params=_COPY_WINDOW_PARAMS,
status=status,
page=1,
page_size=10,
sort_desc=True,
session_cursor=None,
cap=100,
)
assert tuple(session.key for session in session_page.sessions) == expected
assert (session_page.total, session_page.total_is_capped) == (len(expected), False)

327
tests/unit/proxy/spend_tracking/test_session_page.py Normal file
View file

@ -0,0 +1,327 @@
import itertools
import random
import re
from collections.abc import Sequence
from typing import Final, TypeAlias
import pytest
from litellm.proxy.spend_tracking.session_page import (
Session,
_parse_session_cursor,
_Read,
_ReadRow,
_Request,
_session_page,
read_session_page,
)
_Row: TypeAlias = tuple[int, str, str] # (startTime, session key, api_key)
_Ranked: TypeAlias = tuple[int, str, str] # (last activity, session key, api_key)
def _session_of(row: _Row) -> tuple[str, str]:
return row[1], row[2]
def _ranked(rows: Sequence[_Row], request: _Request) -> tuple[_Ranked, ...]:
"""The sessions in ``rows`` in page order, after the request's cursor or offset."""
by_session: Final = itertools.groupby(sorted(rows, key=_session_of), key=_session_of)
ranked: Final = tuple(
sorted(
((max(start for start, _, _ in group), key, api_key) for (key, api_key), group in by_session),
reverse=request.sort_desc,
)
)
if request.cursor is None:
return ranked[request.offset :]
cursor: Final = (int(request.cursor[0]), request.cursor[1], request.cursor[2])
return tuple(session for session in ranked if (session < cursor if request.sort_desc else session > cursor))
class _Reads:
"""Reads like the page query, whose LIMIT cuts tied rows lowest key first, the opposite of page order. It fails a
read no larger than the one before, since only growing reads end the loop."""
def __init__(self, rows: Sequence[_Row], request: _Request, every_row: bool = False) -> None:
self.newest_first: Final = tuple(sorted(rows, key=lambda row: (-row[0], row[1], row[2])))
self.request: Final = request
self.every_row: Final = every_row
self.budgets: tuple[int | None, ...] = ()
async def __call__(self, budget: int | None) -> _Read:
assert not self.budgets or (self.budgets[-1] is not None and (budget is None or budget > self.budgets[-1])), (
self.budgets,
budget,
)
self.budgets = (*self.budgets, budget)
taken: Final = self.newest_first if budget is None or self.every_row else self.newest_first[:budget]
page: Final = _ranked(taken, self.request)[: self.request.page_limit + 1] if self.request.page_limit > 0 else ()
return _Read(
candidates=tuple(Session(key, api_key, f"{at:06d}") for at, key, api_key in page),
last_settled=bool(page) and page[-1][0] > taken[-1][0],
rows_read=len(taken),
session_count=len(frozenset(_session_of(row) for row in taken)),
saw_every_row=budget is None or self.every_row or len(taken) < budget,
)
def _session_rows(rng: random.Random, index: int, clock: int) -> tuple[_Row, ...]:
start, length = rng.randrange(0, clock), rng.choice((1, 1, 2, 3, 8, 40))
return tuple((start + rng.randrange(0, 30), f"s{index:03d}", rng.choice(("k1", "k2"))) for _ in range(length))
def _sessions(rng: random.Random) -> tuple[_Row, ...]:
"""Sessions of one to forty rows over a short clock, so many rows share a time."""
clock: Final = rng.choice((3, 40, 400))
return tuple(
itertools.chain.from_iterable(_session_rows(rng, index, clock) for index in range(rng.randrange(1, 120)))
)
def _random_cursor(rng: random.Random, ranked: Sequence[_Ranked]) -> tuple[str, str, str]:
at, key, api_key = ranked[rng.randrange(len(ranked))]
return f"{at:06d}", key, api_key
def _request(rng: random.Random, rows: Sequence[_Row]) -> _Request:
cap, page_size, sort_desc = rng.randrange(2, 60), rng.randrange(1, 12), rng.random() < 0.8
mode: Final = rng.choice(("first", "offset", "cursor"))
full: Final = _ranked(rows, _Request(None, 0, 0, sort_desc, None, cap))
cursor: Final = _random_cursor(rng, full) if mode == "cursor" and full else None
offset: Final = (rng.randrange(1, 2 * cap // page_size + 3) - 1) * page_size if mode == "offset" else 0
return _Request(cursor, offset, min(page_size, cap - offset), sort_desc, None, cap)
@pytest.mark.parametrize("seed", range(60))
@pytest.mark.asyncio
async def test_a_page_read_from_the_newest_rows_matches_grouping_every_row(seed: int) -> None:
rng: Final = random.Random(seed)
rows: Final = _sessions(rng)
for request in tuple(_request(rng, rows) for _ in range(12)):
grouped = await _session_page(_Reads(rows, request, every_row=True), request)
assert await _session_page(_Reads(rows, request), request) == grouped, request
@pytest.mark.parametrize(
"rows_per_session, expected_budgets",
[
pytest.param(1, (102,), id="one-read-holds-the-cap"),
pytest.param(2, (102, 253), id="sized-by-the-first-reads-density"),
pytest.param(50, (102, 4293, 8586), id="doubles-when-the-density-misleads"),
pytest.param(500, (102, 12878, 62533), id="re-sized-by-each-reads-density"),
],
)
@pytest.mark.asyncio
async def test_reads_grow_by_the_last_reads_rows_per_session_until_one_settles_the_page(
rows_per_session: int, expected_budgets: tuple[int, ...]
) -> None:
"""150 sessions a row per second apart. A cap of 100 starts at 100 + 2 rows."""
rows: Final = tuple((index, f"s{index // rows_per_session:05d}", "k") for index in range(150 * rows_per_session))
request: Final = _Request(cursor=None, offset=0, page_limit=10, sort_desc=True, status=None, cap=100)
reads: Final = _Reads(rows, request)
await _session_page(reads, request)
assert reads.budgets == expected_budgets
@pytest.mark.asyncio
async def test_a_session_tied_with_the_reads_oldest_row_is_not_trusted() -> None:
"""Thirty one-row sessions share one time and the first read takes 12 of them. The page's sixth session sits at the
read's oldest time, where the LIMIT may have cut a session that ranks before it, so the page reads on."""
rows: Final = tuple((0, f"s{index:02d}", "k") for index in range(30))
request: Final = _Request(cursor=None, offset=0, page_limit=5, sort_desc=True, status=None, cap=10)
reads: Final = _Reads(rows, request)
page: Final = await _session_page(reads, request)
assert tuple(session.key for session in page.sessions) == ("s29", "s28", "s27", "s26", "s25")
assert reads.budgets == (12, 24, 48)
@pytest.mark.parametrize("session_count, expected_total", [(100, (100, False)), (101, (100, True))])
@pytest.mark.asyncio
async def test_the_total_is_capped_only_past_the_cap(session_count: int, expected_total: tuple[int, bool]) -> None:
rows: Final = tuple((index, f"s{index:05d}", "k") for index in range(session_count))
request: Final = _Request(cursor=None, offset=0, page_limit=10, sort_desc=True, status=None, cap=100)
page: Final = await _session_page(_Reads(rows, request), request)
assert (page.total, page.total_is_capped) == expected_total
@pytest.mark.asyncio
async def test_a_page_number_near_the_bigint_limit_reads_once() -> None:
rows: Final = tuple((index, f"s{index:05d}", "k") for index in range(500))
offset: Final = (9_300_000_000_000_000 - 1) * 1000
request: Final = _Request(
cursor=None, offset=offset, page_limit=min(1000, 100 - offset), sort_desc=True, status=None, cap=100
)
reads: Final = _Reads(rows, request)
page: Final = await _session_page(reads, request)
assert reads.budgets == (102,)
assert (page.sessions, page.total, page.total_is_capped) == ((), 100, True)
@pytest.mark.parametrize(
"sort_desc, status",
[pytest.param(False, None, id="oldest-first"), pytest.param(True, "failure", id="filtered-by-outcome")],
)
@pytest.mark.asyncio
async def test_a_page_that_needs_every_row_reads_them_once(sort_desc: bool, status: str | None) -> None:
rows: Final = tuple((index, f"s{index:05d}", "k") for index in range(500))
request: Final = _Request(cursor=None, offset=0, page_limit=10, sort_desc=sort_desc, status=status, cap=100)
reads: Final = _Reads(rows, request)
await _session_page(reads, request)
assert reads.budgets == (None,)
def test_parse_session_cursor() -> None:
assert _parse_session_cursor(None) is None
assert _parse_session_cursor("no-delimiter") is None
assert _parse_session_cursor("2026-08-29 10:00:00|sess-1") is None
assert _parse_session_cursor("|hashed-key|sess-1") is None
assert _parse_session_cursor("2026-08-29 10:00:00|hashed-key|") is None
assert _parse_session_cursor("2026-08-29 10:00:00|hashed-key|sess-1") == (
"2026-08-29 10:00:00",
"sess-1",
"hashed-key",
)
assert _parse_session_cursor("2026-08-29 10:00:00.123|hashed-key|sess|with|pipes") == (
"2026-08-29 10:00:00.123",
"sess|with|pipes",
"hashed-key",
)
def _bound(sql: str, args: tuple[object, ...], pattern: str) -> tuple[object, ...]:
"""The values bound to the ``$N`` placeholders that ``pattern``'s groups capture in ``sql``."""
found: Final = re.search(pattern, sql)
assert found, f"{pattern!r} not found in:\n{sql}"
assert all(1 <= int(index) <= len(args) for index in found.groups()), f"{found.groups()} outside $1..${len(args)}"
return tuple(args[int(index) - 1] for index in found.groups())
class _TableDb:
"""A table of ``rows`` matching rows in ``sessions`` sessions, whose reads return no page row."""
def __init__(self, rows: int, sessions: int) -> None:
self.rows: Final = rows
self.sessions: Final = sessions
self.calls: tuple[tuple[str, tuple[object, ...]], ...] = ()
async def query_raw(self, query: str, *args: object) -> tuple[_ReadRow, ...]:
self.calls = (*self.calls, (query, args))
budget: Final = re.search(r"DESC LIMIT (\d+)", query)
rows_read: Final = min(int(budget.group(1)), self.rows) if budget else self.rows
return (
_ReadRow(
session_key=None,
api_key=None,
last_activity=None,
settled=None,
rows_read=rows_read,
session_count=self.sessions,
),
)
def _page_bounds(sql: str) -> tuple[int, int]:
"""The LIMIT and OFFSET of the read's sessions."""
found: Final = re.search(r"LIMIT (\d+)(?: OFFSET (\d+))?\s*\) AS page", sql)
assert found, sql
return int(found.group(1)), int(found.group(2) or 0)
def _read_size(sql: str) -> int | None:
found: Final = re.search(r"DESC LIMIT (\d+)", sql)
return int(found.group(1)) if found else None
@pytest.mark.asyncio
async def test_a_read_that_fills_its_size_is_not_taken_for_every_row() -> None:
"""A read that took every row it asked for may have stopped short, so its 50 sessions are no total: the reads
grow until one takes fewer rows than it asked for."""
db: Final = _TableDb(rows=9_000, sessions=50)
page: Final = await read_session_page(
db,
where_clause="TRUE",
params=(),
page=1,
page_size=10,
status=None,
sort_desc=True,
session_cursor=None,
cap=100,
)
assert tuple(_read_size(sql) for sql, _ in db.calls) == (102, 258, 652, 1647, 4159, 10502)
assert (page.total, page.total_is_capped) == (50, False)
@pytest.mark.parametrize("page_number", [pytest.param(5, id="at-the-cap"), pytest.param(6, id="past-the-cap")])
@pytest.mark.asyncio
async def test_a_page_starting_at_or_past_the_cap_reads_once_and_shows_no_session(page_number: int) -> None:
db: Final = _TableDb(rows=1_000, sessions=150)
page: Final = await read_session_page(
db,
where_clause="TRUE",
params=(),
page=page_number,
page_size=25,
status=None,
sort_desc=True,
session_cursor=None,
cap=100,
)
assert len(db.calls) == 1, "the first read already holds more sessions than the cap"
assert _page_bounds(db.calls[0][0])[0] == 0
assert (page.total, page.total_is_capped) == (100, True)
@pytest.mark.asyncio
async def test_the_read_binds_the_filters_and_the_cursor_to_their_placeholders() -> None:
db: Final = _TableDb(rows=50, sessions=40)
where, params = "model = $1 AND api_key = $2", ("gpt-4o", "key-a")
for page, page_size, session_cursor in ((1, 3, "2026-10-28 09:00:00|key-a|sess|9"), (3, 1, None)):
await read_session_page(
db,
where_clause=where,
params=params,
page=page,
page_size=page_size,
status=None,
sort_desc=True,
session_cursor=session_cursor,
cap=100,
)
(cursor_sql, cursor_args), (offset_sql, _) = db.calls
for sql, args in db.calls:
assert {int(index) for index in re.findall(r"\$(\d+)", sql)} == set(range(1, len(args) + 1)), (
"every argument needs a placeholder and every placeholder an argument"
)
assert _bound(sql, args, r"model = \$(\d+) AND api_key = \$(\d+)") == params
assert _bound(
cursor_sql, cursor_args, r"\(last_activity, session_key, api_key\) < \(\$(\d+)::timestamp, \$(\d+), \$(\d+)\)"
) == ("2026-10-28 09:00:00", "sess|9", "key-a")
assert _page_bounds(cursor_sql) == (4, 0)
assert _page_bounds(offset_sql) == (2, 2)
@pytest.mark.asyncio
async def test_a_status_filtered_read_binds_the_status_between_the_filters_and_the_cursor() -> None:
db: Final = _TableDb(rows=50, sessions=40)
await read_session_page(
db,
where_clause="model = $1",
params=("gpt-4o",),
status="failure",
page=1,
page_size=3,
sort_desc=True,
session_cursor="2026-10-28 09:00:00|key-a|sess-9",
cap=100,
)
((sql, args),) = db.calls
assert {int(index) for index in re.findall(r"\$(\d+)", sql)} == set(range(1, len(args) + 1))
assert _bound(sql, args, r"model = \$(\d+)") == ("gpt-4o",)
assert _bound(sql, args, r"COALESCE\(status, 'success'\) = \$(\d+)") == ("failure",)
assert _bound(sql, args, r"\(last_activity, session_key, api_key\) < \(\$(\d+)::timestamp, \$(\d+), \$(\d+)\)") == (
"2026-10-28 09:00:00",
"sess-9",
"key-a",
)

103
tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -5,7 +5,7 @@ import hashlib
import json
import re
import sqlite3
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import timezone
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -7151,37 +7151,19 @@ def test_scoped_spend_report_range_at_max_allowed(client, monkeypatch):
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def test_parse_session_cursor():
from litellm.proxy.spend_tracking.spend_management_endpoints import (
_parse_session_cursor,
)
assert _parse_session_cursor(None) is None
assert _parse_session_cursor("no-delimiter") is None
assert _parse_session_cursor("2026-08-29 10:00:00|sess-1") is None
assert _parse_session_cursor("|hashed-key|sess-1") is None
assert _parse_session_cursor("2026-08-29 10:00:00|hashed-key|") is None
assert _parse_session_cursor("2026-08-29 10:00:00|hashed-key|sess-1") == (
"2026-08-29 10:00:00",
"sess-1",
"hashed-key",
)
assert _parse_session_cursor("2026-08-29 10:00:00.123|hashed-key|sess|with|pipes") == (
"2026-08-29 10:00:00.123",
"sess|with|pipes",
"hashed-key",
)
SESSION_GROUP_KEY_SQL = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
def _session_grouped_mock_prisma(session_page_rows, session_total, representative_rows):
def _session_grouped_mock_prisma(
session_page_rows: Sequence[Mapping[str, object]],
session_total: int,
representative_rows: Sequence[Mapping[str, object]],
) -> MagicMock:
"""Mock prisma dispatching the raw queries the keyset grouped path emits."""
async def mock_query_raw(sql_query, *params):
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": session_total}]
async def mock_query_raw(sql_query: str, *params: object) -> Sequence[Mapping[str, object]]:
if "rows_read" in sql_query:
return _session_read_rows(session_page_rows, session_total)
if "DISTINCT ON" in sql_query:
return representative_rows
if "COALESCE(SUM(spend)" in sql_query:
@ -7199,7 +7181,7 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ
"session_models": ["gpt-4o"],
}
]
return session_page_rows
raise AssertionError(f"unexpected query:\n{sql_query}")
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
@ -7228,23 +7210,33 @@ 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, counted_total=None):
"""Mock prisma serving the grouped page query out of ``sessions``, honoring the LIMIT and OFFSET it asks for."""
def _session_read_rows(
page_rows: Sequence[Mapping[str, object]], session_count: int
) -> tuple[Mapping[str, object], ...]:
"""The read query's rows for ``page_rows``, as if the read held every matching row."""
summary: Final = {"rows_read": 0, "session_count": session_count}
if not page_rows:
return ({"session_key": None, "api_key": None, "last_activity": None, "settled": None, **summary},)
return tuple({**row, "settled": True, **summary} for row in page_rows)
async def mock_query_raw(sql_query, *params):
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}]
def _session_grouped_paginating_prisma(sessions: Sequence[tuple[str, str]]) -> MagicMock:
"""Mock prisma serving the grouped page read out of ``sessions``, honoring the LIMIT and OFFSET it asks for."""
async def mock_query_raw(sql_query: str, *params: object) -> Sequence[Mapping[str, object]]:
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:
return []
bounds = re.search(r"LIMIT \$(\d+)(?: OFFSET \$(\d+))?", sql_query)
limit = params[int(bounds.group(1)) - 1]
offset = params[int(bounds.group(2)) - 1] if bounds.group(2) else 0
return [
_session_page_row(session_key, last_activity)
for session_key, last_activity in sessions[offset : offset + limit]
]
bounds: Final = re.search(r"LIMIT (\d+)(?: OFFSET (\d+))?\s+\) AS page", sql_query)
limit, offset = int(bounds.group(1)), int(bounds.group(2) or 0)
return _session_read_rows(
[
_session_page_row(session_key, last_activity)
for session_key, last_activity in sessions[offset : offset + limit]
],
len(sessions),
)
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
@ -7296,20 +7288,15 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
assert data["data"][1]["session_total_count"] == 1
emitted = [call.args for call in mock_prisma.db.query_raw.await_args_list]
assert len(emitted) == 3, "one read picks the page and counts it, then the representatives and their totals"
page_query_sql = emitted[0][0]
assert f"GROUP BY {SESSION_GROUP_KEY_SQL}" in page_query_sql
assert "OFFSET" not in page_query_sql
assert "HAVING" not in page_query_sql
assert emitted[0][-1] == 3, "page query fetches page_size + 1 sessions to detect has_more"
assert "LIMIT 3 " in page_query_sql, "the read fetches page_size + 1 sessions to detect has_more"
count_sql = emitted[1][0]
assert f"GROUP BY {SESSION_GROUP_KEY_SQL}" in count_sql
assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
rep_call = emitted[2]
rep_call: Final = emitted[1]
assert f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in rep_call[0]
assert (
f"ORDER BY {SESSION_GROUP_KEY_SQL}, " + spend_management_endpoints._SESSION_REPRESENTATIVE_ORDER_SQL
f"ORDER BY {SESSION_GROUP_KEY_SQL}, " + spend_management_endpoints.SESSION_REPRESENTATIVE_ORDER_SQL
) in rep_call[0]
assert rep_call[-2] == ["sess-1", "req-solo"]
assert rep_call[-1] == ["hashed-key", "hashed-key"]
@ -7319,7 +7306,7 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
@pytest.mark.asyncio
async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypatch):
"""A session_cursor becomes a HAVING keyset predicate instead of an OFFSET."""
"""A session_cursor becomes a keyset predicate on the read's sessions instead of an OFFSET."""
page_rows = [_session_page_row("sess-2", "2026-08-29 07:00:00")]
reps = [_session_representative_row("req-2", "sess-2")]
mock_prisma = _session_grouped_mock_prisma(page_rows, 3, reps)
@ -7352,7 +7339,7 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat
page_query_call = mock_prisma.db.query_raw.await_args_list[0]
page_query_sql = page_query_call.args[0]
assert f'HAVING (MAX("startTime"), {SESSION_GROUP_KEY_SQL}) <' in page_query_sql
assert "WHERE (last_activity, session_key, api_key) <" in page_query_sql
assert "OFFSET" not in page_query_sql
cursor_index = page_query_call.args.index("2026-08-29 09:00:00")
assert page_query_call.args[cursor_index : cursor_index + 3] == (
@ -7403,9 +7390,9 @@ async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(
@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."""
"""A page that runs out of sessions has read every matching row, so that one read gives the total."""
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)
mock_prisma: Final = _session_grouped_paginating_prisma(sessions)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr(
"litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe",
@ -7429,10 +7416,12 @@ async def test_ui_view_spend_logs_group_by_session_short_page_totals_itself(clie
)
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"] == 10
assert data["total_is_capped"] is False
assert data["total_pages"] == 1
assert len(data["data"]) == 10
reads: Final = tuple(call for call in mock_prisma.db.query_raw.await_args_list if "rows_read" in call.args[0])
assert len(reads) == 1, "a page that ends the list is not read again for its total"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@ -7517,7 +7506,7 @@ async def test_ui_view_spend_logs_group_by_session_page_past_count_cap_is_empty(
assert past_cap_data["data"] == []
assert past_cap_data["has_more"] is False
assert past_cap_data["total"] == cap
assert mock_prisma.db.query_raw.await_count == 1, "only the bounded count query runs past the capped window"
assert mock_prisma.db.query_raw.await_count == 1, "only the bounded read runs past the capped window"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@ -7820,14 +7809,14 @@ def test_session_representative_uses_completed_agent_outcome(parent_status, chil
)
result = connection.execute(
"SELECT request_id, status FROM logs ORDER BY "
+ spend_management_endpoints._SESSION_REPRESENTATIVE_ORDER_SQL
+ spend_management_endpoints.SESSION_REPRESENTATIVE_ORDER_SQL
+ " LIMIT 1"
).fetchone()
assert result == ("parent", expected)
connection.execute("DELETE FROM logs WHERE call_type = 'asend_message'")
fallback = connection.execute(
"SELECT request_id, status FROM logs ORDER BY "
+ spend_management_endpoints._SESSION_REPRESENTATIVE_ORDER_SQL
+ spend_management_endpoints.SESSION_REPRESENTATIVE_ORDER_SQL
+ " LIMIT 1"
).fetchone()
assert fallback == ("llm", "success")

40
tests/unit/proxy/spend_tracking/test_spend_query_optimization.py
View file

@ -6,7 +6,9 @@ GitHub Issue: #17487
"""
import datetime
from collections.abc import Mapping, Sequence
from datetime import timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -508,20 +510,10 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch):
@pytest.mark.asyncio
async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
"""
With group_by_session=true, /spend/logs/ui must page and count SESSIONS,
not raw calls: the page is one representative row per session (DISTINCT
ON the session group key, preferring non-MCP calls, newest first), the
bounded count counts groups, and on the default startTime sort the page is
selected by keyset (no OFFSET) so deep pages do not degrade. Otherwise the
UI collapses a server page of N calls into fewer visible rows while the
footer still claims N (issue #38060).
"""
"""With group_by_session=true, /spend/logs/ui pages and counts sessions, not calls, one representative row per
session, so the footer's total matches the rows shown (issue #38060)."""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.spend_tracking.spend_management_endpoints import (
SPEND_LOGS_PAGINATION_COUNT_CAP,
ui_view_spend_logs,
)
from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs
group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
session_rows = [
@ -533,12 +525,12 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
{"request_id": "req-2", "api_key": "k", "metadata": "{}", "session_id": None},
]
async def mock_query_raw(sql_query, *params):
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": 60}]
async def mock_query_raw(sql_query: str, *params: object) -> Sequence[Mapping[str, object]]:
if "rows_read" in sql_query:
return [{**row, "settled": True, "rows_read": 0, "session_count": 60} for row in session_rows]
if "AS session_representatives" in sql_query:
return representative_rows
return session_rows
return []
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
@ -565,22 +557,16 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
user_api_key_dict=auth,
group_by_session=True,
session_cursor=None,
status_filter=None,
)
emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list]
page_sql = emitted[0][0]
assert f"GROUP BY {group_key}" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}"
assert "GROUP BY session_key, api_key" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}"
assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected"
assert emitted[0][-1] == 51, "the page query fetches page_size + 1 sessions to detect has_more"
assert "LIMIT 51 " in page_sql, "the page read fetches page_size + 1 sessions to detect has_more"
count_call = emitted[1]
count_sql = count_call[0]
assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}"
assert "COUNT(*) OVER ()" not in count_sql
assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
assert count_call[-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1
rep_sql = emitted[2][0]
rep_sql: Final = emitted[1][0]
assert f"DISTINCT ON ({group_key})" in rep_sql, f"page must return one row per session. SQL was:\n{rep_sql}"
assert (
f"ORDER BY {group_key}, (call_type = 'asend_message') DESC, "