mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
perf(spend): narrow user and team spend log reads to the owner's api_keys so the (api_key, startTime) index is used (#51026)
* perf(spend): narrow user and team spend log reads to the owner's api_keys so the (api_key, startTime) index is used Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend): satisfy the type discipline gate for the api_key scope helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend): run the api_key lookup under the statement timeout, skip it for search and empty key lists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
bc3a69ce35
commit
05c04d6261
5 changed files with 396 additions and 15 deletions
|
|
@ -18,7 +18,7 @@ import traceback
|
|||
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
|
|
@ -228,6 +228,7 @@ from litellm.proxy.management_helpers.utils import (
|
|||
management_endpoint_wrapper,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
|
||||
from litellm.proxy.spend_tracking.spend_log_api_key_scope import spend_log_api_keys_for_teams
|
||||
from litellm.proxy.utils import (
|
||||
PrismaClient,
|
||||
ProxyLogging,
|
||||
|
|
@ -7463,9 +7464,11 @@ def aggregated_date_range_error(start_date: str | None, end_date: str | None) ->
|
|||
return None
|
||||
|
||||
|
||||
def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str:
|
||||
def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool, narrow_to_api_keys: bool) -> str:
|
||||
team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count))
|
||||
user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else ""
|
||||
api_key_placeholder: Final = 3 + team_count + (1 if restrict_to_user else 0)
|
||||
api_key_clause: Final = f" AND sl.api_key = ANY(${api_key_placeholder}::text[])" if narrow_to_api_keys else ""
|
||||
return f"""
|
||||
SELECT
|
||||
sl.team_id,
|
||||
|
|
@ -7483,7 +7486,7 @@ def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str:
|
|||
LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = sl."user"
|
||||
WHERE sl."startTime" >= $1::timestamp
|
||||
AND sl."startTime" < $2::timestamp + INTERVAL '1 day'
|
||||
AND sl.team_id IN ({team_placeholders}){user_clause}
|
||||
AND sl.team_id IN ({team_placeholders}){user_clause}{api_key_clause}
|
||||
GROUP BY sl.team_id, sl."user", u.user_email, u.user_alias
|
||||
ORDER BY spend DESC, sl.team_id, sl."user"
|
||||
"""
|
||||
|
|
@ -7555,15 +7558,24 @@ async def get_team_spend_by_user(
|
|||
|
||||
own_user_only: Final = scope.api_key_filter is not None
|
||||
user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else ()
|
||||
scoped_api_keys: Final = await spend_log_api_keys_for_teams(
|
||||
prisma_client, scoped_team_ids, date.fromisoformat(start_date), date.fromisoformat(end_date)
|
||||
)
|
||||
api_key_param: Final = (list(scoped_api_keys),) if scoped_api_keys is not None else ()
|
||||
async with db_span("team_user_spend", "LiteLLM_SpendLogs"):
|
||||
rows: Final = _TEAM_USER_SPEND_ROWS.validate_python(
|
||||
await query_raw_within_statement_timeout(
|
||||
prisma_client,
|
||||
_team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
|
||||
_team_user_spend_sql(
|
||||
team_count=len(scoped_team_ids),
|
||||
restrict_to_user=own_user_only,
|
||||
narrow_to_api_keys=scoped_api_keys is not None,
|
||||
),
|
||||
start_date,
|
||||
end_date,
|
||||
*scoped_team_ids,
|
||||
*user_param,
|
||||
*api_key_param,
|
||||
what="Spend per user",
|
||||
)
|
||||
)
|
||||
|
|
|
|||
83
litellm/proxy/spend_tracking/spend_log_api_key_scope.py
Normal file
83
litellm/proxy/spend_tracking/spend_log_api_key_scope.py
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
"""``LiteLLM_SpendLogs`` has no index on ``"user"`` or ``team_id``, so a read filtered by owner scans the
|
||||
whole date window. The owner's ``LiteLLM_Daily*Spend`` rows (kept after key deletion) plus their live keys
|
||||
name every api_key that produced those logs, and ``(api_key, "startTime")`` is indexed."""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import date, timedelta
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.proxy.utils import PrismaClient, query_raw_within_statement_timeout
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
SPEND_LOG_API_KEY_SCOPE_MAX_KEYS: Final = 5000
|
||||
_WHAT: Final = "Spend log key lookup"
|
||||
|
||||
_USER_API_KEYS_SQL: Final = """
|
||||
SELECT DISTINCT api_key FROM (
|
||||
SELECT api_key FROM "LiteLLM_DailyUserSpend"
|
||||
WHERE user_id = $1 AND date >= $2 AND date <= $3
|
||||
UNION ALL
|
||||
SELECT token AS api_key FROM "LiteLLM_VerificationToken" WHERE user_id = $1
|
||||
) owned
|
||||
LIMIT $4
|
||||
"""
|
||||
|
||||
_TEAM_API_KEYS_SQL: Final = """
|
||||
SELECT DISTINCT api_key FROM (
|
||||
SELECT api_key FROM "LiteLLM_DailyTeamSpend"
|
||||
WHERE team_id = ANY($1::text[]) AND date >= $2 AND date <= $3
|
||||
UNION ALL
|
||||
SELECT token AS api_key FROM "LiteLLM_VerificationToken" WHERE team_id = ANY($1::text[])
|
||||
) owned
|
||||
LIMIT $4
|
||||
"""
|
||||
|
||||
|
||||
class _ApiKeyRow(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
api_key: str
|
||||
|
||||
|
||||
_API_KEY_ROWS: Final = TypeAdapter(tuple[_ApiKeyRow, ...])
|
||||
|
||||
|
||||
def _narrowing_keys(rows: Sequence[Mapping[str, object]]) -> tuple[str, ...] | None:
|
||||
"""None when the list cannot narrow the read: nothing known about the owner, or too many keys to help."""
|
||||
keys: Final = tuple(row.api_key for row in _API_KEY_ROWS.validate_python(rows))
|
||||
return None if not keys or len(keys) > SPEND_LOG_API_KEY_SCOPE_MAX_KEYS else keys
|
||||
|
||||
|
||||
def _daily_dates_around(start: date, end: date) -> tuple[str, str]:
|
||||
return (start - timedelta(days=1)).isoformat(), (end + timedelta(days=1)).isoformat()
|
||||
|
||||
|
||||
async def spend_log_api_keys_for_user(
|
||||
prisma_client: PrismaClient, user_id: str, start: date, end: date
|
||||
) -> tuple[str, ...] | None:
|
||||
"""The api_keys that can carry this user's spend logs in the window, or None when there are too many to help."""
|
||||
first, last = _daily_dates_around(start, end)
|
||||
return _narrowing_keys(
|
||||
await query_raw_within_statement_timeout(
|
||||
prisma_client, _USER_API_KEYS_SQL, user_id, first, last, SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1, what=_WHAT
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def spend_log_api_keys_for_teams(
|
||||
prisma_client: PrismaClient, team_ids: Sequence[str], start: date, end: date
|
||||
) -> tuple[str, ...] | None:
|
||||
first, last = _daily_dates_around(start, end)
|
||||
return _narrowing_keys(
|
||||
await query_raw_within_statement_timeout(
|
||||
prisma_client,
|
||||
_TEAM_API_KEYS_SQL,
|
||||
list(team_ids),
|
||||
first,
|
||||
last,
|
||||
SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1,
|
||||
what=_WHAT,
|
||||
)
|
||||
)
|
||||
|
|
@ -70,6 +70,10 @@ from litellm.proxy.spend_tracking.spend_capture_rate import (
|
|||
# NOTE: Avoid module-level import from common_utils: proxy_server imports this
|
||||
# module while common_utils may pull proxy_server during init, which can leave
|
||||
# those names undefined. Import the helpers locally where they are used.
|
||||
from litellm.proxy.spend_tracking.spend_log_api_key_scope import (
|
||||
spend_log_api_keys_for_teams,
|
||||
spend_log_api_keys_for_user,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
get_spend_by_team,
|
||||
get_spend_by_team_and_customer,
|
||||
|
|
@ -257,6 +261,22 @@ async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object)
|
|||
return await prisma_client.db.query_raw(sql_query, *args)
|
||||
|
||||
|
||||
async def _scope_api_keys(
|
||||
prisma_client: PrismaClient, *, where_conditions: Mapping[str, object], start: datetime | None, end: datetime | None
|
||||
) -> tuple[str, ...] | None:
|
||||
"""The api_keys behind a ``"user"`` / ``team_id`` filter, so the unindexed column is only checked on rows the
|
||||
``(api_key, "startTime")`` index already picked."""
|
||||
if start is None or end is None or isinstance(where_conditions.get("api_key"), str):
|
||||
return None
|
||||
user_id: Final = where_conditions.get("user")
|
||||
if isinstance(user_id, str):
|
||||
return await spend_log_api_keys_for_user(prisma_client, user_id, start.date(), end.date())
|
||||
team_id: Final = where_conditions.get("team_id")
|
||||
if isinstance(team_id, str):
|
||||
return await spend_log_api_keys_for_teams(prisma_client, (team_id,), start.date(), end.date())
|
||||
return None
|
||||
|
||||
|
||||
async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT] | None:
|
||||
"""``_query_raw`` for the call sites that guard the result against ``None``."""
|
||||
return await _query_raw(prisma_client, sql_query, *args)
|
||||
|
|
@ -2950,6 +2970,18 @@ async def ui_view_spend_logs(
|
|||
sql_params.append(f"%{like_escaped_session_id}%")
|
||||
p += 1
|
||||
|
||||
scope_api_keys: Final = (
|
||||
await _scope_api_keys(
|
||||
prisma_client, where_conditions=where_conditions, start=start_date_obj, end=end_date_obj
|
||||
)
|
||||
if not search_owns_window
|
||||
else None
|
||||
)
|
||||
if scope_api_keys is not None:
|
||||
sql_conditions.append(f"api_key = ANY(${p}::text[])")
|
||||
sql_params.append(list(scope_api_keys))
|
||||
p += 1 # rebind-ok: advances the file's shared $N placeholder counter
|
||||
|
||||
# Status filter
|
||||
if status_filter is not None and not (group_by_session is True and not is_search_lookup):
|
||||
if status_filter == "success":
|
||||
|
|
|
|||
|
|
@ -17341,6 +17341,19 @@ def _team_spend_by_user_caller(user_id: str, teams: list[str]) -> LiteLLM_UserTa
|
|||
)
|
||||
|
||||
|
||||
def _team_spend_reads(api_keys: list[str], spend_rows: list[dict] | Exception) -> AsyncMock:
|
||||
"""Raw reads inside the statement-timeout transaction: the key lookup, then the spend log aggregate."""
|
||||
|
||||
async def read(sql: str, *params: object) -> list[dict]:
|
||||
if 'FROM "LiteLLM_DailyTeamSpend"' in sql:
|
||||
return [{"api_key": api_key} for api_key in api_keys]
|
||||
if isinstance(spend_rows, Exception):
|
||||
raise spend_rows
|
||||
return spend_rows
|
||||
|
||||
return AsyncMock(side_effect=read)
|
||||
|
||||
|
||||
def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, requests: int) -> dict:
|
||||
return {
|
||||
"team_id": team_id,
|
||||
|
|
@ -17367,12 +17380,13 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m
|
|||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta])
|
||||
transaction = statement_timeout_reads(
|
||||
mock_db_client,
|
||||
AsyncMock(
|
||||
return_value=[
|
||||
_team_spend_reads(
|
||||
["sk-alpha", "sk-beta"],
|
||||
[
|
||||
_team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3),
|
||||
_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2),
|
||||
_team_spend_by_user_db_row("team-beta", "alice", 0.1, 1),
|
||||
]
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -17386,11 +17400,15 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m
|
|||
assert any("statement_timeout = 30000" in statement for statement in transaction.statements)
|
||||
mock_db_client.db.tx.assert_not_called()
|
||||
sql, *params = transaction.query_raw.call_args.args
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"]
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta", ["sk-alpha", "sk-beta"]]
|
||||
assert 'FROM "LiteLLM_SpendLogs" sl' in sql
|
||||
assert 'sl."startTime" >= $1::timestamp' in sql
|
||||
assert "sl.\"startTime\" < $2::timestamp + INTERVAL '1 day'" in sql
|
||||
assert "sl.team_id IN ($3, $4)" in sql
|
||||
assert "sl.api_key = ANY($5::text[])" in sql
|
||||
lookup_sql, *lookup_params = transaction.query_raw.call_args_list[0].args
|
||||
assert 'FROM "LiteLLM_DailyTeamSpend"' in lookup_sql and 'FROM "LiteLLM_VerificationToken"' in lookup_sql
|
||||
assert lookup_params == [["team-alpha", "team-beta"], "2026-08-31", "2026-09-05", 5001]
|
||||
assert 'GROUP BY sl.team_id, sl."user"' in sql
|
||||
assert 'sl."user" = $' not in sql
|
||||
|
||||
|
|
@ -17412,7 +17430,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien
|
|||
caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||
transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[]))
|
||||
transaction = statement_timeout_reads(mock_db_client, _team_spend_reads(["sk-alpha"], []))
|
||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=_team_spend_by_user_caller("alice", ["team-alpha"])
|
||||
)
|
||||
|
|
@ -17422,8 +17440,9 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien
|
|||
)
|
||||
|
||||
sql, *params = transaction.query_raw.call_args.args
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha"]
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", ["sk-alpha"]]
|
||||
assert 'sl."user" = $' not in sql
|
||||
assert "sl.api_key = ANY($4::text[])" in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -17435,7 +17454,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli
|
|||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
transaction = statement_timeout_reads(
|
||||
mock_db_client, AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)])
|
||||
mock_db_client, _team_spend_reads(["sk-bob"], [_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)])
|
||||
)
|
||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
||||
|
|
@ -17446,9 +17465,10 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli
|
|||
)
|
||||
|
||||
sql, *params = transaction.query_raw.call_args.args
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"]
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob", ["sk-bob"]]
|
||||
assert "sl.team_id IN ($3)" in sql
|
||||
assert 'AND sl."user" = $4' in sql
|
||||
assert "sl.api_key = ANY($5::text[])" in sql
|
||||
assert [(r.user_id, r.spend) for r in response.results] == [("bob", 0.25)]
|
||||
|
||||
|
||||
|
|
@ -17506,7 +17526,7 @@ async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_clie
|
|||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||
statement_timeout_reads(mock_db_client, AsyncMock(side_effect=statement_timeout_error()))
|
||||
statement_timeout_reads(mock_db_client, _team_spend_reads(["sk-alpha"], statement_timeout_error()))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_team_spend_by_user(
|
||||
|
|
@ -17520,6 +17540,30 @@ async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_clie
|
|||
assert "longer than 30s" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("known_key_count", [0, 5001])
|
||||
async def test_get_team_spend_by_user_without_a_usable_key_list_keeps_the_plain_team_filter(
|
||||
mock_db_client, known_key_count: int
|
||||
):
|
||||
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
||||
|
||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||
transaction = statement_timeout_reads(
|
||||
mock_db_client, _team_spend_reads([f"sk-{i}" for i in range(known_key_count)], [])
|
||||
)
|
||||
|
||||
await get_team_spend_by_user(
|
||||
user_api_key_dict=admin, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
|
||||
)
|
||||
|
||||
sql, *params = transaction.query_raw.call_args.args
|
||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha"]
|
||||
assert "sl.team_id IN ($3)" in sql
|
||||
assert "sl.api_key = ANY(" not in sql
|
||||
|
||||
|
||||
class _TeamRowWithOrganization(LiteLLM_TeamTable):
|
||||
litellm_organization_table: LiteLLM_OrganizationTable | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import hashlib
|
|||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
|
||||
from datetime import timezone
|
||||
from functools import partial
|
||||
from types import SimpleNamespace
|
||||
|
|
@ -141,6 +141,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond)
|
||||
status = re.fullmatch(r"status = \$(\d+)", cond)
|
||||
api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond)
|
||||
api_key_any = re.fullmatch(r"api_key = ANY\(\$(\d+)::text\[\]\)", cond)
|
||||
req_or_call = re.fullmatch(r"\(request_id = \$(\d+) OR litellm_call_id = \$\1\)", cond)
|
||||
if req_or_call:
|
||||
where["request_id_or_call_id"] = params[int(req_or_call.group(1)) - 1]
|
||||
|
|
@ -169,6 +170,8 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
where["session_id"] = {"contains": str(params[int(sess.group(1)) - 1]).strip("%")}
|
||||
elif status:
|
||||
where["status"] = {"equals": params[int(status.group(1)) - 1]}
|
||||
elif api_key_any:
|
||||
where["api_key_in"] = list(params[int(api_key_any.group(1)) - 1])
|
||||
elif api_key_not_in:
|
||||
where["api_key_not_in"] = [
|
||||
params[int(api_key_not_in.group(1)) - 1],
|
||||
|
|
@ -218,7 +221,18 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
return where
|
||||
|
||||
|
||||
def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None):
|
||||
def _emulate_owner_api_key_lookup(mock_spend_logs: list[dict], sql_query: str, params: tuple) -> list[dict]:
|
||||
"""The daily spend tables and key table name every api_key an owner's requests ran under, so the
|
||||
lookup answers with the distinct api_keys of the owner's logs."""
|
||||
owner_col = "user" if 'FROM "LiteLLM_DailyUserSpend"' in sql_query else "team_id"
|
||||
owners = {params[0]} if isinstance(params[0], str) else set(params[0])
|
||||
api_keys = {log["api_key"] for log in mock_spend_logs if log.get(owner_col) in owners}
|
||||
return [{"api_key": api_key} for api_key in sorted(api_keys)]
|
||||
|
||||
|
||||
def make_ui_spend_logs_mock_prisma(
|
||||
mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None, owner_api_keys_fn=None
|
||||
):
|
||||
"""
|
||||
Create a MockPrismaClient for /spend/logs/ui endpoint tests.
|
||||
|
||||
|
|
@ -228,6 +242,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No
|
|||
where_conditions, returns the filtered list of logs.
|
||||
team_lookup_fn: Optional async callable for team RBAC (find_unique).
|
||||
If provided, adds litellm_teamtable to db.
|
||||
owner_api_keys_fn: Optional Callable[[tuple], list[dict]] answering the
|
||||
user/team -> api_key lookup instead of deriving it from the logs.
|
||||
"""
|
||||
|
||||
class MockDB:
|
||||
|
|
@ -251,6 +267,10 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No
|
|||
return []
|
||||
if 'SELECT DISTINCT "user", team_id' in sql_query:
|
||||
return _emulate_spend_log_owner_lookup(mock_spend_logs, sql_query, params)
|
||||
if 'FROM "LiteLLM_DailyUserSpend"' in sql_query or 'FROM "LiteLLM_DailyTeamSpend"' in sql_query:
|
||||
if owner_api_keys_fn is not None:
|
||||
return owner_api_keys_fn(params)
|
||||
return _emulate_owner_api_key_lookup(mock_spend_logs, sql_query, params)
|
||||
filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params))
|
||||
total = len(filtered)
|
||||
if "COUNT(*)" in sql_query:
|
||||
|
|
@ -845,6 +865,196 @@ async def test_ui_view_spend_logs_with_user_id(client, monkeypatch):
|
|||
assert data["data"][0]["user"] == "test_user_1"
|
||||
|
||||
|
||||
def _owned_log(request_id: str, api_key: str, user: str, team_id: str) -> dict:
|
||||
return {
|
||||
"id": request_id,
|
||||
"request_id": request_id,
|
||||
"api_key": api_key,
|
||||
"user": user,
|
||||
"team_id": team_id,
|
||||
"spend": 0.01,
|
||||
"startTime": "2026-09-02T10:00:00+00:00",
|
||||
"model": "gpt-4o",
|
||||
}
|
||||
|
||||
|
||||
def _owner_filtered_logs(logs: list[dict]) -> Callable[[dict], list[dict]]:
|
||||
"""What Postgres returns for the page query: the owner filter and, when present, the api_key list both apply."""
|
||||
|
||||
def filter_fn(where: dict) -> list[dict]:
|
||||
return [
|
||||
log
|
||||
for log in logs
|
||||
if log["user"] == where.get("user", log["user"])
|
||||
and log["team_id"] == where.get("team_id", log["team_id"])
|
||||
and log["api_key"] in where.get("api_key_in", [log["api_key"]])
|
||||
]
|
||||
|
||||
return filter_fn
|
||||
|
||||
|
||||
def _owner_api_keys(daily_rows: dict[str, set[str]], live_keys: dict[str, set[str]]) -> Callable[[tuple], list[dict]]:
|
||||
"""The key lookup's two sources, keyed by owner: api_keys on daily spend rows and api_keys of live keys."""
|
||||
|
||||
def lookup(params: tuple) -> list[dict]:
|
||||
owners = {params[0]} if isinstance(params[0], str) else set(params[0])
|
||||
api_keys = set().union(*(daily_rows.get(o, set()) | live_keys.get(o, set()) for o in owners))
|
||||
return [{"api_key": api_key} for api_key in sorted(api_keys)]
|
||||
|
||||
return lookup
|
||||
|
||||
|
||||
def _spend_logs_page_query(queries: list[tuple[str, tuple]]) -> tuple[str, tuple, list[str] | None]:
|
||||
sql, params = next((sql, params) for sql, params in queries if '"LiteLLM_SpendLogs"' in sql and "ORDER BY" in sql)
|
||||
key_param = re.search(r"api_key = ANY\(\$(\d+)::text\[\]\)", sql)
|
||||
return sql, params, None if key_param is None else params[int(key_param.group(1)) - 1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_user_filter_reads_spend_logs_through_the_users_api_keys(client, monkeypatch):
|
||||
from litellm.proxy.spend_tracking.spend_log_api_key_scope import SPEND_LOG_API_KEY_SCOPE_MAX_KEYS
|
||||
|
||||
logs = [
|
||||
_owned_log("req1", "sk-live", "alice", "team1"),
|
||||
_owned_log("req2", "sk-deleted", "alice", "team1"),
|
||||
_owned_log("req3", "sk-bob", "bob", "team1"),
|
||||
]
|
||||
queries: list[tuple[str, tuple]] = []
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(
|
||||
logs,
|
||||
_owner_filtered_logs(logs),
|
||||
query_observer=lambda sql, params: queries.append((sql, params)),
|
||||
owner_api_keys_fn=_owner_api_keys(
|
||||
daily_rows={"alice": {"sk-deleted"}, "bob": {"sk-bob"}}, live_keys={"alice": {"sk-live"}}
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"user_id": "alice", "start_date": "2026-09-01 00:00:00", "end_date": "2026-09-03 23:59:59"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert {row["request_id"] for row in response.json()["data"]} == {"req1", "req2"}
|
||||
lookup_sql, lookup_params = next((sql, params) for sql, params in queries if '"LiteLLM_DailyUserSpend"' in sql)
|
||||
assert '"LiteLLM_VerificationToken"' in lookup_sql
|
||||
assert lookup_params == ("alice", "2026-08-31", "2026-09-04", SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1)
|
||||
page_sql, _, scoped_api_keys = _spend_logs_page_query(queries)
|
||||
assert scoped_api_keys == ["sk-deleted", "sk-live"]
|
||||
assert '"user" = $' in page_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_team_filter_reads_spend_logs_through_the_teams_api_keys(client, monkeypatch):
|
||||
logs = [
|
||||
_owned_log("req1", "sk-alpha-1", "alice", "team-alpha"),
|
||||
_owned_log("req2", "sk-alpha-2", "bob", "team-alpha"),
|
||||
_owned_log("req3", "sk-beta", "carol", "team-beta"),
|
||||
]
|
||||
queries: list[tuple[str, tuple]] = []
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(
|
||||
logs,
|
||||
_owner_filtered_logs(logs),
|
||||
query_observer=lambda sql, params: queries.append((sql, params)),
|
||||
owner_api_keys_fn=_owner_api_keys(
|
||||
daily_rows={"team-alpha": {"sk-alpha-1"}, "team-beta": {"sk-beta"}},
|
||||
live_keys={"team-alpha": {"sk-alpha-2"}},
|
||||
),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe",
|
||||
lambda user_api_key_dict: True,
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
app.dependency_overrides,
|
||||
ps.user_api_key_auth,
|
||||
lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"),
|
||||
)
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"team_id": "team-alpha", "start_date": "2026-09-01 00:00:00", "end_date": "2026-09-03 23:59:59"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert {row["request_id"] for row in response.json()["data"]} == {"req1", "req2"}
|
||||
_, lookup_params = next((sql, params) for sql, params in queries if '"LiteLLM_DailyTeamSpend"' in sql)
|
||||
assert lookup_params[:3] == (["team-alpha"], "2026-08-31", "2026-09-04")
|
||||
page_sql, _, scoped_api_keys = _spend_logs_page_query(queries)
|
||||
assert scoped_api_keys == ["sk-alpha-1", "sk-alpha-2"]
|
||||
assert "team_id = $" in page_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("known_keys", [set(), {f"sk-{i}" for i in range(5001)}])
|
||||
async def test_ui_view_spend_logs_without_a_usable_key_list_keeps_the_plain_user_filter(
|
||||
client, monkeypatch, known_keys: set[str]
|
||||
):
|
||||
logs = [_owned_log("req1", "sk-live", "alice", "team1")]
|
||||
queries: list[tuple[str, tuple]] = []
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(
|
||||
logs,
|
||||
_owner_filtered_logs(logs),
|
||||
query_observer=lambda sql, params: queries.append((sql, params)),
|
||||
owner_api_keys_fn=_owner_api_keys(daily_rows={"alice": known_keys}, live_keys={}),
|
||||
),
|
||||
)
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"user_id": "alice", "start_date": "2026-09-01 00:00:00", "end_date": "2026-09-03 23:59:59"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [row["request_id"] for row in response.json()["data"]] == ["req1"]
|
||||
page_sql, _, scoped_api_keys = _spend_logs_page_query(queries)
|
||||
assert scoped_api_keys is None
|
||||
assert '"user" = $' in page_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_search_with_user_filter_is_not_narrowed_to_the_windows_keys(client, monkeypatch):
|
||||
logs = [_owned_log("req-old", "sk-retired", "alice", "team1")]
|
||||
queries: list[tuple[str, tuple]] = []
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(
|
||||
logs,
|
||||
_owner_filtered_logs(logs),
|
||||
query_observer=lambda sql, params: queries.append((sql, params)),
|
||||
owner_api_keys_fn=_owner_api_keys(daily_rows={}, live_keys={"alice": {"sk-live"}}),
|
||||
),
|
||||
)
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"user_id": "alice",
|
||||
"search": "req-old",
|
||||
"start_date": "2026-09-01 00:00:00",
|
||||
"end_date": "2026-09-03 23:59:59",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [row["request_id"] for row in response.json()["data"]] == ["req-old"]
|
||||
page_sql, _, scoped_api_keys = _spend_logs_page_query(queries)
|
||||
assert scoped_api_keys is None
|
||||
assert "request_id = $" in page_sql and '"user" = $' in page_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"session_id_query,expected_request_ids",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue