From f3e2af348603e2ed6ac646eddbbeb9892044f485 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 23 Oct 2026 15:38:15 -0700 Subject: [PATCH] perf(spend): bound session enrichment reads on cold pages --- .../spend_management_endpoints.py | 20 +++++++++++-------- .../test_spend_management_endpoints.py | 8 ++++---- .../test_spend_query_optimization.py | 2 +- 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 9f6b034804c..e1a41c7e2ab 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3266,8 +3266,8 @@ class _SessionHeadWalk: self.prisma_client, query, *params, - [row["session_key"] for row in candidates], # mutable-ok: prisma serializes array params from a list - [row["api_key"] for row in candidates], # mutable-ok: prisma serializes array params from a list + [row["session_key"] for row in candidates], + [row["api_key"] for row in candidates], ) return frozenset((row["session_key"], row["api_key"]) for row in shown) @@ -3321,9 +3321,9 @@ class _SessionHeadWalk: FROM unnest($1::text[], $2::text[], $3::text[]) AS heads(session_key, api_key, last_activity) ORDER BY last_activity::timestamp DESC, session_key DESC, api_key DESC LIMIT $4""", - [row["session_key"] for row in all_heads], # mutable-ok: prisma serializes array params from a list - [row["api_key"] for row in all_heads], # mutable-ok: prisma serializes array params from a list - [row["last_activity"] for row in all_heads], # mutable-ok: prisma serializes array params from a list + [row["session_key"] for row in all_heads], + [row["api_key"] for row in all_heads], + [row["last_activity"] for row in all_heads], self.page_limit + 1, ) return tuple(page_rows) @@ -5031,9 +5031,13 @@ async def _build_ui_spend_logs_response( 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 - FROM "LiteLLM_SpendLogs" - WHERE session_id = ANY($1::text[]) - AND api_key = ANY($2::text[]) + FROM unnest($1::text[]) AS requested(session_key) + CROSS JOIN LATERAL ( + SELECT * FROM "LiteLLM_SpendLogs" + WHERE session_id = requested.session_key + AND api_key = ANY($2::text[]) + OFFSET 0 + ) AS session_rows GROUP BY session_id, api_key ) s LEFT JOIN LATERAL ( 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 c0893ee95f9..cb79fcfbcb8 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7212,7 +7212,7 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ return _session_count_walk(session_total) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": session_total}] - if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: + if "AS rep" in sql_query or "DISTINCT ON" in sql_query: return representative_rows if "COALESCE(SUM(spend)" in sql_query: return [ @@ -7270,7 +7270,7 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): return _session_count_walk(len(sessions) if counted_total is None else counted_total) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}] - if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: + if "AS rep" in sql_query or "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 [] @@ -7400,7 +7400,7 @@ def _session_walk_mock_prisma(walked_rows, window_rows_per_batch, grouped_rows): return _session_count_walk(3) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 3}] - if "CROSS JOIN LATERAL" in sql_query: + if "AS rep" 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 [] @@ -7488,7 +7488,7 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam return _session_count_walk(3) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 3}] - if "CROSS JOIN LATERAL" in sql_query: + if "AS rep" in sql_query: return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] return [] 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 72c882bcc74..5d6ae37f527 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -539,7 +539,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): return [{"rows_read": 0, "sessions": 60}] if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 60}] - if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: + if "AS rep" in sql_query or "DISTINCT ON" in sql_query: return representative_rows return session_rows