scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 43

fix(proxy): type the concurrent flat-page count as a task and unpin SQL emission order in the offset test

Browse source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
mubashir1osmani 2026-09-19 14:02:16 -04:00
parent a7325e6336
commit 3001970e36
2 changed files with 8 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3,7 +3,7 @@ import asyncio
import collections
import json
import os
from collections.abc import Awaitable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from itertools import groupby
@ -2880,11 +2880,11 @@ async def ui_view_spend_logs(
LIMIT ${p} OFFSET ${p + 1}
"""
)
count_task: Final[Awaitable[Sequence[_SpendLogsCountRow] | None]] = _query_raw_or_none(
prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
count_task: Final[asyncio.Task[Sequence[_SpendLogsCountRow] | None]] = asyncio.create_task(
_query_raw_or_none(prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1)
)
data_task: Final = prisma_client.db.query_raw(sql_query, *sql_params, page_size, skip)
count_rows, data = await asyncio.gather(count_task, data_task)
data: Final = await prisma_client.db.query_raw(sql_query, *sql_params, page_size, skip)
count_rows: Final = await count_task
raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0
total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total

5
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -7717,8 +7717,9 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort
assert "next_session_cursor" not in data
emitted_sql = [call.args[0] for call in mock_prisma.db.query_raw.await_args_list]
assert "HAVING" not in " ".join(emitted_sql)
assert f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in emitted_sql[1]
assert "OFFSET" in emitted_sql[1]
page_sql = [sql for sql in emitted_sql if f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in sql]
assert len(page_sql) == 1
assert "OFFSET" in page_sql[0]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)