diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7b728ede1a4..1136394889e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -888,12 +888,18 @@ def log_guardrail_information(func): async def async_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper self: CustomGuardrail = args[0] - request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} + # apply_guardrail signature: (self, inputs, request_data, input_type, ...) + # Support both keyword and positional callers for request_data (args[2]) + request_data: dict = ( + kwargs.get("data") + or kwargs.get("request_data") + or (args[2] if len(args) > 2 and isinstance(args[2], dict) else {}) + ) event_type = _infer_event_type_from_function_name(func.__name__) - # For apply_guardrail, infer event type from input_type kwarg + # For apply_guardrail, infer event type from input_type (kwarg or positional arg[3]) if event_type is None and func.__name__ == "apply_guardrail": - _input_type = kwargs.get("input_type") + _input_type = kwargs.get("input_type") or (args[3] if len(args) > 3 else None) if _input_type == "request": event_type = GuardrailEventHooks.pre_call elif _input_type == "response": @@ -929,12 +935,18 @@ def log_guardrail_information(func): def sync_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper self: CustomGuardrail = args[0] - request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} + # apply_guardrail signature: (self, inputs, request_data, input_type, ...) + # Support both keyword and positional callers for request_data (args[2]) + request_data: dict = ( + kwargs.get("data") + or kwargs.get("request_data") + or (args[2] if len(args) > 2 and isinstance(args[2], dict) else {}) + ) event_type = _infer_event_type_from_function_name(func.__name__) - # For apply_guardrail, infer event type from input_type kwarg + # For apply_guardrail, infer event type from input_type (kwarg or positional arg[3]) if event_type is None and func.__name__ == "apply_guardrail": - _input_type = kwargs.get("input_type") + _input_type = kwargs.get("input_type") or (args[3] if len(args) > 3 else None) if _input_type == "request": event_type = GuardrailEventHooks.pre_call elif _input_type == "response": diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index e5fd542028a..f86d2d4f47d 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1021,3 +1021,38 @@ class TestLogGuardrailInformationApplyGuardrailEventType: f"guardrail_mode must be a string, got {type(guardrail_mode)}: {guardrail_mode}" ) assert guardrail_mode == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_apply_guardrail_positional_input_type_logs_correct_event_type(self): + """ + When apply_guardrail is called with input_type passed positionally (args[3]), + the decorator must still infer the correct event_type. This guards against + regressions where only kwargs.get("input_type") is checked. + """ + mode = Mode(tags={}, default="pre_call") + + class _TestGuardrail(CustomGuardrail): + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type, + logging_obj=None, + ) -> GenericGuardrailAPIInputs: + return inputs + + guardrail = _TestGuardrail(guardrail_name="test", event_hook=mode) + request_data: dict = {"metadata": {}} + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + + # Pass input_type positionally — it lands in args[3], not kwargs + await guardrail.apply_guardrail(inputs, request_data, "request") + + slg_list = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(slg_list) == 1 + guardrail_mode = slg_list[0]["guardrail_mode"] + assert isinstance(guardrail_mode, str), ( + f"guardrail_mode must be a string, got {type(guardrail_mode)}: {guardrail_mode}" + ) + assert guardrail_mode == GuardrailEventHooks.pre_call