scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(guardrails): show YAML-defined guardrails in Guardrail Monitor

Browse source 
The /guardrails/usage/{overview,detail,logs} endpoints only consulted the
litellm_guardrailstable Prisma table, so guardrails defined in config.yaml
(stored only in IN_MEMORY_GUARDRAIL_HANDLER) were invisible: detail 404'd,
overview rendered them as "Custom"/"Guardrail" via the orphan-metric path or
not at all, and logs missed the logical-name alias.

Mirror the existing list_guardrails_v2 / get_guardrail_info pattern: union DB
rows with IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails() (deduped by
guardrail_id) for overview, fall back to get_guardrail_by_id() for detail and
logs. Adds tests covering DB-only, YAML-only, both, and DB-takes-precedence.

Fixes LIT-2529.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuneng Jiang 2026-04-27 15:22:13 -07:00
parent 82dacfb746
commit 44923f260c
2 changed files with 354 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
litellm/proxy/guardrails/usage_endpoints.py
View file

@ -12,6 +12,8 @@ from pydantic import BaseModel
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.types.guardrails import LitellmParams
router = APIRouter()
@ -146,6 +148,22 @@ def _get_guardrail_attrs(g: Any) -> tuple[Any, str]:
return gid, (name or gid or "")
def _get_guardrail_field(g: Any, field: str) -> Any:
"""Read `field` off a guardrail (Prisma row attr or dict/TypedDict key)."""
if isinstance(g, dict):
return g.get(field)
return getattr(g, field, None)
def _to_dict(value: Any) -> Dict[str, Any]:
"""Coerce a LitellmParams / dict / None into a plain dict."""
if isinstance(value, LitellmParams):
return value.model_dump(exclude_none=True)
if isinstance(value, dict):
return value
return {}
def _guardrail_overview_rows(
guardrails: Any,
agg: Dict[str, Dict[str, Any]],
@ -165,13 +183,9 @@ def _guardrail_overview_rows(
break
req, blocked = a["requests"], a["blocked"]
fail_rate = (100.0 * blocked / req) if req else 0.0
litellm_params = (
(g.litellm_params or {}) if isinstance(g.litellm_params, dict) else {}
)
litellm_params = _to_dict(_get_guardrail_field(g, "litellm_params"))
provider = str(litellm_params.get("guardrail", "Unknown"))
guardrail_info = (
(g.guardrail_info or {}) if isinstance(g.guardrail_info, dict) else {}
)
guardrail_info = _to_dict(_get_guardrail_field(g, "guardrail_info"))
gtype = str(guardrail_info.get("type", "Guardrail"))
prev_fail = 0.0
for k in lookup_keys:
@ -271,8 +285,19 @@ async def guardrails_usage_overview(
start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
try:
# Guardrails from DB
guardrails = await prisma_client.db.litellm_guardrailstable.find_many()
# Guardrails from DB unioned with YAML/in-memory guardrails (deduped by id).
db_guardrails = await prisma_client.db.litellm_guardrailstable.find_many()
seen_ids = {
getattr(g, "guardrail_id", None)
for g in db_guardrails
if getattr(g, "guardrail_id", None)
}
in_memory_guardrails = [
g
for g in IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails()
if g.get("guardrail_id") not in seen_ids
]
guardrails: List[Any] = list(db_guardrails) + in_memory_guardrails
# Daily metrics in range
metrics = await prisma_client.db.litellm_dailyguardrailmetrics.find_many(
@ -338,15 +363,18 @@ async def guardrails_usage_detail(
guardrail = await prisma_client.db.litellm_guardrailstable.find_unique(
where={"guardrail_id": guardrail_id}
)
if not guardrail:
if guardrail is None:
# YAML-defined guardrails live only in the in-memory registry.
guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(
guardrail_id=guardrail_id
)
if guardrail is None:
from fastapi import HTTPException
raise HTTPException(status_code=404, detail="Guardrail not found")
# Metrics are keyed by logical name (from spend log metadata), not UUID
logical_id = getattr(guardrail, "guardrail_name", None) or (
guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None
)
logical_id = _get_guardrail_field(guardrail, "guardrail_name")
metric_ids = [i for i in (logical_id, guardrail_id) if i]
metrics = await prisma_client.db.litellm_dailyguardrailmetrics.find_many(
@ -383,17 +411,9 @@ async def guardrails_usage_detail(
{"date": d, "passed": v["passed"], "blocked": v["blocked"], "score": None}
for d, v in sorted(ts_by_date.items())
]
_litellm_params = getattr(guardrail, "litellm_params", None) or (
guardrail.get("litellm_params") if isinstance(guardrail, dict) else None
)
litellm_params = _litellm_params if isinstance(_litellm_params, dict) else {}
_guardrail_info = getattr(guardrail, "guardrail_info", None) or (
guardrail.get("guardrail_info") if isinstance(guardrail, dict) else None
)
guardrail_info = _guardrail_info if isinstance(_guardrail_info, dict) else {}
_guardrail_name = getattr(guardrail, "guardrail_name", None) or (
guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None
)
litellm_params = _to_dict(_get_guardrail_field(guardrail, "litellm_params"))
guardrail_info = _to_dict(_get_guardrail_field(guardrail, "guardrail_info"))
_guardrail_name = _get_guardrail_field(guardrail, "guardrail_name")
return UsageDetailResponse(
guardrail_id=guardrail_id,
@ -577,8 +597,12 @@ async def guardrails_usage_logs(
guardrail = await prisma_client.db.litellm_guardrailstable.find_unique(
where={"guardrail_id": guardrail_id}
)
if guardrail is None:
guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(
guardrail_id=guardrail_id
)
if guardrail:
logical_name = getattr(guardrail, "guardrail_name", None)
logical_name = _get_guardrail_field(guardrail, "guardrail_name")
if logical_name and logical_name not in effective_guardrail_ids:
effective_guardrail_ids.append(logical_name)

306
tests/test_litellm/proxy/guardrails/test_usage_endpoints.py Normal file
View file

@ -0,0 +1,306 @@
"""
Tests for /guardrails/usage/* endpoints serving the dashboard Guardrail Monitor.
Regression: LIT-2529 — guardrails defined in YAML (config) were invisible to
/guardrails/usage/detail and only partially visible to /guardrails/usage/overview.
"""
import os
import sys
from datetime import datetime
from typing import Optional
from unittest.mock import AsyncMock
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.guardrails.usage_endpoints import (
guardrails_usage_detail,
guardrails_usage_logs,
guardrails_usage_overview,
)
from litellm.types.guardrails import Guardrail, LitellmParams
ADMIN_AUTH = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
# When calling FastAPI handlers directly (not via TestClient), Query() defaults
# don't resolve to None, so explicit date strings are required.
START_DATE = "2026-04-20"
END_DATE = "2026-04-27"
def _make_yaml_guardrail(
guardrail_id: str = "yaml-gr-1",
guardrail_name: str = "my-yaml-pii",
provider: str = "presidio",
info: Optional[dict] = None,
) -> Guardrail:
"""Build a Guardrail TypedDict matching what InMemoryGuardrailHandler stores."""
return Guardrail(
guardrail_id=guardrail_id,
guardrail_name=guardrail_name,
litellm_params=LitellmParams(guardrail=provider, mode="pre_call"),
guardrail_info=info or {"type": "PII", "description": "YAML-defined"},
)
def _make_db_guardrail(
guardrail_id: str = "db-gr-1",
guardrail_name: str = "db-aim",
provider: str = "aim",
):
"""Build a mock Prisma row (object with attribute access)."""
class _Row:
pass
row = _Row()
row.guardrail_id = guardrail_id
row.guardrail_name = guardrail_name
row.litellm_params = {"guardrail": provider, "mode": "pre_call"}
row.guardrail_info = {"type": "ContentSafety", "description": "DB-defined"}
row.created_at = datetime.now()
row.updated_at = datetime.now()
return row
def _make_metric_row(
guardrail_id: str,
date: str = "2026-04-25",
requests: int = 10,
passed: int = 8,
blocked: int = 2,
flagged: int = 0,
):
class _M:
pass
m = _M()
m.guardrail_id = guardrail_id
m.date = date
m.requests_evaluated = requests
m.passed_count = passed
m.blocked_count = blocked
m.flagged_count = flagged
return m
@pytest.fixture
def mock_prisma(mocker):
client = mocker.Mock()
client.db = mocker.Mock()
client.db.litellm_guardrailstable = mocker.Mock()
client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
client.db.litellm_dailyguardrailmetrics = mocker.Mock()
client.db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=[])
client.db.litellm_spendlogguardrailindex = mocker.Mock()
client.db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=[])
client.db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0)
client.db.litellm_spendlogs = mocker.Mock()
client.db.litellm_spendlogs.find_many = AsyncMock(return_value=[])
mocker.patch("litellm.proxy.proxy_server.prisma_client", client)
return client
@pytest.fixture
def mock_in_memory(mocker):
handler = mocker.Mock()
handler.list_in_memory_guardrails.return_value = []
handler.get_guardrail_by_id.return_value = None
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
handler,
)
# Also patch the symbol re-imported into usage_endpoints (after the fix lands)
mocker.patch(
"litellm.proxy.guardrails.usage_endpoints.IN_MEMORY_GUARDRAIL_HANDLER",
handler,
create=True,
)
return handler
# ---- /guardrails/usage/detail ----------------------------------------------
@pytest.mark.asyncio
async def test_usage_detail_returns_yaml_guardrail(mock_prisma, mock_in_memory):
"""YAML guardrail (DB miss, in-memory hit) should return 200 with details."""
yaml_gr = _make_yaml_guardrail()
mock_in_memory.get_guardrail_by_id.return_value = yaml_gr
response = await guardrails_usage_detail(
guardrail_id="yaml-gr-1", user_api_key_dict=ADMIN_AUTH
)
assert response.guardrail_id == "yaml-gr-1"
assert response.guardrail_name == "my-yaml-pii"
assert response.provider == "presidio"
assert response.type == "PII"
assert response.description == "YAML-defined"
@pytest.mark.asyncio
async def test_usage_detail_db_takes_precedence_over_yaml(mock_prisma, mock_in_memory):
"""When both DB and in-memory have the same id, DB wins (parity with get_guardrail_info)."""
db_row = _make_db_guardrail(guardrail_id="shared-id", provider="bedrock")
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=db_row)
mock_in_memory.get_guardrail_by_id.return_value = _make_yaml_guardrail(
guardrail_id="shared-id", provider="presidio"
)
response = await guardrails_usage_detail(
guardrail_id="shared-id", user_api_key_dict=ADMIN_AUTH
)
assert response.provider == "bedrock"
# In-memory handler should not be consulted when DB lookup hits
mock_in_memory.get_guardrail_by_id.assert_not_called()
@pytest.mark.asyncio
async def test_usage_detail_404_when_neither_db_nor_yaml(mock_prisma, mock_in_memory):
"""If both miss, raise 404."""
with pytest.raises(HTTPException) as exc:
await guardrails_usage_detail(
guardrail_id="ghost", user_api_key_dict=ADMIN_AUTH
)
assert exc.value.status_code == 404
# ---- /guardrails/usage/overview --------------------------------------------
@pytest.mark.asyncio
async def test_usage_overview_includes_yaml_guardrails_with_no_metrics(
mock_prisma, mock_in_memory
):
"""YAML guardrails must appear in overview rows even with zero metrics."""
yaml_gr = _make_yaml_guardrail(
guardrail_id="yaml-gr-1",
guardrail_name="my-yaml-pii",
provider="presidio",
info={"type": "PII", "description": "YAML-defined"},
)
mock_in_memory.list_in_memory_guardrails.return_value = [yaml_gr]
response = await guardrails_usage_overview(
start_date=START_DATE, end_date=END_DATE, user_api_key_dict=ADMIN_AUTH
)
matching = [r for r in response.rows if r.id == "yaml-gr-1"]
assert len(matching) == 1, f"Expected one row for yaml-gr-1, got {response.rows}"
row = matching[0]
assert row.name == "my-yaml-pii"
assert row.provider == "presidio"
assert row.type == "PII"
assert row.requestsEvaluated == 0
@pytest.mark.asyncio
async def test_usage_overview_yaml_guardrail_metric_lookup_by_logical_name(
mock_prisma, mock_in_memory
):
"""Metrics are keyed by logical name; row must pick them up via name lookup."""
yaml_gr = _make_yaml_guardrail(
guardrail_id="yaml-uuid-xyz", guardrail_name="my-yaml-pii"
)
mock_in_memory.list_in_memory_guardrails.return_value = [yaml_gr]
# metric row keyed by the logical name (how the spend-log writer keys YAML guardrails)
mock_prisma.db.litellm_dailyguardrailmetrics.find_many = AsyncMock(
side_effect=[
[_make_metric_row("my-yaml-pii", requests=10, blocked=2)], # current
[], # previous
]
)
response = await guardrails_usage_overview(
start_date=START_DATE, end_date=END_DATE, user_api_key_dict=ADMIN_AUTH
)
matching = [r for r in response.rows if r.id == "yaml-uuid-xyz"]
assert len(matching) == 1
row = matching[0]
assert row.requestsEvaluated == 10
assert row.failRate == 20.0
assert row.provider == "presidio"
# Should NOT also produce an orphan "Custom" row keyed by the logical name
orphan_rows = [r for r in response.rows if r.id == "my-yaml-pii"]
assert (
orphan_rows == []
), "Logical-name row should be merged into the YAML guardrail row"
@pytest.mark.asyncio
async def test_usage_overview_dedupe_when_guardrail_in_both_db_and_yaml(
mock_prisma, mock_in_memory
):
"""If the same guardrail_id appears in DB and in-memory, only one row is emitted."""
db_row = _make_db_guardrail(
guardrail_id="shared-id", guardrail_name="shared-name", provider="bedrock"
)
yaml_gr = _make_yaml_guardrail(
guardrail_id="shared-id",
guardrail_name="shared-name",
provider="presidio",
)
mock_prisma.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[db_row])
mock_in_memory.list_in_memory_guardrails.return_value = [yaml_gr]
response = await guardrails_usage_overview(
start_date=START_DATE, end_date=END_DATE, user_api_key_dict=ADMIN_AUTH
)
matching = [r for r in response.rows if r.id == "shared-id"]
assert len(matching) == 1
# DB row wins
assert matching[0].provider == "bedrock"
# ---- /guardrails/usage/logs ------------------------------------------------
@pytest.mark.asyncio
async def test_usage_logs_includes_logical_name_for_yaml_guardrail(
mock_prisma, mock_in_memory
):
"""For a YAML guardrail (DB miss), the logs query must include the logical name alias."""
yaml_gr = _make_yaml_guardrail(
guardrail_id="yaml-uuid-xyz", guardrail_name="my-yaml-pii"
)
mock_in_memory.get_guardrail_by_id.return_value = yaml_gr
captured: dict = {}
async def _capture_find_many(**kwargs):
captured.update(kwargs)
return []
mock_prisma.db.litellm_spendlogguardrailindex.find_many = AsyncMock(
side_effect=_capture_find_many
)
await guardrails_usage_logs(
guardrail_id="yaml-uuid-xyz",
policy_id=None,
page=1,
page_size=50,
action=None,
start_date=None,
end_date=None,
user_api_key_dict=ADMIN_AUTH,
)
where = captured.get("where", {})
gid_filter = where.get("guardrail_id")
# The query should accept either the UUID or the logical name
assert (
isinstance(gid_filter, dict) and "in" in gid_filter
), f"Expected 'in' filter to include UUID + logical name, got {gid_filter}"
assert "yaml-uuid-xyz" in gid_filter["in"]
assert "my-yaml-pii" in gid_filter["in"]