mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
perf(spend): group /spend/logs summary by day in Postgres instead of per-row Prisma group_by (#39351)
* perf(spend): group /spend/logs summary by day in Postgres instead of per-row Prisma group_by Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(spend): simplify /spend/logs daily summary aggregation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend): preserve spend logs response schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend): compare spend log range bounds as naive UTC timestamps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: cover spend logs summary edge cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: cover spend summary request filters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet lint budgets after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
39a17898ff
commit
cf3af0f486
5 changed files with 303 additions and 114 deletions
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 14074
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2214
|
||||
"limit": 2206
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 319
|
||||
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5601
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15287
|
||||
"limit": 15285
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44362
|
||||
"limit": 44360
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38323
|
||||
"limit": 38311
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19624
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29861
|
||||
"limit": 29847
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 111
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ import collections
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -16,7 +17,6 @@ from typing import (
|
|||
TypeAlias,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings
|
||||
)
|
||||
|
||||
import fastapi
|
||||
|
|
@ -201,16 +201,12 @@ class _SessionSpendStats(NamedTuple):
|
|||
_SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats]
|
||||
|
||||
|
||||
class _SpendSumAggregate(TypedDict, total=False):
|
||||
spend: ReadOnly[float]
|
||||
|
||||
|
||||
class _SpendGroupByRow(TypedDict):
|
||||
class _SpendDailySummaryRow(TypedDict):
|
||||
day: ReadOnly[str]
|
||||
api_key: ReadOnly[str]
|
||||
user: ReadOnly[str | None]
|
||||
model: ReadOnly[str]
|
||||
startTime: ReadOnly[object]
|
||||
_sum: ReadOnly[_SpendSumAggregate]
|
||||
spend: ReadOnly[float]
|
||||
|
||||
|
||||
async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
|
||||
|
|
@ -251,6 +247,66 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken
|
|||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _spend_logs_daily_summary_sql(
|
||||
*,
|
||||
start_date_iso: str,
|
||||
end_date_iso: str,
|
||||
api_key: str | None,
|
||||
request_id: str | None,
|
||||
user_id: str | None,
|
||||
) -> tuple[str, tuple[object, ...]]:
|
||||
filter_params: Final[tuple[tuple[str, object], ...]] = tuple(
|
||||
(column, value)
|
||||
for column, value in (
|
||||
("api_key", api_key),
|
||||
("request_id", request_id),
|
||||
('"user"', user_id),
|
||||
)
|
||||
if value is not None
|
||||
)
|
||||
filter_clauses: Final[tuple[str, ...]] = tuple(
|
||||
f"AND {column} = ${index}" for index, (column, _) in enumerate(filter_params, start=3)
|
||||
)
|
||||
filter_sql: Final = "\n".join(filter_clauses)
|
||||
sql_query: Final = f"""
|
||||
SELECT
|
||||
to_char(date_trunc('day', "startTime"), 'YYYY-MM-DD') AS day,
|
||||
api_key,
|
||||
"user",
|
||||
model,
|
||||
SUM(spend) AS spend
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND "startTime" <= ($2::timestamptz AT TIME ZONE 'UTC')
|
||||
{filter_sql}
|
||||
GROUP BY 1, 2, 3, 4
|
||||
ORDER BY 1
|
||||
"""
|
||||
params: Final[tuple[object, ...]] = (
|
||||
start_date_iso,
|
||||
end_date_iso,
|
||||
*(value for _, value in filter_params),
|
||||
)
|
||||
return sql_query, params
|
||||
|
||||
|
||||
def _sum_spend_by(
|
||||
rows: Sequence[_SpendDailySummaryRow], column: Literal["api_key", "user", "model"]
|
||||
) -> Mapping[str | None, float]:
|
||||
keys: Final = frozenset(row[column] for row in rows)
|
||||
return {key: sum(float(row["spend"]) for row in rows if row[column] == key) for key in keys}
|
||||
|
||||
|
||||
def _daily_summary_item(summary_date: date, rows: Sequence[_SpendDailySummaryRow]) -> Mapping[str, object]:
|
||||
api_key_spend: Final = {key: value for key, value in _sum_spend_by(rows, "api_key").items() if key is not None}
|
||||
return {
|
||||
**api_key_spend,
|
||||
"startTime": summary_date,
|
||||
"spend": sum(float(row["spend"]) for row in rows),
|
||||
"users": _sum_spend_by(rows, "user"),
|
||||
"models": _sum_spend_by(rows, "model"),
|
||||
}
|
||||
|
||||
|
||||
async def _find_spend_logs(
|
||||
prisma_client: PrismaClient,
|
||||
where: Mapping[str, object],
|
||||
|
|
@ -3266,18 +3322,22 @@ async def view_spend_logs(
|
|||
start_date_iso: Final = start_date_obj.isoformat()
|
||||
end_date_iso: Final = end_date_obj.isoformat()
|
||||
|
||||
filter_query: Final = {
|
||||
filter_query: Final[
|
||||
dict[str, object]
|
||||
] = { # mutable-ok: legacy filters are extended for optional parameters
|
||||
"startTime": {
|
||||
"gte": start_date_iso, # Greater than or equal to Start Date
|
||||
"lte": end_date_iso, # Less than or equal to End Date
|
||||
}
|
||||
}
|
||||
|
||||
summary_api_key: Final[str | None] = (
|
||||
prisma_client.hash_token(token=api_key)
|
||||
if api_key is not None and api_key.startswith("sk-")
|
||||
else api_key
|
||||
)
|
||||
if api_key is not None and isinstance(api_key, str):
|
||||
if api_key.startswith("sk-"):
|
||||
filter_query["api_key"] = prisma_client.hash_token(token=api_key)
|
||||
else:
|
||||
filter_query["api_key"] = api_key
|
||||
filter_query["api_key"] = summary_api_key
|
||||
if request_id is not None and isinstance(request_id, str):
|
||||
filter_query["request_id"] = request_id
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
|
|
@ -3296,58 +3356,34 @@ async def view_spend_logs(
|
|||
return data
|
||||
|
||||
# Legacy behavior: return summarized data (when summarize=true)
|
||||
# SQL query
|
||||
response: Final = await SpendLogsRepository(prisma_client).table.group_by(
|
||||
by=["api_key", "user", "model", "startTime"],
|
||||
where=filter_query,
|
||||
sum={
|
||||
"spend": True,
|
||||
},
|
||||
summary_sql_and_params: Final = _spend_logs_daily_summary_sql(
|
||||
start_date_iso=start_date_iso,
|
||||
end_date_iso=end_date_iso,
|
||||
api_key=summary_api_key,
|
||||
request_id=request_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
sql_query, params = summary_sql_and_params
|
||||
rows: Final[Sequence[_SpendDailySummaryRow]] = await _query_raw(prisma_client, sql_query, *params)
|
||||
if len(rows) == 0:
|
||||
return [] # pyright: ignore[reportUnknownVariableType] # empty summary has no element type
|
||||
|
||||
if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
|
||||
spend_rows: Final = cast(Sequence[_SpendGroupByRow], response) # cast-ok: by/sum fix the shape
|
||||
result: Final[dict] = {}
|
||||
for record in spend_rows:
|
||||
dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
|
||||
date = dt_object.date()
|
||||
if date not in result:
|
||||
result[date] = {"users": {}, "models": {}}
|
||||
api_key = record["api_key"]
|
||||
user_id = record["user"]
|
||||
model = record["model"]
|
||||
result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0)
|
||||
result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0)
|
||||
result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get(
|
||||
"spend", 0
|
||||
)
|
||||
result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get(
|
||||
"spend", 0
|
||||
)
|
||||
return_list: Final = []
|
||||
final_date = None
|
||||
for k, v in sorted(result.items()):
|
||||
return_list.append({**v, "startTime": k})
|
||||
final_date = k
|
||||
|
||||
end_date_date: Final = end_date_obj.date()
|
||||
if final_date is not None and final_date < end_date_date:
|
||||
current_date = final_date + timedelta(days=1)
|
||||
while current_date <= end_date_date:
|
||||
# Represent current_date as string because original response has it this way
|
||||
return_list.append(
|
||||
{
|
||||
"startTime": current_date,
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
) # If no data, will stay as zero
|
||||
current_date += timedelta(days=1) # Move on to the next day
|
||||
|
||||
return return_list
|
||||
|
||||
return response
|
||||
summary_items: Final = tuple(
|
||||
_daily_summary_item(date.fromisoformat(day), tuple(day_rows))
|
||||
for day, day_rows in groupby(rows, key=lambda row: row["day"])
|
||||
)
|
||||
final_date: Final = date.fromisoformat(rows[-1]["day"])
|
||||
end_date_date: Final = end_date_obj.date()
|
||||
padding: Final[tuple[Mapping[str, object], ...]] = tuple(
|
||||
{
|
||||
"startTime": final_date + timedelta(days=offset),
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
for offset in range(1, (end_date_date - final_date).days + 1)
|
||||
)
|
||||
return [*summary_items, *padding]
|
||||
|
||||
else:
|
||||
scoped_filter: Final[dict[str, str]] = {}
|
||||
|
|
|
|||
|
|
@ -96,7 +96,7 @@
|
|||
"limit": 10
|
||||
},
|
||||
"DTZ007": {
|
||||
"limit": 17
|
||||
"limit": 6
|
||||
},
|
||||
"DTZ011": {
|
||||
"limit": 3
|
||||
|
|
|
|||
|
|
@ -5,14 +5,12 @@ import hashlib
|
|||
import json
|
||||
import re
|
||||
from datetime import timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
|
|
@ -3325,7 +3323,7 @@ def _compare_nested_dicts(
|
|||
return differences
|
||||
|
||||
# Check for keys in actual but not in expected
|
||||
for key in actual.keys():
|
||||
for key in actual:
|
||||
current_path = f"{path}.{key}" if path else key
|
||||
if current_path not in ignore_keys and key not in expected:
|
||||
differences.append(f"Extra key in actual: {current_path}")
|
||||
|
|
@ -3495,24 +3493,22 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch):
|
|||
# Return individual log entries when summarize=false
|
||||
return mock_spend_logs
|
||||
|
||||
async def group_by(self, *args, **kwargs):
|
||||
# Return grouped data when summarize=true
|
||||
# Simplified mock response for grouped data
|
||||
async def query_raw(self, sql_query, *params):
|
||||
yesterday = datetime.datetime.now(timezone.utc) - timedelta(days=1)
|
||||
return [
|
||||
{
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"model": "gpt-3.5-turbo",
|
||||
"startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
"_sum": {"spend": 0.05},
|
||||
"day": yesterday.date().isoformat(),
|
||||
"spend": 0.05,
|
||||
},
|
||||
{
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"model": "gpt-4",
|
||||
"startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
"_sum": {"spend": 0.10},
|
||||
"day": yesterday.date().isoformat(),
|
||||
"spend": 0.10,
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -3850,47 +3846,30 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
"""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# This simulates the summarized data that Prisma's `group_by` would return.
|
||||
mock_summarized_response = [
|
||||
{
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"model": "gpt-4",
|
||||
"startTime": (datetime.now(timezone.utc) - timedelta(days=1)).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S.%fZ"
|
||||
),
|
||||
"_sum": {"spend": 0.15},
|
||||
"day": (datetime.now(timezone.utc) - timedelta(days=1)).date().isoformat(),
|
||||
"spend": 0.15,
|
||||
}
|
||||
]
|
||||
|
||||
# This mock class will replace the real Prisma client.
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = self
|
||||
|
||||
async def group_by(self, *args, **kwargs):
|
||||
# We assert that the `gte` and `lte` values are strings in ISO format.
|
||||
# If they were datetime objects, this test would fail.
|
||||
where_clause = kwargs.get("where", {})
|
||||
start_time_filter = where_clause.get("startTime", {})
|
||||
|
||||
assert "gte" in start_time_filter
|
||||
assert "lte" in start_time_filter
|
||||
assert isinstance(start_time_filter["gte"], str)
|
||||
assert isinstance(start_time_filter["lte"], str)
|
||||
assert "T" in start_time_filter["gte"] # Check for ISO format 'T' separator
|
||||
|
||||
# If the assertions pass, return the mock response.
|
||||
async def query_raw(self, sql_query, *params):
|
||||
assert isinstance(params[0], str)
|
||||
assert isinstance(params[1], str)
|
||||
assert "T" in params[0]
|
||||
assert "T" in params[1]
|
||||
return mock_summarized_response
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
# Apply the monkeypatch to replace the real prisma_client with our mock.
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
|
||||
|
||||
# Define a date range for the test.
|
||||
start_date = (datetime.now(timezone.utc) - timedelta(days=2)).strftime("%Y-%m-%d")
|
||||
end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
|
||||
|
|
@ -3898,8 +3877,6 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
# Call the endpoint with both start and end dates.
|
||||
# We don't need `summarize=true` as it's the default.
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
|
|
@ -3909,11 +3886,9 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
# ASSERTIONS
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check that the response is not empty and has the summarized structure.
|
||||
assert isinstance(data, list)
|
||||
assert len(data) > 0
|
||||
assert "startTime" in data[0]
|
||||
|
|
@ -3924,6 +3899,183 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatch):
|
||||
mock_rows = [
|
||||
{
|
||||
"day": "2024-01-01",
|
||||
"api_key": "hashed::sk-abc",
|
||||
"user": "u1",
|
||||
"model": "gpt-4",
|
||||
"spend": 0.1,
|
||||
},
|
||||
{
|
||||
"day": "2024-01-01",
|
||||
"api_key": "hashed::sk-abc",
|
||||
"user": "u1",
|
||||
"model": "gpt-4o",
|
||||
"spend": 0.2,
|
||||
},
|
||||
]
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.captured_sql = None
|
||||
self.captured_params = None
|
||||
|
||||
async def query_raw(self, sql_query, *params):
|
||||
self.captured_sql = sql_query
|
||||
self.captured_params = params
|
||||
return mock_rows
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
def hash_token(self, token):
|
||||
return "hashed::" + token
|
||||
|
||||
mock_prisma_client = MockPrismaClient()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
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-03",
|
||||
"api_key": "sk-abc",
|
||||
"request_id": "req-123",
|
||||
"user_id": "u1",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
sql = mock_prisma_client.db.captured_sql
|
||||
assert "date_trunc('day'" in sql
|
||||
assert "GROUP BY" in sql
|
||||
assert "find_many" not in sql
|
||||
assert not hasattr(mock_prisma_client.db, "group_by")
|
||||
assert mock_prisma_client.db.captured_params == (
|
||||
"2024-01-01T00:00:00+00:00",
|
||||
"2024-01-03T00:00:00+00:00",
|
||||
"hashed::sk-abc",
|
||||
"req-123",
|
||||
"u1",
|
||||
)
|
||||
assert len(data) == 3
|
||||
assert data[0]["startTime"] == "2024-01-01"
|
||||
assert data[0]["spend"] == pytest.approx(0.3)
|
||||
assert data[0]["models"] == {"gpt-4": 0.1, "gpt-4o": 0.2}
|
||||
assert data[0]["users"] == {"u1": pytest.approx(0.3)}
|
||||
assert data[0]["hashed::sk-abc"] == pytest.approx(0.3)
|
||||
assert data[1] == {
|
||||
"startTime": "2024-01-02",
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
assert data[2] == {
|
||||
"startTime": "2024-01-03",
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch):
|
||||
class MockDB:
|
||||
async def query_raw(self, sql_query, *params):
|
||||
return []
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
|
||||
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-01"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == []
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client, monkeypatch):
|
||||
mock_rows = [
|
||||
{
|
||||
"day": "2024-01-01",
|
||||
"api_key": "plain-key",
|
||||
"user": "u1",
|
||||
"model": "gpt-4",
|
||||
"spend": 0.4,
|
||||
}
|
||||
]
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.captured_params = None
|
||||
|
||||
async def query_raw(self, sql_query, *params):
|
||||
self.captured_params = params
|
||||
return mock_rows
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
mock_prisma_client = MockPrismaClient()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
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-01",
|
||||
"api_key": "plain-key",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert mock_prisma_client.db.captured_params == (
|
||||
"2024-01-01T00:00:00+00:00",
|
||||
"2024-01-01T00:00:00+00:00",
|
||||
"plain-key",
|
||||
)
|
||||
assert data == [
|
||||
{
|
||||
"startTime": "2024-01-01",
|
||||
"spend": pytest.approx(0.4),
|
||||
"plain-key": pytest.approx(0.4),
|
||||
"users": {"u1": pytest.approx(0.4)},
|
||||
"models": {"gpt-4": pytest.approx(0.4)},
|
||||
}
|
||||
]
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_with_error_code(client):
|
||||
"""Test filtering spend logs by error code"""
|
||||
|
|
@ -4832,13 +4984,14 @@ class _CaptureFilterDB:
|
|||
def __init__(self):
|
||||
self.litellm_spendlogs = self
|
||||
self.captured_where = None
|
||||
self.captured_params = None
|
||||
|
||||
async def find_many(self, *args, **kwargs):
|
||||
self.captured_where = kwargs.get("where")
|
||||
return []
|
||||
|
||||
async def group_by(self, *args, **kwargs):
|
||||
self.captured_where = kwargs.get("where")
|
||||
async def query_raw(self, sql_query, *params):
|
||||
self.captured_params = params
|
||||
return []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 22328
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26758
|
||||
"limit": 26750
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 261
|
||||
|
|
@ -27,10 +27,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16470
|
||||
"limit": 16468
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5516
|
||||
"limit": 5514
|
||||
},
|
||||
"LIT012": {
|
||||
"limit": 4489
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue