scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 416

perf(spend): bound session enrichment reads on cold pages

Browse source
This commit is contained in:
shrey kharbanda 2026-10-23 15:38:15 -07:00
parent d0c14bdda4
commit f3e2af3486
3 changed files with 17 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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 (

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

@ -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 []

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

@ -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