mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix: paginate customer team activity in SQL
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kenneth
parent
d9823f8610
commit
11cf10e360
4 changed files with 152 additions and 18 deletions
|
|
@ -1,11 +1,12 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import StrictStr, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm import constants
|
||||
|
|
@ -18,7 +19,7 @@ from litellm.repositories.daily_activity_sql import (
|
|||
adjust_dates_for_timezone,
|
||||
build_aggregated_sql,
|
||||
build_cache_leakage_keys_sql,
|
||||
build_customer_filtered_keys_sql,
|
||||
build_customer_daily_rows_sql,
|
||||
build_customer_team_rollup_sql,
|
||||
build_entity_rollup_sql,
|
||||
build_export_sql,
|
||||
|
|
@ -106,7 +107,6 @@ class DailyActivityDatabase(Protocol):
|
|||
|
||||
_GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...])
|
||||
_CUSTOMER_TEAM_ADAPTER: Final = TypeAdapter(tuple[CustomerTeamRollupRow, ...])
|
||||
_FILTERED_KEYS_ADAPTER: Final = TypeAdapter(tuple[dict[str, StrictStr], ...])
|
||||
_ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...])
|
||||
_GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...])
|
||||
_KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...])
|
||||
|
|
@ -116,6 +116,44 @@ _EXPORT_ADAPTER: Final = TypeAdapter(tuple[ExportRow, ...])
|
|||
_METADATA_TAGS_ADAPTER: Final = TypeAdapter(list[StrictStr])
|
||||
|
||||
|
||||
class _CustomerDailyRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
date: str
|
||||
api_key: str
|
||||
end_user_id: str | None
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
compression_saved_tokens: int
|
||||
compression_savings_spend: float
|
||||
prompt_caching_savings_spend: float
|
||||
gateway_injected_caching_savings_spend: float
|
||||
autorouter_savings_spend: float
|
||||
spend: float
|
||||
api_requests: int
|
||||
successful_requests: int
|
||||
failed_requests: int
|
||||
total_response_time_ms: int
|
||||
timed_requests: int
|
||||
model: str | None = None
|
||||
model_group: str | None = None
|
||||
custom_llm_provider: str | None = None
|
||||
mcp_namespaced_tool_name: str | None = None
|
||||
endpoint: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CustomerDailyRowsResult:
|
||||
total_count: int
|
||||
rows: tuple[_CustomerDailyRow, ...]
|
||||
|
||||
|
||||
_CUSTOMER_DAILY_ROWS_ADAPTER: Final = TypeAdapter(tuple[_CustomerDailyRowsResult])
|
||||
|
||||
|
||||
def _metadata_tags(value: object) -> tuple[str, ...]:
|
||||
stable_value: Final = value
|
||||
if not isinstance(value, list):
|
||||
|
|
@ -329,6 +367,11 @@ class DailyActivityRepository:
|
|||
)
|
||||
|
||||
async def daily_rows(self, scope: DailyActivityScope, *, page: int, page_size: int) -> DailyRowsPage:
|
||||
if scope.table is DailyActivityTable.CUSTOMER and scope.team_ids is not None:
|
||||
result: Final = _CUSTOMER_DAILY_ROWS_ADAPTER.validate_python(
|
||||
await self._query(build_customer_daily_rows_sql(scope, page=page, page_size=page_size))
|
||||
)[0]
|
||||
return DailyRowsPage(total_count=result.total_count, rows=result.rows)
|
||||
table: Final = _daily_rows_table(self._prisma_client, scope.table)
|
||||
adjusted_start, adjusted_end = adjust_dates_for_timezone(
|
||||
scope.start_date,
|
||||
|
|
@ -346,16 +389,6 @@ class DailyActivityRepository:
|
|||
if scope.exclude_entity_ids
|
||||
else {}
|
||||
)
|
||||
filtered_keys: Final = (
|
||||
tuple(
|
||||
row["api_key"]
|
||||
for row in _FILTERED_KEYS_ADAPTER.validate_python(
|
||||
await self._query(build_customer_filtered_keys_sql(scope))
|
||||
)
|
||||
)
|
||||
if scope.table is DailyActivityTable.CUSTOMER and scope.team_ids is not None
|
||||
else scope.api_keys
|
||||
)
|
||||
conditions: Final = {
|
||||
"date": {"gte": adjusted_start, "lte": adjusted_end},
|
||||
**({scope.entity_id_field: {"in": list(scope.entity_ids)}} if scope.entity_ids is not None else {}),
|
||||
|
|
@ -367,7 +400,7 @@ class DailyActivityRepository:
|
|||
if scope.model
|
||||
else {}
|
||||
),
|
||||
**({"api_key": {"in": list(filtered_keys)}} if filtered_keys is not None else {}),
|
||||
**({"api_key": {"in": list(scope.api_keys)}} if scope.api_keys is not None else {}),
|
||||
}
|
||||
count, rows = await asyncio.gather(
|
||||
table.count(where=conditions),
|
||||
|
|
|
|||
|
|
@ -189,11 +189,18 @@ def build_customer_team_rollup_sql(scope: DailyActivityScope, *, row_ids: tuple[
|
|||
)
|
||||
|
||||
|
||||
def build_customer_filtered_keys_sql(scope: DailyActivityScope) -> SqlQuery:
|
||||
def build_customer_daily_rows_sql(scope: DailyActivityScope, *, page: int, page_size: int) -> SqlQuery:
|
||||
where_clause, params = build_where_clause(scope)
|
||||
return SqlQuery(
|
||||
sql=f'SELECT DISTINCT api_key FROM "{PRISMA_TO_PG_TABLE[scope.table]}" WHERE {where_clause}',
|
||||
params=params,
|
||||
sql=f'''WITH filtered AS MATERIALIZED (
|
||||
SELECT * FROM "{PRISMA_TO_PG_TABLE[scope.table]}" WHERE {where_clause}
|
||||
)
|
||||
SELECT (SELECT COUNT(*) FROM filtered)::bigint AS total_count,
|
||||
COALESCE((SELECT jsonb_agg(page_rows) FROM (
|
||||
SELECT * FROM filtered ORDER BY date DESC, id ASC
|
||||
OFFSET ${len(params) + 1} LIMIT ${len(params) + 2}
|
||||
) page_rows), '[]'::jsonb) AS rows''',
|
||||
params=(*params, (page - 1) * page_size, page_size),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import uuid
|
||||
from hashlib import sha256
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from time import perf_counter
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -119,3 +120,66 @@ def test_customer_spend_by_current_team_and_team_filter(gateway: Gateway, route:
|
|||
legacy_teams: Final = legacy_body.results[0].breakdown.teams or {}
|
||||
assert set(legacy_teams) == {team}
|
||||
assert legacy_teams[team].metrics.spend == legacy_body.metadata.total_spend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ("/customer/daily/activity", "/customer/daily/activity/aggregated"))
|
||||
def test_large_customer_team_reports_keep_filters_and_pagination(gateway: Gateway, route: str) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
team: Final = scenario.team(team_alias="large-customer-team")
|
||||
customer: Final = f"large-customer-{uuid.uuid4().hex}"
|
||||
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_VerificationToken" WHERE team_id=%s', (team,))
|
||||
scenario.cleanups.callback(
|
||||
write_rows, 'DELETE FROM "LiteLLM_DailyEndUserSpend" WHERE end_user_id=%s', (customer,)
|
||||
)
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_VerificationToken" (token,models,team_id) '
|
||||
"SELECT encode(sha256(convert_to(%s||i,'UTF8')),'hex'),ARRAY[]::text[],%s "
|
||||
"FROM generate_series(1,33000) i",
|
||||
(customer, team),
|
||||
)
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_DailyEndUserSpend" '
|
||||
"(id,end_user_id,date,api_key,model,spend,api_requests,updated_at) "
|
||||
"SELECT %s||lpad(i::text,5,'0'),%s,'2026-02-03',"
|
||||
"encode(sha256(convert_to(%s||i,'UTF8')),'hex'),"
|
||||
"CASE WHEN i%%2=1 THEN 'large-model-a' ELSE 'large-model-b' END,"
|
||||
"CASE WHEN i%%2=1 THEN 0.25 ELSE 0.5 END,1,now() FROM generate_series(1,33000) i",
|
||||
(customer, customer, customer),
|
||||
)
|
||||
write_rows('ANALYZE "LiteLLM_VerificationToken"', ())
|
||||
write_rows('ANALYZE "LiteLLM_DeletedVerificationToken"', ())
|
||||
write_rows('ANALYZE "LiteLLM_DailyEndUserSpend"', ())
|
||||
key: Final = sha256(f"{customer}1".encode()).hexdigest()
|
||||
base: Final = (
|
||||
f"{route}?end_user_ids={customer}&team_ids={team}&start_date=2026-02-03&end_date=2026-02-03&page_size=10"
|
||||
)
|
||||
aggregated: Final = route.endswith("aggregated")
|
||||
for suffix, spend, pages in (
|
||||
("", 12375.0 if aggregated else 3.75, 3300),
|
||||
("&model=large-model-a", 4125.0 if aggregated else 2.5, 1650),
|
||||
(f"&api_key={key}", 0.25, 1),
|
||||
(f"&api_key={key}&model=large-model-a", 0.25, 1),
|
||||
(f"&api_key={key}&model=large-model-b", 0.0, 0),
|
||||
):
|
||||
started: Final = perf_counter()
|
||||
response: Final = gateway.request("GET", base + suffix)
|
||||
print(f"large-team {route} {suffix!r}: {response.status_code}, {perf_counter() - started:.3f}s")
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(response.content)
|
||||
assert body.metadata.total_spend == spend
|
||||
assert (
|
||||
sum(
|
||||
sum(bucket.metrics.spend for bucket in (day.breakdown.teams or {}).values()) for day in body.results
|
||||
)
|
||||
== spend
|
||||
)
|
||||
if not aggregated:
|
||||
assert body.metadata.total_pages == pages
|
||||
assert body.metadata.has_more == (pages > 1)
|
||||
if not aggregated:
|
||||
second: Final = gateway.request("GET", base + "&page=2")
|
||||
assert second.status_code == 200, second.text
|
||||
second_body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(second.content)
|
||||
assert second_body.metadata.total_spend == 3.75
|
||||
expected_keys: Final = {sha256(f"{customer}{i}".encode()).hexdigest() for i in range(11, 21)}
|
||||
assert set(second_body.results[0].breakdown.api_keys) == expected_keys
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -662,6 +662,36 @@ async def test_daily_rows_selects_the_table_and_applies_filters_and_pagination(
|
|||
assert sum(len(daily_table.find_many_calls) for daily_table in tables.values()) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("empty", (False, True))
|
||||
async def test_customer_team_page_returns_loaded_metrics_and_retains_count_on_empty_pages(empty: bool) -> None:
|
||||
stored: Final = {
|
||||
**_key_spend_row("key-a"),
|
||||
"id": "row-a",
|
||||
"end_user_id": "customer-a",
|
||||
"date": "2026-01-01",
|
||||
"compression_saved_tokens": 2,
|
||||
"compression_savings_spend": 0.5,
|
||||
"prompt_caching_savings_spend": 0.25,
|
||||
"gateway_injected_caching_savings_spend": 0.125,
|
||||
"autorouter_savings_spend": 0.0,
|
||||
"total_response_time_ms": 10,
|
||||
"timed_requests": 1,
|
||||
"created_at": "2026-01-01T00:00:00Z",
|
||||
"updated_at": "2026-01-01T00:00:00Z",
|
||||
}
|
||||
database: Final = _FakeDatabase((({"total_count": 33000, "rows": [] if empty else [stored]},),))
|
||||
repository, _ = _repository(database)
|
||||
scope: Final = replace(_scope(table=DailyActivityTable.CUSTOMER), team_ids=("team-a",))
|
||||
|
||||
result: Final = await repository.daily_rows(scope, page=3301 if empty else 1, page_size=10)
|
||||
|
||||
assert result.total_count == 33000
|
||||
assert tuple((row.id, row.api_key, row.spend, row.prompt_tokens) for row in result.rows) == (
|
||||
() if empty else (("row-a", "key-a", 1.0, 10),)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_rows_exclusion_without_entity_filter_keeps_null_entity_rows() -> None:
|
||||
database = _FakeDatabase()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue