mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(spend): preserve database ordering across tied sessions
This commit is contained in:
parent
44792295f0
commit
b7755a48ac
3 changed files with 57 additions and 1 deletions
|
|
@ -3276,6 +3276,19 @@ class _SessionHeadWalk:
|
|||
rows_read=total_read,
|
||||
size=next_size,
|
||||
)
|
||||
if len(frozenset(row["sort_time"] for row in all_heads)) < len(all_heads):
|
||||
page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw(
|
||||
self.prisma_client,
|
||||
"""SELECT session_key, api_key, last_activity
|
||||
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
|
||||
self.page_limit + 1,
|
||||
)
|
||||
return tuple(page_rows)
|
||||
in_page_order: Final = sorted(
|
||||
all_heads, key=lambda row: (row["sort_time"], row["session_key"], row["api_key"]), reverse=True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from typing import Final
|
|||
import psycopg
|
||||
import pytest
|
||||
from integration._support.client import Gateway, string_value
|
||||
from integration._support.database import ROWS
|
||||
from integration._support.database import ROWS, read_rows
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
|
|
@ -469,3 +469,44 @@ def test_session_view_filtered_page_only_lists_matching_sessions(gateway: Gatewa
|
|||
(f"{small_seed}-shared-b", f"{small_seed}-key-b"),
|
||||
]
|
||||
assert page["total"] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("identities", ((("a", "K"), ("B", "K")), (("shared", "a"), ("shared", "B"))))
|
||||
def test_session_view_cursor_pages_preserve_database_order_for_tied_sessions(
|
||||
gateway: Gateway, identities: tuple[tuple[str, str], tuple[str, str]]
|
||||
) -> None:
|
||||
marker: Final = f"intg-sesscollation-{uuid.uuid4().hex}"
|
||||
rows: Final = tuple(
|
||||
_SeedRow(f"{marker}-{index}", f"{marker}-{session}", SMALL_WINDOW_START, api_key=key)
|
||||
for index, (session, key) in enumerate(identities)
|
||||
)
|
||||
|
||||
with gateway.scenario() as scenario:
|
||||
scenario.cleanups.callback(_delete_small_rows, marker)
|
||||
_insert_rows(rows)
|
||||
expected: Final = tuple(
|
||||
(string_value(row["request_id"]), string_value(row["api_key"]))
|
||||
for row in read_rows(
|
||||
'SELECT request_id, api_key FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s '
|
||||
'ORDER BY "startTime" DESC, session_id DESC, api_key DESC',
|
||||
(f"{marker}-%",),
|
||||
)
|
||||
)
|
||||
|
||||
first: Final = _small_page(gateway, session_id=marker, page_size="1", page="1")
|
||||
second: Final = _small_page(
|
||||
gateway,
|
||||
session_id=marker,
|
||||
page_size="1",
|
||||
page="2",
|
||||
session_cursor=string_value(first["next_session_cursor"]),
|
||||
)
|
||||
walked: Final = tuple(
|
||||
(string_value(row["request_id"]), string_value(row["api_key"]))
|
||||
for row in _data_rows(first) + _data_rows(second)
|
||||
)
|
||||
|
||||
assert walked == expected, "cursor pages must match database ordering without omissions or duplicates"
|
||||
assert first["total"] == second["total"] == len(rows)
|
||||
assert first["has_more"] is True
|
||||
assert second["has_more"] is False
|
||||
|
|
|
|||
|
|
@ -7303,6 +7303,8 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam
|
|||
tied = "2026-08-29 10:00:00"
|
||||
|
||||
async def mock_query_raw(sql_query, *params):
|
||||
if "AS heads(session_key, api_key, last_activity)" in sql_query:
|
||||
return [_session_page_row(key, tied) for key in ("c", "b")]
|
||||
if "WITH batch AS MATERIALIZED" in sql_query:
|
||||
if '("startTime", request_id) <' not in sql_query:
|
||||
return [
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue