diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ba131ecac4c..1a9086ccee7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5681 }, "reportMissingTypeArgument": { - "limit": 15609 + "limit": 15608 }, "reportMissingTypeStubs": { "limit": 40 diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index c25cf07b4d4..727f767469f 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -3,14 +3,16 @@ Track guardrail and policy usage for the dashboard: upsert daily metrics and insert into SpendLogGuardrailIndex when spend logs are written. """ +import asyncio import json from collections import defaultdict -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone +from functools import partial from itertools import groupby from operator import itemgetter from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple +from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient @@ -24,6 +26,12 @@ if TYPE_CHECKING: from prisma import types as prisma_types +_UPSERT_RETRY_TIMES: Final = 3 + +_RowKey = TypeVar("_RowKey") +_RowValue = TypeVar("_RowValue") + + class _UsageUnitKey(NamedTuple): guardrail_id: str date: str @@ -32,6 +40,46 @@ class _UsageUnitKey(NamedTuple): usage_unit: str +class _MetricsKey(NamedTuple): + guardrail_id: str + date: str + + +async def _attempt_upsert( + upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], key: _RowKey, value: _RowValue +) -> Exception | None: + try: + await upsert_row(key, value) + except Exception as error: + return error + return None + + +async def _upsert_rows_with_retry( + rows: Mapping[_RowKey, _RowValue], + upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], + label: str, + sleep: Callable[[float], Awaitable[None]], + retries_left: int = _UPSERT_RETRY_TIMES, +) -> None: + outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()} + failed: Final = MappingProxyType({key: rows[key] for key, error in outcomes.items() if error is not None}) + if not failed: + return + if retries_left == 0: + for key in failed: + verbose_proxy_logger.warning( + "Guardrail usage tracking: %s upsert failed for %s after %d retries (non-fatal): %s", + label, + key, + _UPSERT_RETRY_TIMES, + outcomes[key], + ) + return + await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left)) + await _upsert_rows_with_retry(failed, upsert_row, label, sleep, retries_left - 1) + + def _guardrail_status_to_action(status: str | None) -> str: """Map StandardLogging guardrail_status to blocked/passed/flagged.""" if not status: @@ -131,9 +179,33 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) +async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None: + n: Final = int(agg["requests_evaluated"]) + await DailyGuardrailMetricsRepository(prisma_client).table.upsert( + where={"guardrail_id_date": {"guardrail_id": key.guardrail_id, "date": key.date}}, + data={ + "create": { + "guardrail_id": key.guardrail_id, + "date": key.date, + "requests_evaluated": n, + "passed_count": int(agg["passed_count"]), + "blocked_count": int(agg["blocked_count"]), + "flagged_count": int(agg["flagged_count"]), + }, + "update": { + "requests_evaluated": {"increment": n}, + "passed_count": {"increment": int(agg["passed_count"])}, + "blocked_count": {"increment": int(agg["blocked_count"])}, + "flagged_count": {"increment": int(agg["flagged_count"])}, + }, + }, + ) + + async def process_spend_logs_guardrail_usage( prisma_client: PrismaClient, logs_to_process: list[dict[str, Any]], + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, ) -> None: """ After spend logs are written: update DailyGuardrailMetrics and insert @@ -142,7 +214,7 @@ async def process_spend_logs_guardrail_usage( if not logs_to_process: return # Aggregate daily metrics by (guardrail_id, date). Latency/score metrics dropped. - daily_guardrail: Final[dict[tuple, dict[str, Any]]] = defaultdict( + daily_guardrail: Final[dict[_MetricsKey, dict[str, Any]]] = defaultdict( lambda: { "requests_evaluated": 0, "passed_count": 0, @@ -163,7 +235,7 @@ async def process_spend_logs_guardrail_usage( guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" if not guardrail_id: continue - key = (guardrail_id, date_key) + key = _MetricsKey(guardrail_id, date_key) daily_guardrail[key]["requests_evaluated"] += 1 action = _guardrail_status_to_action(entry.get("guardrail_status")) if action == "passed": @@ -199,51 +271,12 @@ async def process_spend_logs_guardrail_usage( verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e) # Upsert daily guardrail metrics (counts only; latency/score dropped) - for (guardrail_id, date_key), agg in daily_guardrail.items(): - n = int(agg["requests_evaluated"]) - if n == 0: - continue - try: - await DailyGuardrailMetricsRepository(prisma_client).table.upsert( - where={ - "guardrail_id_date": { - "guardrail_id": guardrail_id, - "date": date_key, - } - }, - data={ - "create": { - "guardrail_id": guardrail_id, - "date": date_key, - "requests_evaluated": n, - "passed_count": int(agg["passed_count"]), - "blocked_count": int(agg["blocked_count"]), - "flagged_count": int(agg["flagged_count"]), - }, - "update": { - "requests_evaluated": {"increment": n}, - "passed_count": {"increment": int(agg["passed_count"])}, - "blocked_count": {"increment": int(agg["blocked_count"])}, - "flagged_count": {"increment": int(agg["flagged_count"])}, - }, - }, - ) - except Exception as metrics_error: - verbose_proxy_logger.warning( - "Guardrail usage tracking: daily metrics upsert failed for %s on %s (non-fatal): %s", - guardrail_id, - date_key, - metrics_error, - ) - - for unit_key, units in usage_unit_totals.items(): - try: - await _upsert_usage_unit_row(prisma_client, unit_key, units) - except Exception as unit_error: - verbose_proxy_logger.warning( - "Guardrail usage tracking: usage unit upsert failed for %s (non-fatal): %s", - unit_key, - unit_error, - ) + metrics_rows: Final = MappingProxyType( + {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0} + ) + await _upsert_rows_with_retry(metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep) + await _upsert_rows_with_retry( + usage_unit_totals, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep + ) except Exception as e: verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 15396a95632..ffbb8dc2332 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -78,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 313 + "limit": 312 }, "D419": { "limit": 6 diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 4f606d489d5..62a4bbbbe6e 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -80,24 +80,70 @@ async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): } +def _fake_sleep() -> tuple[AsyncMock, list[float]]: + delays: list[float] = [] + sleep = AsyncMock(side_effect=lambda delay: delays.append(delay)) + return sleep, delays + + @pytest.mark.asyncio async def test_one_failing_upsert_does_not_drop_remaining_writes(): """ A DB error on one daily-metrics or usage-unit upsert must not cancel the remaining upserts in the flushed batch, or the usage endpoints would - permanently under-report billable counters (batches are never retried). + permanently under-report billable counters. """ prisma = _prisma() prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") - prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None] + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None, None] + sleep, _ = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), ] - await process_spend_logs_guardrail_usage(prisma, logs) + await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) - assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 2 + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, + ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1, + } + + +@pytest.mark.asyncio +async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only(): + """ + A transient DB error must not permanently drop billed units from the + aggregates: only the rows that failed are re-sent, after exponential + backoff, and the batch ends once every row has landed. + """ + prisma = _prisma() + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("blip"), None, None] + sleep, delays = _fake_sleep() + logs = [ + _payload("r1", usage={"topicPolicyUnits": 1}), + _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) + + calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list + assert len(calls) == 3 + assert calls[2].kwargs["where"] == calls[0].kwargs["where"] + assert delays == [1] + + +@pytest.mark.asyncio +async def test_persistent_upsert_failure_stops_after_three_retries(): + prisma = _prisma() + prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") + sleep, delays = _fake_sleep() + + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep) + + assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4 + assert delays == [1, 2, 4] + assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1 @pytest.mark.asyncio diff --git a/type-discipline-budget.json b/type-discipline-budget.json index d0ec219b85c..3d2923c39ac 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22903 + "limit": 22900 }, "LIT002": { - "limit": 26894 + "limit": 26893 }, "LIT003": { "limit": 269