diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 601a1376cb3..485389b8b39 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -48,11 +48,15 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ): self.flush_lock = asyncio.Lock() kwargs.setdefault("guardrail_name", "rubrik") - # Use `or` rather than setdefault because callers (e.g. - # initialize_guardrail) always pass event_hook explicitly, possibly as - # None when the user omits `mode` from the guardrail config. + # `initialize_guardrail` always passes these kwargs explicitly, with + # value `None` when the user omits `mode` / `default_on` from the + # guardrail config. Coerce None (omitted) to the desired default + # while preserving any explicit value the caller did set -- + # in particular `default_on=False` if the user wants the guardrail + # off by default. kwargs["event_hook"] = kwargs.get("event_hook") or GuardrailEventHooks.post_call - kwargs.setdefault("default_on", True) + if kwargs.get("default_on") is None: + kwargs["default_on"] = True super().__init__( flush_lock=self.flush_lock, **kwargs, diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 113868480c5..56799641d90 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -204,6 +204,27 @@ class TestInitialization: handler = RubrikLogger(event_hook=GuardrailEventHooks.pre_call) assert handler.event_hook == GuardrailEventHooks.pre_call + def test_default_on_defaults_to_true_when_none_passed(self, mock_env): + """`initialize_guardrail` always passes ``default_on=litellm_params.default_on`` + (which is ``None`` when the user omits ``default_on``). The logger must + coerce a None ``default_on`` to True, otherwise ``should_run_guardrail`` + (which checks ``self.default_on is True``) silently skips the guardrail.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=None) + assert handler.default_on is True + + def test_explicit_default_on_false_preserved(self, mock_env): + """A user explicitly setting ``default_on: false`` in their guardrail + config must NOT be silently overridden to True.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=False) + assert handler.default_on is False + + def test_explicit_default_on_true_preserved(self, mock_env): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=True) + assert handler.default_on is True + def test_headers_with_api_key(self, handler): assert handler._headers["Authorization"] == "Bearer test-api-key" assert handler._headers["Content-Type"] == "application/json"