mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): bound legacy spend log summaries
This commit is contained in:
KirschQAQ
parent
476566d260
commit
c14b34edaf
2 changed files with 171 additions and 6 deletions
|
|
@ -42,6 +42,8 @@ router: Final = APIRouter()
|
|||
SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000
|
||||
LEGACY_SPEND_LOGS_MAX_PAGE_SIZE: Final = 1000
|
||||
LEGACY_SPEND_LOGS_MAX_SKIP: Final = SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS: Final = SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
LEGACY_SPEND_LOGS_MAX_SUMMARY_DAYS: Final = SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
|
||||
_RowT = TypeVar("_RowT")
|
||||
|
||||
|
|
@ -294,9 +296,20 @@ async def _find_legacy_spend_log_page(
|
|||
take,
|
||||
skip,
|
||||
)
|
||||
_hydrate_spend_log_metadata(rows)
|
||||
_hydrate_spend_log_request_tags(rows)
|
||||
return rows
|
||||
response_rows: Final = tuple(
|
||||
{
|
||||
**row,
|
||||
"messages": None,
|
||||
"response": None,
|
||||
"proxy_server_request": None,
|
||||
}
|
||||
if isinstance(row, dict)
|
||||
else row
|
||||
for row in rows
|
||||
)
|
||||
_hydrate_spend_log_metadata(response_rows)
|
||||
_hydrate_spend_log_request_tags(response_rows)
|
||||
return response_rows
|
||||
|
||||
|
||||
async def _summarize_legacy_spend_logs(
|
||||
|
|
@ -326,8 +339,14 @@ async def _summarize_legacy_spend_logs(
|
|||
WHERE {where_sql}
|
||||
GROUP BY spend_date, api_key, "user", model
|
||||
ORDER BY spend_date ASC, api_key ASC, "user" ASC, model ASC
|
||||
LIMIT ${len(query_params) + 1}
|
||||
"""
|
||||
rows: Final[Sequence[_LegacySpendSummaryRow]] = await _query_raw(prisma_client, sql_query, *query_params)
|
||||
rows: Final[Sequence[_LegacySpendSummaryRow]] = await _query_raw(
|
||||
prisma_client,
|
||||
sql_query,
|
||||
*query_params,
|
||||
LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS + 1,
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
|
|
@ -3128,6 +3147,15 @@ async def view_spend_logs(
|
|||
start_date_obj: Final = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
end_date_obj: Final = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
summary_day_count: Final = (end_date_obj.date() - start_date_obj.date()).days + 1
|
||||
if summarize and summary_day_count > LEGACY_SPEND_LOGS_MAX_SUMMARY_DAYS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail=(
|
||||
f"summarize=true date ranges must be no more than {LEGACY_SPEND_LOGS_MAX_SUMMARY_DAYS} days"
|
||||
),
|
||||
)
|
||||
|
||||
# Convert to ISO format strings for Prisma
|
||||
start_date_iso: Final = start_date_obj.isoformat()
|
||||
end_date_iso: Final = end_date_obj.isoformat()
|
||||
|
|
@ -3181,6 +3209,14 @@ async def view_spend_logs(
|
|||
request_id=effective_request_id,
|
||||
)
|
||||
|
||||
if len(response) > LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail=(
|
||||
"summarized spend logs exceed the server result limit; narrow the date range or add a filter"
|
||||
),
|
||||
)
|
||||
|
||||
if len(response) > 0:
|
||||
daily_spend: Final[dict[date, float]] = {} # mutable-ok: aggregate database groups by day
|
||||
daily_api_spend: Final[dict[date, dict[str, float]]] = {} # mutable-ok: aggregate groups by key
|
||||
|
|
|
|||
|
|
@ -3037,6 +3037,8 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch):
|
|||
async def query_raw(self, sql_query, *params):
|
||||
if "GROUP BY" not in sql_query:
|
||||
return mock_spend_logs
|
||||
assert len(params) == 3
|
||||
assert params[-1] == spend_management_endpoints.LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS + 1
|
||||
yesterday = datetime.datetime.now(timezone.utc) - timedelta(days=1)
|
||||
return [
|
||||
{
|
||||
|
|
@ -3316,8 +3318,9 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
async def query_raw(self, sql_query, *params):
|
||||
assert "date_trunc('day', \"startTime\")" in sql_query
|
||||
assert "GROUP BY" in sql_query
|
||||
assert len(params) == 2
|
||||
assert all(isinstance(value, datetime) for value in params)
|
||||
assert len(params) == 3
|
||||
assert all(isinstance(value, datetime) for value in params[:2])
|
||||
assert params[-1] == spend_management_endpoints.LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS + 1
|
||||
return [
|
||||
{
|
||||
"spend_date": (datetime.now(timezone.utc) - timedelta(days=1)).date(),
|
||||
|
|
@ -4048,6 +4051,7 @@ def _assert_legacy_list_query(mock_prisma, expected_where, expected_take, expect
|
|||
query_args = mock_prisma.db.query_raw.await_args.args
|
||||
assert _reconstruct_ui_where_from_sql(query_args[0], query_args[1:]) == expected_where
|
||||
assert query_args[-2:] == (expected_take, expected_skip)
|
||||
assert 'ORDER BY "startTime" DESC, request_id DESC' in query_args[0]
|
||||
assert "messages" not in query_args[0]
|
||||
assert "response" not in query_args[0]
|
||||
assert "proxy_server_request" not in query_args[0]
|
||||
|
|
@ -4087,6 +4091,9 @@ def test_legacy_spend_logs_list_query_omits_heavy_payload_columns(
|
|||
response, mock_prisma = legacy_spend_logs_request(spend_logs)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()[0]["messages"] is None
|
||||
assert response.json()[0]["response"] is None
|
||||
assert response.json()[0]["proxy_server_request"] is None
|
||||
mock_prisma.db.query_raw.assert_awaited_once()
|
||||
sql_query = mock_prisma.db.query_raw.await_args.args[0]
|
||||
assert "messages" not in sql_query
|
||||
|
|
@ -4235,6 +4242,128 @@ def test_legacy_spend_logs_summary_aggregates_by_day_in_database(
|
|||
mock_prisma.db.group_by.assert_not_awaited()
|
||||
|
||||
|
||||
def test_legacy_spend_logs_summary_rejects_high_cardinality_groups(client, monkeypatch):
|
||||
mock_prisma = _LegacySpendLogsPaginationPrismaClient([])
|
||||
summary_row = {
|
||||
"spend_date": datetime.date(2024, 1, 1),
|
||||
"api_key": "key-a",
|
||||
"user": "user-a",
|
||||
"model": "model-a",
|
||||
"spend": 1.0,
|
||||
}
|
||||
mock_prisma.db.query_raw.side_effect = None
|
||||
mock_prisma.db.query_raw.return_value = [
|
||||
{**summary_row, "model": f"model-{index}"}
|
||||
for index in range(spend_management_endpoints.LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS + 1)
|
||||
]
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": "2024-01-01",
|
||||
"end_date": "2024-01-02",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 422
|
||||
mock_prisma.db.query_raw.assert_awaited_once()
|
||||
query_args = mock_prisma.db.query_raw.await_args.args
|
||||
assert query_args[-1] == spend_management_endpoints.LEGACY_SPEND_LOGS_MAX_SUMMARY_GROUPS + 1
|
||||
assert "LIMIT $3" in query_args[0]
|
||||
|
||||
|
||||
def test_legacy_spend_logs_summary_rejects_excessive_date_range(client, monkeypatch):
|
||||
mock_prisma = _LegacySpendLogsPaginationPrismaClient([])
|
||||
mock_prisma.db.query_raw.side_effect = RuntimeError("summary query must not run")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": "2024-01-01",
|
||||
"end_date": "2051-05-20",
|
||||
"summarize": "true",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 422
|
||||
mock_prisma.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
def test_legacy_spend_logs_summary_accepts_maximum_date_range(client, monkeypatch):
|
||||
mock_prisma = _LegacySpendLogsPaginationPrismaClient([])
|
||||
mock_prisma.db.query_raw.side_effect = None
|
||||
mock_prisma.db.query_raw.return_value = []
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
start_date = datetime.date(2024, 1, 1)
|
||||
end_date = start_date + datetime.timedelta(
|
||||
days=spend_management_endpoints.LEGACY_SPEND_LOGS_MAX_SUMMARY_DAYS - 1
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": start_date.isoformat(),
|
||||
"end_date": end_date.isoformat(),
|
||||
"summarize": "true",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == []
|
||||
mock_prisma.db.query_raw.assert_awaited_once()
|
||||
|
||||
|
||||
def test_legacy_spend_logs_unsummarized_range_remains_bounded_past_summary_day_cap(
|
||||
client, monkeypatch
|
||||
):
|
||||
mock_prisma = _LegacySpendLogsPaginationPrismaClient([])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": "2024-01-01",
|
||||
"end_date": "2025-01-02",
|
||||
"summarize": "false",
|
||||
"page_size": 1,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert isinstance(response.json(), list)
|
||||
assert mock_prisma.db.query_raw.await_args.args[-2:] == (1, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_spend_logs_direct_call_uses_pagination_defaults(monkeypatch):
|
||||
spend_logs = [
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue