diff --git a/litellm/proxy/spend_tracking/session_page.py b/litellm/proxy/spend_tracking/session_page.py new file mode 100644 index 00000000000..810afbfa7b7 --- /dev/null +++ b/litellm/proxy/spend_tracking/session_page.py @@ -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 diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 89200792c0b..fcd719b24b9 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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 - ``'||'`` 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 diff --git a/tests/integration/spend/test_session_grouped_logs_pages.py b/tests/integration/spend/test_session_grouped_logs_pages.py new file mode 100644 index 00000000000..19854d60bd3 --- /dev/null +++ b/tests/integration/spend/test_session_grouped_logs_pages.py @@ -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) diff --git a/tests/unit/proxy/spend_tracking/test_session_page.py b/tests/unit/proxy/spend_tracking/test_session_page.py new file mode 100644 index 00000000000..a86990bddbb --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_session_page.py @@ -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", + ) 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 6598b94e60f..f8d0ebb0f77 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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") diff --git a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py index 93fae093340..7a38c1d8457 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -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, "