fix(guardrails): degrade usage units to empty when the units table is missing

GET /guardrails/usage/overview and GET /guardrails/usage/detail/{id} 500ed on a
database that has not applied 20260817143646_add_daily_guardrail_usage_units yet
(pip installs on litellm-proxy-extras 0.4.86 with DISABLE_SCHEMA_UPDATE=true).
Both endpoints now return their metrics with empty units and log one warning
until the migration lands.
This commit is contained in:
mateo-berri 2026-08-17 19:38:43 -07:00
parent 35176fa64a
commit 15823b1be3
2 changed files with 50 additions and 1 deletions

View file

@ -14,6 +14,7 @@ from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.table_repositories import (
@ -111,7 +112,16 @@ async def _find_daily_guardrail_usage_units(
prisma_client: "PrismaClient",
where: "prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput",
) -> "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where)
from prisma.errors import TableNotFoundError
try:
return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where)
except TableNotFoundError as e:
verbose_proxy_logger.warning(
"Guardrail usage units are unavailable until the LiteLLM_DailyGuardrailUsageUnits migration is applied: %s",
e,
)
return ()
def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str:

View file

@ -19,6 +19,7 @@ import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from fastapi import HTTPException
from prisma.errors import TableNotFoundError
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
@ -295,6 +296,44 @@ async def test_detail_breaks_units_down_by_day_team_and_key():
assert units_where == {"guardrail_id": {"in": ["yaml-pii", "yaml-1"]}, "date": {"gte": START, "lte": END}}
def _units_table_missing() -> TableNotFoundError:
return TableNotFoundError(
data={"user_facing_error": {"meta": {"table": "public.LiteLLM_DailyGuardrailUsageUnits"}}}
)
@pytest.mark.asyncio
async def test_overview_degrades_units_to_empty_when_units_table_is_missing():
prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
row = next(r for r in resp.rows if r.id == "yaml-uuid")
assert (row.requestsEvaluated, row.usageUnits) == (4, {})
assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {})
@pytest.mark.asyncio
async def test_detail_degrades_units_to_empty_when_units_table_is_missing():
prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
handler = _config_handler(_yaml_guardrail())
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_detail(
guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
)
assert (resp.requestsEvaluated, resp.failRate) == (4, 25.0)
assert (resp.usage_units, list(resp.usage_units_daily), resp.usage_units_by_team, resp.usage_units_by_key) == (
{},
[],
{},
{},
)
# ---- logs -------------------------------------------------------------------