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

fix(guardrails): record guardrail information for undecorated custom apply_guardrail overrides (#39727)

Browse source 
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-05 11:39:24 -07:00 • committed by GitHub
parent 7672399c26
commit 4df284e16d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 216 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/constants.py
View file

@ -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)

9
litellm/integrations/custom_guardrail.py
View file

@ -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

168
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -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."""

31
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -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."""

6
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -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