From d8c067413552a6ad57b48023dc9be3e69087b429 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 27 Apr 2026 16:42:29 -0700 Subject: [PATCH] fix(guardrails): preserve guardrail_info when storing YAML guardrails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR review feedback. InMemoryGuardrailHandler.initialize_guardrail was constructing the parsed Guardrail without guardrail_info, so even after the endpoint fallback fix, /guardrails/usage/{overview,detail} couldn't render type or description for YAML-defined guardrails — both fields would silently default to "Guardrail" / None at runtime. Pass guardrail_info through into IN_MEMORY_GUARDRAILS, and add a regression test that exercises the real handler (not a mock) end-to-end. Also refactors _get_guardrail_attrs to use the new _get_guardrail_field helper for consistency. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../proxy/guardrails/guardrail_registry.py | 1 + litellm/proxy/guardrails/usage_endpoints.py | 18 +++---- .../proxy/guardrails/test_usage_endpoints.py | 54 +++++++++++++++++++ 3 files changed, 62 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 868b23756d2..bef43851fa6 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -487,6 +487,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 diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 8ee28cfadeb..7cde1136289 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -137,17 +137,6 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: ] -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 - ) - 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): @@ -155,6 +144,13 @@ def _get_guardrail_field(g: Any, field: str) -> Any: return getattr(g, field, None) +def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: + """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" + gid = _get_guardrail_field(g, "guardrail_id") + name = _get_guardrail_field(g, "guardrail_name") + return gid, (name or gid or "") + + def _to_dict(value: Any) -> Dict[str, Any]: """Coerce a LitellmParams / dict / None into a plain dict.""" if isinstance(value, LitellmParams): diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index b1aca532854..a4dc2b19bc7 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -304,3 +304,57 @@ async def test_usage_logs_includes_logical_name_for_yaml_guardrail( ), 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"] + + +# ---- Integration: real InMemoryGuardrailHandler ---------------------------- + + +@pytest.mark.asyncio +async def test_usage_detail_with_real_in_memory_handler_preserves_guardrail_info( + mock_prisma, mocker +): + """ + Regression: `initialize_guardrail` must persist `guardrail_info` into + IN_MEMORY_GUARDRAILS so that /usage/detail can render `type` and `description` + for YAML-defined guardrails. Exercises the real handler (not a Mock). + """ + from litellm.proxy.guardrails.guardrail_registry import ( + IN_MEMORY_GUARDRAIL_HANDLER, + ) + + # Bypass callback initialization; we only care about what gets stored. + mocker.patch.object( + IN_MEMORY_GUARDRAIL_HANDLER, + "initialize_custom_guardrail", + return_value=None, + ) + + yaml_input = { + "guardrail_id": "real-handler-yaml", + "guardrail_name": "real-handler-pii", + # `module.Class` form routes to initialize_custom_guardrail (mocked above) + "litellm_params": { + "guardrail": "my_module.MyCustomGuardrail", + "mode": "pre_call", + }, + "guardrail_info": {"type": "PII", "description": "Real-handler-defined"}, + } + try: + IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(guardrail=yaml_input) + + response = await guardrails_usage_detail( + guardrail_id="real-handler-yaml", + start_date=START_DATE, + end_date=END_DATE, + user_api_key_dict=ADMIN_AUTH, + ) + + assert response.guardrail_id == "real-handler-yaml" + assert response.provider == "my_module.MyCustomGuardrail" + assert response.type == "PII" + assert response.description == "Real-handler-defined" + finally: + IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.pop("real-handler-yaml", None) + IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.pop( + "real-handler-yaml", None + )