diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index f7457dd4ff9..45b20211ac5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -256,10 +256,8 @@ class CustomGuardrail(CustomLogger): scan_raw_request: bool = False, only_scan_new_messages: bool = False, timeout: float | None = None, - turn_off_message_logging: bool = False, - message_logging: bool = True, - **kwargs: object, - ) -> None: + **kwargs, + ): """ Initialize the CustomGuardrail class @@ -305,7 +303,7 @@ class CustomGuardrail(CustomLogger): self.run_in_parallel: bool = run_in_parallel self.scan_raw_request: bool = scan_raw_request self.only_scan_new_messages: bool = only_scan_new_messages - stream_scope_arg: Final[object] = kwargs.pop("stream_scope", None) + stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None)) # cast-ok: config self.apply_stream_scope(stream_scope_arg) self.logging_only_scope = None self.logging_only_continue_on_input_failure = False @@ -325,11 +323,7 @@ class CustomGuardrail(CustomLogger): "(default) to enforce validation and fail at startup.", validation_error, ) - super().__init__( - turn_off_message_logging=turn_off_message_logging, - message_logging=message_logging, - **kwargs, - ) + super().__init__(**kwargs) def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str: """Return a custom violation message if template is configured."""