diff --git a/litellm/proxy/example_config_yaml/custom_guardrail.py b/litellm/proxy/example_config_yaml/custom_guardrail.py index 1056c1397fe..330bdd63342 100644 --- a/litellm/proxy/example_config_yaml/custom_guardrail.py +++ b/litellm/proxy/example_config_yaml/custom_guardrail.py @@ -4,7 +4,10 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.exceptions import GuardrailRaisedException -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import CallTypesLiteral @@ -79,6 +82,7 @@ class myCustomGuardrail(CustomGuardrail): return data + @log_guardrail_information async def async_moderation_hook( self, data: dict, @@ -107,6 +111,7 @@ class myCustomGuardrail(CustomGuardrail): blocked_content=True, ) + @log_guardrail_information async def async_post_call_success_hook( self, data: dict, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 35a99b759d4..a2fc6944df5 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1889,22 +1889,28 @@ class TestGuardrailInterventionClassification: async def test_example_config_guardrail_blocks_with_blocked_content(self): """The example_config_yaml copy is what otel_test_config.yaml loads; it must block the same way.""" guardrail = ExampleConfigGuardrail(guardrail_name="example-config") + input_request: dict = {"metadata": {}, "messages": [{"role": "user", "content": "tell me about litellm"}]} + output_request: dict = {"metadata": {}} with pytest.raises(GuardrailRaisedException) as input_block: await guardrail.async_moderation_hook( - data={"messages": [{"role": "user", "content": "tell me about litellm"}]}, + data=input_request, user_api_key_dict=UserAPIKeyAuth(), call_type="completion", ) assert input_block.value.blocked_content is True + input_slg = input_request["metadata"]["standard_logging_guardrail_information"][0] + assert input_slg["guardrail_status"] == "guardrail_intervened" with pytest.raises(GuardrailRaisedException) as output_block: await guardrail.async_post_call_success_hook( - data={}, + data=output_request, user_api_key_dict=UserAPIKeyAuth(), response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="coffee"))]), ) assert output_block.value.blocked_content is True + output_slg = output_request["metadata"]["standard_logging_guardrail_information"][0] + assert output_slg["guardrail_status"] == "guardrail_intervened" class _ApplyStyleGuardrail(CustomGuardrail):