From 0fe3145e83d8cd02e2bebae472d654d6ca1de199 Mon Sep 17 00:00:00 2001 From: Miguel Armenta <37154380+ma-armenta@users.noreply.github.com> Date: Sat, 21 Feb 2026 14:39:17 -0600 Subject: [PATCH] Fix/presidio controls (#21798) * check should_run_guardrail in sync logging hook path * Add tests for CustomGuardrail logging behavior Added tests to ensure CustomGuardrail logging behavior based on the guardrail execution state. --------- Co-authored-by: Miguel Armenta --- litellm/litellm_core_utils/litellm_logging.py | 20 ++- .../test_litellm_logging.py | 119 ++++++++++++++++++ 2 files changed, 138 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c673ec1219..021b14f67a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1962,7 +1962,24 @@ class Logging(LiteLLMLoggingBaseClass): ) ## LOGGING HOOK ## for callback in callbacks: - if isinstance(callback, CustomLogger): + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks + + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + continue + + self.model_call_details, result = callback.logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + elif isinstance(callback, CustomLogger): self.model_call_details, result = callback.logging_hook( kwargs=self.model_call_details, result=result, @@ -5454,3 +5471,4 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: model_parameters={"stream": True}, hidden_params=hidden_params, ) + diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0c9adecf5c..283a1351ac 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -406,6 +406,125 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call dummy_logger.log_stream_event.assert_not_called() +def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj): + """Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False.""" + import datetime + + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.guardrails import GuardrailEventHooks + + class DummyGuardrail(CustomGuardrail): + pass + + class DummyLogger(CustomLogger): + pass + + logging_obj.stream = False + + model_response = ModelResponse( + id="resp-guardrail-skip", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + guardrail = DummyGuardrail( + guardrail_name="dummy-guardrail", + event_hook=GuardrailEventHooks.logging_only, + ) + guardrail.should_run_guardrail = MagicMock(return_value=False) + guardrail.logging_hook = MagicMock( + return_value=(logging_obj.model_call_details, model_response) + ) + + dummy_logger = DummyLogger() + dummy_logger.logging_hook = MagicMock( + return_value=(logging_obj.model_call_details, model_response) + ) + + with patch.object( + logging_obj, + "get_combined_callback_list", + return_value=[guardrail, dummy_logger], + ): + logging_obj.success_handler( + result=model_response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + cache_hit=False, + ) + + guardrail.should_run_guardrail.assert_called_once() + guardrail_call_kwargs = guardrail.should_run_guardrail.call_args.kwargs + assert guardrail_call_kwargs["event_type"] == GuardrailEventHooks.logging_only + guardrail.logging_hook.assert_not_called() + dummy_logger.logging_hook.assert_called_once() + + +def test_success_handler_runs_guardrail_logging_hook_when_enabled(logging_obj): + """Ensure CustomGuardrail logging_hook runs when should_run_guardrail is True.""" + import datetime + + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class DummyGuardrail(CustomGuardrail): + pass + + logging_obj.stream = False + + model_response = ModelResponse( + id="resp-guardrail-run", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + guardrail = DummyGuardrail( + guardrail_name="dummy-guardrail", + event_hook=GuardrailEventHooks.logging_only, + ) + guardrail.should_run_guardrail = MagicMock(return_value=True) + + def _guardrail_logging_hook(kwargs, result, call_type): + updated_kwargs = dict(kwargs) + updated_kwargs["guardrail_hook_ran"] = True + return updated_kwargs, result + + guardrail.logging_hook = MagicMock(side_effect=_guardrail_logging_hook) + + with patch.object( + logging_obj, + "get_combined_callback_list", + return_value=[guardrail], + ): + logging_obj.success_handler( + result=model_response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + cache_hit=False, + ) + + guardrail.should_run_guardrail.assert_called_once() + guardrail_call_kwargs = guardrail.should_run_guardrail.call_args.kwargs + assert guardrail_call_kwargs["event_type"] == GuardrailEventHooks.logging_only + guardrail.logging_hook.assert_called_once() + assert logging_obj.model_call_details.get("guardrail_hook_ran") is True + + def test_get_user_agent_tags(): from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup