mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
35176fa64a
commit
15823b1be3
2 changed files with 50 additions and 1 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 -------------------------------------------------------------------
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue