diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index e9eee3a1a8a..1f2c9e0c182 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -494,6 +494,7 @@ class InMemoryGuardrailHandler: guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail["guardrail_name"], litellm_params=litellm_params, + guardrail_info=guardrail.get("guardrail_info"), ) # store references to the guardrail in memory @@ -612,6 +613,27 @@ class InMemoryGuardrailHandler: """ return self._sources.get(guardrail_id) + def list_config_guardrails(self) -> List[Guardrail]: + """ + List in-memory guardrails owned by config.yaml. + + DB-sourced entries are excluded: a read surface that also queries the DB + would double-count live ones, and a DB-sourced entry that's missing from + the DB is stale (deleted on another pod, awaiting reconciliation here). + """ + return [g for gid, g in self.IN_MEMORY_GUARDRAILS.items() if self._sources.get(gid) == "config"] + + def get_config_guardrail_by_id(self, guardrail_id: str) -> Optional[Guardrail]: + """ + Get a config-owned in-memory guardrail by its ID, or None. + + Mirrors the fallback in get_guardrail_info: a DB-sourced in-memory entry + that missed the DB lookup is stale and must not be surfaced. + """ + if self._sources.get(guardrail_id) != "config": + return None + return self.IN_MEMORY_GUARDRAILS.get(guardrail_id) + def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]: """ Drop in-memory entries that originated from the DB but are no longer diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index e03bdbb95d2..f56b22ddd49 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -137,10 +137,26 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: return [{"date": d, "passed": v["passed"], "blocked": v["blocked"]} for d, v in sorted(chart_by_date.items())] +def _get_guardrail_field(g: Any, field: str) -> Any: + """Read `field` off a guardrail whether it's a Prisma row (attr) or a 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 pydantic model (e.g. LitellmParams) / dict value into a plain dict.""" + if isinstance(value, BaseModel): + return value.model_dump(exclude_none=True) + if isinstance(value, dict): + return value + return {} + + def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" - gid = getattr(g, "guardrail_id", None) or (g.get("guardrail_id") if isinstance(g, dict) else None) - name = getattr(g, "guardrail_name", None) or (g.get("guardrail_name") if isinstance(g, dict) else None) + gid = _get_guardrail_field(g, "guardrail_id") + name = _get_guardrail_field(g, "guardrail_name") return gid, (name or gid or "") @@ -163,9 +179,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: @@ -262,9 +278,15 @@ async def guardrails_usage_overview( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + try: - # Guardrails from DB - guardrails = await GuardrailsRepository(prisma_client).table.find_many() + db_guardrails = await GuardrailsRepository(prisma_client).table.find_many() + seen_ids = {gid for g in db_guardrails if (gid := _get_guardrail_field(g, "guardrail_id")) is not None} + config_guardrails = [ + g for g in IN_MEMORY_GUARDRAIL_HANDLER.list_config_guardrails() if g.get("guardrail_id") not in seen_ids + ] + guardrails: List[Any] = [*db_guardrails, *config_guardrails] # Daily metrics in range metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( @@ -321,16 +343,18 @@ async def guardrails_usage_detail( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") - guardrail = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) - if not guardrail: + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + if guardrail is None: + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_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 DailyGuardrailMetricsRepository(prisma_client).table.find_many( @@ -367,17 +391,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, @@ -548,11 +564,15 @@ async def guardrails_usage_logs( # Query by both so we match regardless of which was written. effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else [] if guardrail_id: - guardrail = await GuardrailsRepository(prisma_client).table.find_unique( + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) + if guardrail is None: + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_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) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 0ef9ad857f9..26feddadf79 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -44,9 +44,7 @@ def test_update_in_memory_guardrail(): "123", Guardrail( guardrail_name="test-guardrail", - litellm_params=LitellmParams( - guardrail="test-guardrail", mode="pre_call", default_on=True - ), + litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True), ), ) @@ -56,10 +54,7 @@ def test_update_in_memory_guardrail(): ) is True ) - assert ( - handler.guardrail_id_to_custom_guardrail["123"].event_hook - is GuardrailEventHooks.pre_call - ) + assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: @@ -135,6 +130,34 @@ def test_delete_in_memory_guardrail_clears_source_marker(): assert handler.get_source("a") is None +def test_list_config_guardrails_excludes_db_sourced(): + """LIT-2529: read surfaces union DB rows with config guardrails; db-sourced + in-memory entries would double-count (or resurrect stale ones), so exclude them.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one") + handler._sources["db"] = "db" + + config_guardrails = handler.list_config_guardrails() + + assert [g["guardrail_id"] for g in config_guardrails] == ["cfg"] + + +def test_get_config_guardrail_by_id_returns_config_only(): + """LIT-2529: the detail/logs fallback must return config-owned guardrails and + treat a db-sourced (stale) or missing id as a miss.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one") + handler._sources["db"] = "db" + + assert handler.get_config_guardrail_by_id("cfg")["guardrail_name"] == "config-one" + assert handler.get_config_guardrail_by_id("db") is None + assert handler.get_config_guardrail_by_id("missing") is None + + def test_initialize_guardrail_early_return_updates_source_marker(): """ When initialize_guardrail is called for a guardrail that already exists @@ -152,9 +175,7 @@ def test_initialize_guardrail_early_return_updates_source_marker(): g = Guardrail( guardrail_id="collide", guardrail_name="bedrock", - litellm_params=LitellmParams( - guardrail="bedrock", mode="pre_call", default_on=False - ), + litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call", default_on=False), ) handler.initialize_guardrail(guardrail=g, source="config") @@ -331,10 +352,7 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances(): def distinct_runner_instances() -> int: seen = set() for callback in litellm.logging_callback_manager._get_all_callbacks(): - if ( - isinstance(callback, CustomGuardrail) - and getattr(callback, "guardrail_name", None) == name - ): + if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == name: seen.add(id(callback)) return len(seen) diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index a511229942a..83593c20110 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -5,9 +5,7 @@ from unittest.mock import MagicMock, patch import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -36,8 +34,31 @@ def test_initialize_presidio_guardrail(): ) assert result["guardrail_name"] == "test_presidio_guardrail" - assert ( - result["litellm_params"].guardrail - == SupportedGuardrailIntegrations.PRESIDIO.value - ) + assert result["litellm_params"].guardrail == SupportedGuardrailIntegrations.PRESIDIO.value assert result["litellm_params"].mode == "pre_call" + + +def test_initialize_guardrail_preserves_guardrail_info(): + """ + Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the + stored in-memory Guardrail. Dropping it left the Guardrail Monitor's usage + endpoints unable to render type/description for YAML-defined guardrails. + """ + test_guardrail = { + "guardrail_name": "test_presidio_with_info", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + }, + "guardrail_info": {"type": "PII", "description": "masks PII"}, + } + + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + assert result is not None + assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"} + stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]] + assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"} diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py new file mode 100644 index 00000000000..bf7b1b3b238 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -0,0 +1,239 @@ +""" +Tests for the /guardrails/usage/* endpoints backing the dashboard Guardrail Monitor. + +Regression (LIT-2529): guardrails defined in config.yaml live only in +IN_MEMORY_GUARDRAIL_HANDLER, so the monitor's overview/detail/logs endpoints — +which read the litellm_guardrailstable Prisma table — could not see them: +detail 404'd, overview omitted them (or rendered them as Custom/Guardrail +orphans), and logs missed their logical-name alias. +""" + +import os +import sys +from datetime import datetime +from typing import Any, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.proxy.guardrails.usage_endpoints import ( + guardrails_usage_detail, + guardrails_usage_logs, + guardrails_usage_overview, +) +from litellm.types.guardrails import Guardrail, LitellmParams + +ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) +# Query() defaults don't resolve to None when the handler is called directly. +START, END = "2026-04-20", "2026-04-27" + + +def _config_handler(*guardrails: Guardrail) -> InMemoryGuardrailHandler: + """A real handler seeded with config-sourced YAML guardrails (no callbacks).""" + handler = InMemoryGuardrailHandler() + for g in guardrails: + gid = g["guardrail_id"] + handler.IN_MEMORY_GUARDRAILS[gid] = g + handler._sources[gid] = "config" + return handler + + +def _yaml_guardrail( + guardrail_id: str = "yaml-1", + name: str = "yaml-pii", + provider: str = "presidio", + info: Optional[dict] = None, +) -> Guardrail: + return Guardrail( + guardrail_id=guardrail_id, + guardrail_name=name, + litellm_params=LitellmParams(guardrail=provider, mode="pre_call"), + guardrail_info=info if info is not None else {"type": "PII", "description": "yaml-defined"}, + ) + + +def _db_row(guardrail_id: str = "db-1", name: str = "db-guard", provider: str = "aim") -> Any: + """A Prisma-style row: attribute access, litellm_params/guardrail_info as plain dicts.""" + row = MagicMock(spec=["guardrail_id", "guardrail_name", "litellm_params", "guardrail_info"]) + row.guardrail_id = guardrail_id + row.guardrail_name = name + row.litellm_params = {"guardrail": provider, "mode": "pre_call"} + row.guardrail_info = {"type": "ContentSafety", "description": "db-defined"} + return row + + +def _metric(guardrail_id: str, date: str = "2026-04-25", requests: int = 10, passed: int = 8, blocked: int = 2) -> Any: + m = MagicMock() + m.guardrail_id = guardrail_id + m.date = date + m.requests_evaluated = requests + m.passed_count = passed + m.blocked_count = blocked + m.flagged_count = 0 + return m + + +def _prisma( + *, + find_many=None, + find_unique=None, + metrics=None, + index_find_many=None, +) -> MagicMock: + client = MagicMock() + db = client.db + db.litellm_guardrailstable.find_many = AsyncMock(return_value=find_many or []) + db.litellm_guardrailstable.find_unique = AsyncMock(return_value=find_unique) + db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=metrics or []) + db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=index_find_many or []) + db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0) + db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + return client + + +def _patches(prisma: MagicMock, handler: InMemoryGuardrailHandler): + return ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler), + ) + + +# ---- detail ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_detail_returns_yaml_guardrail_when_db_misses(): + prisma = _prisma(find_unique=None) + 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.guardrail_id == "yaml-1" + assert resp.guardrail_name == "yaml-pii" + assert resp.provider == "presidio" # coerced from the LitellmParams pydantic model + assert resp.type == "PII" # from guardrail_info + assert resp.description == "yaml-defined" + + +@pytest.mark.asyncio +async def test_detail_404_when_neither_db_nor_config(): + prisma = _prisma(find_unique=None) + handler = _config_handler() # empty + p1, p2 = _patches(prisma, handler) + with p1, p2, pytest.raises(HTTPException) as exc: + await guardrails_usage_detail(guardrail_id="ghost", start_date=START, end_date=END, user_api_key_dict=ADMIN) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_detail_does_not_surface_db_sourced_in_memory_entry(): + """A stale in-memory entry (source=db, gone from DB) must 404, not resurface.""" + prisma = _prisma(find_unique=None) + handler = InMemoryGuardrailHandler() + stale = _yaml_guardrail(guardrail_id="stale-1", name="stale") + handler.IN_MEMORY_GUARDRAILS["stale-1"] = stale + handler._sources["stale-1"] = "db" + p1, p2 = _patches(prisma, handler) + with p1, p2, pytest.raises(HTTPException) as exc: + await guardrails_usage_detail(guardrail_id="stale-1", start_date=START, end_date=END, user_api_key_dict=ADMIN) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_detail_db_row_still_resolves(): + prisma = _prisma(find_unique=_db_row(guardrail_id="db-1", provider="aim")) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.provider == "aim" + assert resp.type == "ContentSafety" + + +# ---- overview --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_overview_includes_yaml_guardrail_with_no_metrics(): + """The core bug: a YAML guardrail with zero metrics must still appear as a row.""" + prisma = _prisma(find_many=[]) # no DB guardrails + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + rows = [r for r in resp.rows if r.id == "yaml-1"] + assert len(rows) == 1 + assert rows[0].name == "yaml-pii" + assert rows[0].provider == "presidio" + assert rows[0].type == "PII" + assert rows[0].requestsEvaluated == 0 + + +@pytest.mark.asyncio +async def test_overview_yaml_metrics_matched_by_logical_name(): + """Daily metrics are keyed by logical name; the YAML row must pick them up.""" + prisma = _prisma( + find_many=[], + metrics=[_metric("yaml-pii", requests=10, blocked=2)], # keyed by name, not uuid + ) + 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) + rows = [r for r in resp.rows if r.id == "yaml-uuid"] + assert len(rows) == 1 + assert rows[0].requestsEvaluated == 10 + assert rows[0].failRate == 20.0 + # must not also emit an orphan row keyed by the logical name + assert [r for r in resp.rows if r.id == "yaml-pii"] == [] + + +@pytest.mark.asyncio +async def test_overview_excludes_db_sourced_in_memory_entry(): + """union must not resurrect a stale db-sourced in-memory guardrail.""" + prisma = _prisma(find_many=[]) + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _yaml_guardrail(guardrail_id="cfg", name="cfg-guard") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["stale"] = _yaml_guardrail(guardrail_id="stale", name="stale-guard") + handler._sources["stale"] = "db" + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + ids = {r.id for r in resp.rows} + assert "cfg" in ids + assert "stale" not in ids + + +# ---- logs ------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_logs_resolves_config_guardrail_logical_name(): + """The index query must include the YAML guardrail's logical name alias.""" + prisma = _prisma(find_unique=None) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + await guardrails_usage_logs( + guardrail_id="yaml-uuid", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + where = prisma.db.litellm_spendlogguardrailindex.find_many.call_args.kwargs["where"] + assert where["guardrail_id"] == {"in": ["yaml-uuid", "yaml-pii"]}