diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index fe0ce208ee..7bc02bb3bb 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -243,14 +243,14 @@ class CustomGuardrail(CustomLogger): def _is_valid_response_type(self, result: Any) -> bool: """ Check if result is a valid LLMResponseTypes instance. - + Safely handles TypedDict types which don't support isinstance checks. For non-LiteLLM responses (like passthrough httpx.Response), returns True to allow them through. """ if result is None: return False - + try: # Try isinstance check on valid types that support it response_types = get_args(LLMResponseTypes) @@ -506,6 +506,7 @@ class CustomGuardrail(CustomLogger): duration: Optional[float] = None, masked_entity_count: Optional[Dict[str, int]] = None, guardrail_provider: Optional[str] = None, + event_type: Optional[GuardrailEventHooks] = None, ) -> None: """ Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc. @@ -514,14 +515,18 @@ class CustomGuardrail(CustomLogger): guardrail_json_response = str(guardrail_json_response) from litellm.types.utils import GuardrailMode + # Use event_type if provided, otherwise fall back to self.event_hook + if event_type is not None: + guardrail_mode = event_type + elif isinstance(self.event_hook, Mode): + guardrail_mode = GuardrailMode(**self.event_hook.model_dump()) + else: + guardrail_mode = self.event_hook + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, - guardrail_mode=( - GuardrailMode(**self.event_hook.model_dump()) # type: ignore - if isinstance(self.event_hook, Mode) - else self.event_hook - ), + guardrail_mode=guardrail_mode, guardrail_response=guardrail_json_response, guardrail_status=guardrail_status, start_time=start_time, @@ -589,6 +594,7 @@ class CustomGuardrail(CustomLogger): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -605,6 +611,7 @@ class CustomGuardrail(CustomLogger): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) return response @@ -615,6 +622,7 @@ class CustomGuardrail(CustomLogger): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -628,6 +636,7 @@ class CustomGuardrail(CustomLogger): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) raise e @@ -712,16 +721,32 @@ def log_guardrail_information(func): Logs for: - pre_call - during_call - - TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run + - post_call """ import asyncio import functools + def _infer_event_type_from_function_name( + func_name: str, + ) -> Optional[GuardrailEventHooks]: + """Infer the actual event type from the function name""" + if func_name == "async_pre_call_hook": + return GuardrailEventHooks.pre_call + elif func_name == "async_moderation_hook": + return GuardrailEventHooks.during_call + elif func_name in ( + "async_post_call_success_hook", + "async_post_call_streaming_hook", + ): + return GuardrailEventHooks.post_call + return None + @functools.wraps(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 {} + event_type = _infer_event_type_from_function_name(func.__name__) try: response = await func(*args, **kwargs) return self._process_response( @@ -730,6 +755,7 @@ def log_guardrail_information(func): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) except Exception as e: return self._process_error( @@ -738,6 +764,7 @@ def log_guardrail_information(func): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) @functools.wraps(func) @@ -745,18 +772,21 @@ def log_guardrail_information(func): 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 {} + event_type = _infer_event_type_from_function_name(func.__name__) try: response = func(*args, **kwargs) return self._process_response( response=response, request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) except Exception as e: return self._process_error( e=e, request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) @functools.wraps(func) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a719d102a7..a322dfe9a2 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -498,7 +498,7 @@ class TestPassthroughCallTypeHandling: def test_get_pre_call_type_with_allm_passthrough_route(self): """ Test that _get_pre_call_type correctly maps allm_passthrough_route. - + This tests Fix #1: allm_passthrough_route was not being handled, causing call_type to be None. """ from litellm.proxy.common_request_processing import ( @@ -509,14 +509,14 @@ class TestPassthroughCallTypeHandling: result = ProxyBaseLLMRequestProcessing._get_pre_call_type( route_type="allm_passthrough_route" ) - + # Should return allm_passthrough_route, not None assert result == "allm_passthrough_route" def test_get_pre_call_type_preserves_standard_mappings(self): """ Test that _get_pre_call_type still correctly maps standard route types. - + Ensures Fix #1 didn't break existing functionality. """ from litellm.proxy.common_request_processing import ( @@ -536,3 +536,235 @@ class TestPassthroughCallTypeHandling: ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses") == "responses" ) + + +class TestEventTypeLogging: + """Tests for event_type logging in guardrail information.""" + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call + from async_pre_call_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_pre_call_hook(self, data: dict, **kwargs): + return {"result": "pre_call_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_pre_call_hook(data=request_data) + + # Check that the guardrail_mode was set to pre_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call + from async_post_call_success_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_post_call_success_hook(self, data: dict, **kwargs): + return {"result": "post_call_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_post_call_success_hook(data=request_data) + + # Check that the guardrail_mode was set to post_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call + from async_moderation_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_moderation_hook(self, data: dict, **kwargs): + return {"result": "moderation_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_moderation_hook(data=request_data) + + # Check that the guardrail_mode was set to during_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call + from async_post_call_streaming_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_post_call_streaming_hook(self, data: dict, **kwargs): + return {"result": "streaming_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_post_call_streaming_hook(data=request_data) + + # Check that the guardrail_mode was set to post_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_returns_none_for_unknown_function_name( + self, + ): + """ + Test that log_guardrail_information decorator returns None for event_type + when function name doesn't match known patterns, and falls back to self.event_hook. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + @log_guardrail_information + async def some_other_hook(self, data: dict, **kwargs): + return {"result": "other_hook_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.some_other_hook(data=request_data) + + # Check that the guardrail_mode falls back to self.event_hook + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call + + def test_add_standard_logging_uses_event_type_over_event_hook(self): + """ + Test that add_standard_logging_guardrail_information_to_request_data + prioritizes event_type parameter over self.event_hook. + """ + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + ) + + request_data = {"metadata": {}} + + # Call with explicit event_type + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"result": "ok"}, + request_data=request_data, + guardrail_status="success", + event_type=GuardrailEventHooks.post_call, + ) + + # Should use the provided event_type (post_call), not the full event_hook list + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none( + self, + ): + """ + Test that add_standard_logging_guardrail_information_to_request_data + falls back to self.event_hook when event_type is None. + """ + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + request_data = {"metadata": {}} + + # Call with event_type=None + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"result": "ok"}, + request_data=request_data, + guardrail_status="success", + event_type=None, + ) + + # Should fall back to self.event_hook + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call