From 4df284e16dcc02ceabad4576df6f2c976f20d839 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:39:24 -0700 Subject: [PATCH] fix(guardrails): record guardrail information for undecorated custom apply_guardrail overrides (#39727) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 3 + litellm/integrations/custom_guardrail.py | 9 + .../integrations/test_custom_guardrail.py | 168 ++++++++++++++++++ .../test_openai_guardrail_handler.py | 31 ++++ .../test_bedrock_guardrails.py | 6 +- 5 files changed, 216 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 7d6de612349..e4fb2a00297 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -215,6 +215,9 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # so the deployment-level hook does not re-run them for the same request PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" +# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again +LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9bb613654e4..c462b1edb98 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -46,6 +46,7 @@ dc: Final = DualCache() from litellm.constants import ( GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + LOGS_GUARDRAIL_INFORMATION_MARKER, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) from litellm.exceptions import ( @@ -151,6 +152,13 @@ class CustomGuardrail(CustomLogger): records_own_guardrail_information: ClassVar[bool] = False + def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks + super().__init_subclass__(**kwargs) + own_apply_guardrail: Final = cls.__dict__.get("apply_guardrail") + if own_apply_guardrail is None or LOGS_GUARDRAIL_INFORMATION_MARKER in vars(own_apply_guardrail): + return + cls.apply_guardrail = log_guardrail_information(own_apply_guardrail) + def __init__( self, guardrail_name: str | None = None, @@ -1559,4 +1567,5 @@ def log_guardrail_information(func): return async_wrapper(*args, **kwargs) return sync_wrapper(*args, **kwargs) + vars(wrapper)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the wrapper this call just built return wrapper diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6edc2c9bf77..49a52157c8a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,5 @@ import asyncio +from typing import TYPE_CHECKING, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -11,6 +12,9 @@ from litellm.integrations.custom_guardrail import ( from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class TestCustomGuardrailDeploymentHook: @@ -2239,6 +2243,170 @@ class TestRecordsOwnGuardrailInformation: assert _guardrail_entries(request_data) == [] +class _UndecoratedGuardrail(CustomGuardrail): + """apply_guardrail written like the docs example: no @log_guardrail_information.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + from litellm.exceptions import GuardrailRaisedException + + if any("forbidden" in text for text in inputs.get("texts") or []): + raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message="Content blocked") + return inputs + + +class _UndecoratedSelfRecordingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"custom": True}, + request_data=request_data, + guardrail_status="success", + start_time=0.0, + end_time=0.0, + duration=0.0, + ) + return inputs + + +class _InheritedApplyGuardrail(_UndecoratedGuardrail): + pass + + +class TestUndecoratedApplyGuardrailIsLogged: + """LIT-5983 regression: a custom guardrail that overrides apply_guardrail without the + @log_guardrail_information decorator must still record guardrail information, and the + auto-wrap must not double-record decorated or self-recording implementations.""" + + @pytest.mark.asyncio + async def test_undecorated_success_is_recorded(self): + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = _UndecoratedGuardrail(guardrail_name="docs-style", event_hook=GuardrailEventHooks.pre_call) + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "docs-style" + assert entries[0]["guardrail_mode"] == "pre_call" + assert entries[0]["guardrail_status"] == "success" + + @pytest.mark.asyncio + async def test_undecorated_block_is_recorded_and_reraised(self): + from litellm.exceptions import GuardrailRaisedException + + guardrail = _UndecoratedGuardrail(guardrail_name="docs-style") + request_data: dict = {"model": "gpt-4o"} + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["forbidden"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "docs-style" + assert entries[0]["guardrail_status"] == "guardrail_intervened" + + @pytest.mark.asyncio + async def test_undecorated_bare_exception_is_recorded_as_failed_to_respond(self): + class _BareExceptionGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + raise Exception("Content blocked: policy violation") + + guardrail = _BareExceptionGuardrail(guardrail_name="docs-style") + request_data: dict = {"model": "gpt-4o"} + + with pytest.raises(Exception, match="Content blocked"): + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["x"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" + + @pytest.mark.asyncio + async def test_inherited_apply_guardrail_is_recorded_once(self): + guardrail = _InheritedApplyGuardrail(guardrail_name="child") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + assert len(_guardrail_entries(request_data)) == 1 + + @pytest.mark.asyncio + async def test_undecorated_self_recording_apply_guardrail_is_recorded_once(self): + guardrail = _UndecoratedSelfRecordingGuardrail(guardrail_name="self-recording") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_response"] == {"custom": True} + + @pytest.mark.asyncio + async def test_base_apply_guardrail_is_not_recorded(self): + guardrail = CustomGuardrail(guardrail_name="base") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + assert _guardrail_entries(request_data) == [] + + def test_subclass_keywords_reach_cooperative_init_subclass(self): + class _LabelMixin: + seen_label: str = "" + + def __init_subclass__(cls, label: str = "", **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + cls.seen_label = label + + class _Labelled(CustomGuardrail, _LabelMixin, label="docs-style"): + pass + + assert _Labelled.seen_label == "docs-style" + + class _ApplyOnlyObserver(CustomGuardrail): """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index cebab2512d0..36e715d5804 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1075,6 +1075,37 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert result == responses_so_far +class TestUndecoratedGuardrailIsRecorded: + """LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail + without @log_guardrail_information must still end up in the request's guardrail + information on both the request and response paths.""" + + @pytest.mark.asyncio + async def test_request_path_records_undecorated_guardrail(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="docs-style") + data = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + + await handler.process_input_messages(data, guardrail) + + entries = data["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")] + + @pytest.mark.asyncio + async def test_response_path_records_undecorated_guardrail(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="docs-style") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))] + ) + request_data: dict = {"metadata": {}} + + await handler.process_output_response(response, guardrail, request_data=request_data) + + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")] + + class TestGetStructuredMessages: """Test the get_structured_messages method.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 9842d88e8d1..7da55f22bda 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1137,7 +1137,11 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): mock_api.assert_called_once() kwargs = mock_api.call_args.kwargs assert kwargs["source"] == "OUTPUT" - assert kwargs["request_data"] == {"model": "gpt-4o"} + assert kwargs["request_data"]["model"] == "gpt-4o" + recorded = kwargs["request_data"]["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in recorded] == [ + (guardrail.guardrail_name, "success") + ] synthetic = kwargs["response"] assert isinstance(synthetic, ModelResponse) assert len(synthetic.choices) == 2