scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(usage): fall back only on database outages, on every daily activity route

Browse source 
Data-layer Prisma errors such as a rejected query or a missing table return 500 again, a key whose details failed to load reports key_exists null instead of false, and the non-aggregated daily activity routes keep spend when the key or entity metadata lookup hits an outage
This commit is contained in:
mateo-berri 2026-11-06 06:48:23 -08:00
parent ee95ac14a5
commit 2c32b2d554
2 changed files with 124 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -191,9 +191,11 @@ class _AggregatedSpendData(TypedDict):
totals: ReadOnly[SpendMetrics]
def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str | None) -> KeyMetadata:
def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow] | None, api_key: str | None) -> KeyMetadata:
if not api_key:
return KeyMetadata(key_alias="Unattributed", key_exists=False)
if api_key_metadata is None:
return KeyMetadata()
meta: Final = api_key_metadata.get(api_key)
return KeyMetadata(
key_alias=meta.key_alias if meta is not None else None,
@ -309,7 +311,7 @@ def update_breakdown_metrics(
record: DailySpendRecord,
model_metadata: Mapping[str, dict[str, object]],
provider_metadata: Mapping[str, dict[str, object]],
api_key_metadata: Mapping[str, KeyMetadataRow],
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
entity_id_field: str | None = None,
entity_metadata_field: Mapping[str, dict[str, object]] | None = None,
) -> BreakdownMetrics:
@ -629,7 +631,7 @@ async def get_api_key_metadata(
def _aggregate_spend_records_sync(
*,
records: Sequence[DailySpendRecord],
api_key_metadata: Mapping[str, KeyMetadataRow],
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
entity_id_field: str | None,
entity_metadata_field: Mapping[str, dict[str, object]] | None,
) -> _AggregatedSpendData:
@ -688,22 +690,10 @@ async def _aggregate_spend_records(
The per-row loop is offloaded to a worker thread via asyncio.to_thread so
a large result set doesn't peg the event loop.
"""
api_keys: Final[set[str]] = {
record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY
}
api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = (
await repository.key_metadata(
frozenset(api_keys), spend_logs_window(frozenset(record.date for record in records))
)
if api_keys
else MappingProxyType({})
)
return await asyncio.to_thread(
_aggregate_spend_records_sync,
records=records,
api_key_metadata=api_key_metadata,
api_key_metadata=await _key_metadata_unless_database_fails(repository, records),
entity_id_field=entity_id_field,
entity_metadata_field=entity_metadata_field,
)
@ -786,7 +776,7 @@ def _record_to_spend_metrics(record: RollupMetricsRow) -> SpendMetrics:
def _aggregate_grouping_sets_records_sync(
*,
records: Sequence[GroupingSetsRow],
api_key_metadata: Mapping[str, KeyMetadataRow],
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
) -> _AggregatedSpendData:
"""Build the response from rollup rows produced by the GROUPING SETS query.
@ -926,22 +916,20 @@ async def unless_database_fails(lookup: Awaitable[_T], fallback: _T, left_out: s
try:
return await lookup
except Exception as exc:
if not (
PrismaDBExceptionHandler.is_prisma_error(exc)
or PrismaDBExceptionHandler.is_database_service_unavailable_error(exc)
):
if not PrismaDBExceptionHandler.is_database_service_unavailable_error(exc):
raise
verbose_proxy_logger.warning("Returning daily activity without %s: %s", left_out, exc)
return fallback
async def _key_metadata_unless_database_fails(
repository: DailyActivityRepository, rows: Sequence[GroupingSetsRow] | Sequence[EntityRollupRow]
) -> Mapping[str, KeyMetadataRow]:
repository: DailyActivityRepository,
rows: Sequence[GroupingSetsRow] | Sequence[EntityRollupRow] | Sequence[DailySpendRecord],
) -> Mapping[str, KeyMetadataRow] | None:
api_keys: Final = frozenset(row.api_key for row in rows if row.api_key and row.api_key != PTU_SENTINEL_API_KEY)
return await unless_database_fails(
repository.key_metadata(api_keys, spend_logs_window(frozenset(row.date for row in rows))),
MappingProxyType({}),
None,
"its key metadata",
)
@ -1010,7 +998,13 @@ async def get_daily_activity(
if resolve_entity_metadata is not None:
resolved_entity_metadata = {
**(entity_metadata_field or {}),
**(await resolve_entity_metadata(daily_spend_data)),
**(
await unless_database_fails(
resolve_entity_metadata(daily_spend_data),
MappingProxyType[str, dict[str, object]]({}),
"its entity metadata",
)
),
}
aggregated: Final = await _aggregate_spend_records(
repository=repository,
@ -1057,7 +1051,7 @@ def _fold_entity_rollups_sync(
*,
results: Sequence[DailySpendData],
entity_rows: Sequence[EntityRollupRow],
api_key_metadata: Mapping[str, KeyMetadataRow],
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
entity_metadata_field: Mapping[str, Mapping[str, object]] | None,
) -> None:
"""Write breakdown.entities onto the already-built per-day results."""

131
tests/unit/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -1,15 +1,14 @@
from collections.abc import Mapping, Sequence
from dataclasses import asdict, fields
from datetime import date, datetime
from types import SimpleNamespace
from types import FunctionType, SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from fastapi import HTTPException
from prisma.engine.utils import handle_response_errors
from prisma.errors import DataError, PrismaError
from prisma.errors import DataError, PrismaError, TableNotFoundError
import litellm.proxy.management_endpoints.common_daily_activity as common_daily_activity_module
from litellm.constants import USAGE_TOP_API_KEYS_DEFAULT
@ -49,22 +48,16 @@ from litellm.types.repositories.daily_activity import (
)
def _raise_missing_error_meta() -> None:
raise AttributeError("'NoneType' object has no attribute 'get'")
def _query_engine_crash_on_closed_connection() -> AttributeError:
raise_inside_query_engine: Final = FunctionType(
_raise_missing_error_meta.__code__, {"__name__": "prisma.engine.utils"}
)
with pytest.raises(AttributeError) as crash:
handle_response_errors(
MagicMock(),
[
{
"error": "Server has closed the connection.",
"user_facing_error": {
"is_panic": False,
"message": "Server has closed the connection.",
"meta": None,
"error_code": "P1017",
},
}
],
)
raise_inside_query_engine()
return crash.value
@ -74,7 +67,6 @@ DATABASE_FAILURES: Final = pytest.mark.parametrize(
DataError(
{"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server at `db:5432`"}}
),
DataError({"user_facing_error": {"error_code": "P2010", "message": "Raw query failed"}}),
PrismaError("metadata unavailable"),
httpx.ReadTimeout("the query engine stopped answering"),
httpx.ConnectError("the query engine refused the connection"),
@ -82,7 +74,6 @@ DATABASE_FAILURES: Final = pytest.mark.parametrize(
],
ids=[
"unreachable-database",
"rejected-query",
"prisma-error",
"engine-read-timeout",
"engine-connect-error",
@ -90,6 +81,25 @@ DATABASE_FAILURES: Final = pytest.mark.parametrize(
],
)
NOT_DATABASE_OUTAGES: Final = pytest.mark.parametrize(
"bug",
[
TypeError("not a database failure"),
AttributeError("not a database failure"),
DataError({"user_facing_error": {"error_code": "P2010", "message": "not a database failure"}}),
TableNotFoundError(
{
"user_facing_error": {
"error_code": "P2021",
"message": "not a database failure",
"meta": {"table": "LiteLLM_VerificationToken"},
}
}
),
],
ids=["type-error", "attribute-error", "rejected-query", "missing-table"],
)
@pytest.mark.asyncio
@DATABASE_FAILURES
@ -141,8 +151,8 @@ async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata(
{} if user_failure or empty else {"user_email": "user@example.com", "user_alias": "User"}
)
assert entity.api_key_breakdown["key"].metrics.spend == 2.0
assert entity.api_key_breakdown["key"].metadata == KeyMetadata(
key_alias=None if key_failure or empty else "Alias", key_exists=not (key_failure or empty)
assert entity.api_key_breakdown["key"].metadata == (
KeyMetadata() if key_failure else KeyMetadata(key_alias=None if empty else "Alias", key_exists=not empty)
)
@ -192,7 +202,7 @@ async def test_usage_keeps_spend_when_the_first_key_lookup_hits_a_database_failu
day: Final = result.results[0]
assert result.metadata.total_spend == day.metrics.spend == 2.0
assert day.breakdown.api_keys["key"].metrics.spend == 2.0
assert day.breakdown.api_keys["key"].metadata == KeyMetadata(key_alias=None, key_exists=False)
assert day.breakdown.api_keys["key"].metadata == KeyMetadata()
@pytest.mark.asyncio
@ -209,12 +219,10 @@ async def test_usage_keeps_spend_when_the_guardrail_lookup_hits_a_database_failu
@pytest.mark.asyncio
@pytest.mark.parametrize("lookup", ["key_metadata", "user_metadata", "guardrail_spend"])
@pytest.mark.parametrize("bug", [TypeError, AttributeError])
async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(
lookup: str, bug: type[Exception]
) -> None:
@NOT_DATABASE_OUTAGES
async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(lookup: str, bug: Exception) -> None:
repository: Final = _usage_repository()
setattr(repository, lookup, AsyncMock(side_effect=bug("not a database failure")))
setattr(repository, lookup, AsyncMock(side_effect=bug))
with pytest.raises(HTTPException) as error:
await _get_daily_activity_aggregated(
@ -225,6 +233,75 @@ async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(
assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"}
def _daily_user_spend_prisma(key_lookup: AsyncMock) -> MagicMock:
prisma: Final = MagicMock()
records: Final = [_daily_user_spend_record(user_id="user", api_key="key", spend=7.0)]
prisma.db.litellm_dailyuserspend.count = AsyncMock(return_value=len(records))
prisma.db.litellm_dailyuserspend.find_many = AsyncMock(return_value=records)
prisma.db.litellm_verificationtoken.find_many = key_lookup
prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
return prisma
async def _daily_user_activity(prisma: MagicMock, user_emails: AsyncMock) -> SpendAnalyticsPaginatedResponse:
return await get_daily_activity(
prisma_client=prisma,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
entity_metadata_field=None,
start_date="2024-01-01",
end_date="2024-01-01",
model=None,
api_key=None,
page=1,
page_size=1000,
resolve_entity_metadata=user_emails,
)
@pytest.mark.asyncio
@DATABASE_FAILURES
@pytest.mark.parametrize(("key_failure", "email_failure"), [(True, False), (False, True)])
async def test_daily_activity_keeps_spend_when_a_metadata_lookup_hits_a_database_failure(
failure: Exception, key_failure: bool, email_failure: bool
) -> None:
key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None)
prisma: Final = _daily_user_spend_prisma(
AsyncMock(side_effect=failure if key_failure else None, return_value=[key])
)
user_emails: Final = AsyncMock(
side_effect=failure if email_failure else None, return_value={"user": {"user_email": "user@example.com"}}
)
result: Final = await _daily_user_activity(prisma, user_emails)
day: Final = result.results[0]
assert result.metadata.total_spend == day.metrics.spend == 7.0
assert day.breakdown.api_keys["key"].metadata == (
KeyMetadata() if key_failure else KeyMetadata(key_alias="Alias", key_exists=True)
)
assert day.breakdown.entities["user"].metadata == ({} if email_failure else {"user_email": "user@example.com"})
@pytest.mark.asyncio
@pytest.mark.parametrize("lookup", ["key metadata", "user emails"])
@NOT_DATABASE_OUTAGES
async def test_daily_activity_still_reports_a_metadata_lookup_bug_that_is_not_a_database_failure(
lookup: str, bug: Exception
) -> None:
prisma: Final = _daily_user_spend_prisma(
AsyncMock(side_effect=bug if lookup == "key metadata" else None, return_value=[])
)
user_emails: Final = AsyncMock(side_effect=bug if lookup == "user emails" else None, return_value={})
with pytest.raises(HTTPException) as error:
await _daily_user_activity(prisma, user_emails)
assert error.value.status_code == 500
assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"}
async def _run_aggregated_daily_activity(
*,
prisma_client: PrismaClient,