scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

perf(spend): keep the per-session lookups off the window scan under generic plans

Browse source 
After five executions Postgres switches the representative and cursor
lookups to a generic plan, which ANDed each session's rows with a scan of
the whole date window on the startTime index. With 50 sessions of 150 rows
that took about 500 ms instead of about 10 ms. Find the session's rows by
its key first behind OFFSET 0, then apply the window and filters to them.
This commit is contained in:
shrey-berri 2026-09-26 20:18:19 +00:00
parent 9adfe67c34
commit 8729832b8a
No known key found for this signature in database
1 changed files with 13 additions and 7 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

@ -2974,6 +2974,16 @@ def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] |
return (last_activity, session_key, api_key)
def _session_rows_sql(session_key: str, api_key: str) -> str:
"""Every row that can belong to one session, found by its key alone. The OFFSET 0 keeps the caller's window
and filters out of this scan, so a generic plan cannot combine it with a scan of the whole window."""
return f"""SELECT *
FROM "LiteLLM_SpendLogs"
WHERE (request_id = {session_key} OR session_id = {session_key})
AND api_key = {api_key}
OFFSET 0"""
async def _fetch_session_representatives(
prisma_client: "PrismaClient",
where_clause: str,
@ -2987,10 +2997,8 @@ async def _fetch_session_representatives(
FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[]) AS requested(session_key, key)
CROSS JOIN LATERAL (
SELECT {_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs"
WHERE {where_clause}
AND (request_id = requested.session_key OR session_id = requested.session_key)
AND api_key = requested.key
AND {_SESSION_KEY_EXPR} = requested.session_key
ORDER BY call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
LIMIT 1
@ -3116,11 +3124,9 @@ class _SessionHeadWalk:
FROM unnest(${len(params) + 1}::text[], ${len(params) + 2}::text[]) AS requested(session_key, key)
CROSS JOIN LATERAL (
SELECT 1
FROM "LiteLLM_SpendLogs"
FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs"
WHERE {self.where_clause}
AND api_key = requested.key
AND (session_id = requested.session_key
OR (request_id = requested.session_key AND NULLIF(session_id, '') IS NULL))
AND {_SESSION_KEY_EXPR} = requested.session_key
AND ("startTime", requested.session_key, requested.key) >= {self._cursor_bounds()}
LIMIT 1
) AS newer_row