From 6888d34ea5372bd17086e92ed30554baf3d7cdbd Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:20:24 +0900 Subject: [PATCH 1/4] feat: Log actual executed event type in guardrail logging --- litellm/integrations/custom_guardrail.py | 46 +++- .../integrations/test_custom_guardrail.py | 238 +++++++++++++++++- 2 files changed, 273 insertions(+), 11 deletions(-) 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 From 8bf0b9e19bb79c57b784cf74b88de5b26d1c28ad Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:16:30 +0900 Subject: [PATCH 2/4] feat: add event_type parameter to add_standard_logging_guardrail_information_to_request_data --- .../guardrail_hooks/bedrock_guardrails.py | 8 +++++ .../guardrail_hooks/dynamoai/dynamoai.py | 6 ++++ .../ibm_guardrails/ibm_detector.py | 12 +++++++ .../guardrail_hooks/javelin/javelin.py | 6 +++- .../guardrail_hooks/lakera_ai_v2.py | 4 +++ .../model_armor/model_armor.py | 2 ++ .../guardrails/guardrail_hooks/noma/noma.py | 32 +++++++++++++++---- .../panw_prisma_airs/panw_prisma_airs.py | 24 ++++++++++++-- 8 files changed, 84 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 62c997659b..c89decfd01 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -449,6 +449,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) + event_type = ( + GuardrailEventHooks.pre_call + if source == "INPUT" + else GuardrailEventHooks.post_call + ) + try: httpx_response = await self.async_handler.post( url=prepared_request.url, @@ -469,6 +475,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) # Re-raise the exception to maintain existing behavior raise @@ -486,6 +493,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) ######################################################### if httpx_response.status_code == 200: diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 6915286a2d..5938114980 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -97,6 +97,7 @@ class DynamoAIGuardrails(CustomGuardrail): async def _call_dynamoai_guardrails( self, messages: List[Dict[str, Any]], + event_type: GuardrailEventHooks, text_type: str = "input", request_data: Optional[dict] = None, ) -> DynamoAIResponse: @@ -157,6 +158,7 @@ class DynamoAIGuardrails(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json @@ -177,6 +179,7 @@ class DynamoAIGuardrails(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -332,6 +335,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=_messages, text_type="input", request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -380,6 +384,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=_messages, text_type="input", request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -460,6 +465,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=dynamoai_messages, text_type="output", request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 55fa17c21e..18228ae4b4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -108,6 +108,7 @@ class IBMGuardrailDetector(CustomGuardrail): async def _call_detector_server( self, contents: List[str], + event_type: GuardrailEventHooks, request_data: Optional[dict] = None, ) -> List[List[IBMDetectorDetection]]: """ @@ -172,6 +173,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json @@ -192,6 +194,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -199,6 +202,7 @@ class IBMGuardrailDetector(CustomGuardrail): async def _call_orchestrator( self, content: str, + event_type: GuardrailEventHooks, request_data: Optional[dict] = None, ) -> List[IBMDetectorDetection]: """ @@ -258,6 +262,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json.get("detections", []) @@ -278,6 +283,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -472,6 +478,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -500,6 +507,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -557,6 +565,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -585,6 +594,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -673,6 +683,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( @@ -702,6 +713,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index 6d4ed08981..953275acf1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -83,6 +83,7 @@ class JavelinGuardrail(CustomGuardrail): async def call_javelin_guard( self, request: JavelinGuardRequest, + event_type: GuardrailEventHooks, ) -> JavelinGuardResponse: """ Call the Javelin guard API. @@ -158,6 +159,7 @@ class JavelinGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) async def async_pre_call_hook( @@ -208,7 +210,9 @@ class JavelinGuardrail(CustomGuardrail): config=self.config if self.config else {}, ) - javelin_response = await self.call_javelin_guard(request=javelin_guard_request) + javelin_response = await self.call_javelin_guard( + request=javelin_guard_request, event_type=GuardrailEventHooks.pre_call + ) assessments = javelin_response.get("assessments", []) reject_prompt = "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 6d98866ead..732331349e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -70,6 +70,7 @@ class LakeraAIGuardrail(CustomGuardrail): self, messages: List[AllMessageValues], request_data: Dict, + event_type: GuardrailEventHooks, ) -> Tuple[LakeraAIResponse, Dict]: """ Call the Lakera AI v2 guard API. @@ -128,6 +129,7 @@ class LakeraAIGuardrail(CustomGuardrail): end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), masked_entity_count=masked_entity_count, + event_type=event_type, ) def _mask_pii_in_messages( @@ -214,6 +216,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( messages=new_messages, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) ######################################################### @@ -279,6 +282,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( messages=new_messages, request_data=data, + event_type=GuardrailEventHooks.during_call, ) ######################################################### diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 51136c29ec..4b56fa5fc9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -327,6 +327,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Override to store only the Model Armor API response, not the entire data dict. @@ -351,6 +352,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index a0ea90ccf2..04f88e1a3d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -163,6 +163,7 @@ class NomaGuardrail(CustomGuardrail): self, request_data: dict, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[str]: """Shared logic for processing user message checks""" start_time = datetime.now() @@ -213,6 +214,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if self.monitor_mode: @@ -242,6 +244,7 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[str]: """Shared logic for processing LLM response checks""" @@ -293,6 +296,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if self.monitor_mode: @@ -602,7 +606,9 @@ class NomaGuardrail(CustomGuardrail): return data try: - return await self._check_user_message(data, user_api_key_dict) + return await self._check_user_message( + data, user_api_key_dict, GuardrailEventHooks.pre_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -619,6 +625,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}") @@ -650,7 +657,9 @@ class NomaGuardrail(CustomGuardrail): return data try: - return await self._check_user_message(data, user_api_key_dict) + return await self._check_user_message( + data, user_api_key_dict, GuardrailEventHooks.during_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -667,6 +676,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}") @@ -700,7 +710,9 @@ class NomaGuardrail(CustomGuardrail): return response try: - return await self._check_llm_response(data, response, user_api_key_dict) + return await self._check_llm_response( + data, response, user_api_key_dict, GuardrailEventHooks.post_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_llm_response_check with "blocked" status raise @@ -717,6 +729,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}") @@ -728,9 +741,12 @@ class NomaGuardrail(CustomGuardrail): self, request_data: dict, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Union[Exception, str, dict, None]: """Check user message for policy violations""" - user_message = await self._process_user_message_check(request_data, user_auth) + user_message = await self._process_user_message_check( + request_data, user_auth, event_type + ) if not user_message: return request_data @@ -741,10 +757,11 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Any: """Check LLM response for policy violations""" content = await self._process_llm_response_check( - request_data, response, user_auth + request_data, response, user_auth, event_type ) if not content: return response @@ -858,7 +875,10 @@ class NomaGuardrail(CustomGuardrail): if isinstance(assembled_model_response, ModelResponse): try: processed_response = await self._check_llm_response( - request_data, assembled_model_response, user_api_key_dict + request_data, + assembled_model_response, + user_api_key_dict, + GuardrailEventHooks.post_call, ) except NomaBlockedMessage: raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 88145ae9e4..02e481acdd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral, ModelResponse if TYPE_CHECKING: @@ -523,6 +524,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): scan_result: Dict[str, Any], data: Dict[str, Any], start_time: datetime, + event_type: GuardrailEventHooks, is_response: bool = False, ) -> Optional[Dict[str, Any]]: """Handle API errors with fail-open/fail-closed logic.""" @@ -542,6 +544,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if scan_result.get("_always_block"): @@ -735,7 +738,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): return self._handle_api_error_with_logging( - scan_result, data, start_time, is_response=False + scan_result, + data, + start_time, + is_response=False, + event_type=GuardrailEventHooks.pre_call, ) end_time = datetime.now() @@ -749,6 +756,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=GuardrailEventHooks.pre_call, ) action = scan_result.get("action", "block") @@ -872,7 +880,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( - scan_result, data, start_time, is_response=True + scan_result, + data, + start_time, + is_response=True, + event_type=GuardrailEventHooks.post_call, ) return response @@ -887,6 +899,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=GuardrailEventHooks.post_call, ) action = scan_result.get("action", "block") @@ -1066,7 +1079,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( - scan_result, request_data, start_time, is_response=True + scan_result, + request_data, + start_time, + is_response=True, + event_type=EventHooks.post_call, ) for chunk in all_chunks: yield chunk @@ -1083,6 +1100,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=EventHooks.post_call, ) # Add guardrail to applied guardrails header for observability From 7d673c308da5da3c3d8c889e71a30894a7d83a74 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:23:07 +0900 Subject: [PATCH 3/4] lint --- .../guardrail_hooks/bedrock_guardrails.py | 18 +++++++++--------- .../ibm_guardrails/ibm_detector.py | 1 - .../guardrail_hooks/model_armor/model_armor.py | 4 +++- .../guardrails/guardrail_hooks/noma/noma.py | 5 +---- 4 files changed, 13 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index c89decfd01..c8cef6e279 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -613,10 +613,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. - If `self.mask_request_content` or `self.mask_response_content` is set to `True`, + If `self.mask_request_content` or `self.mask_response_content` is set to `True`, then use the output from the guardrail to mask the request or response content. - - However, even with masking enabled, content with action="BLOCKED" should still + + However, even with masking enabled, content with action="BLOCKED" should still raise an exception, only content with action="ANONYMIZED" should be masked. """ @@ -739,9 +739,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data @@ -811,9 +811,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 18228ae4b4..2fc0521364 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -143,7 +143,6 @@ class IBMGuardrailDetector(CustomGuardrail): ) try: - response = await self.async_handler.post( url=self.api_url, json=payload, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 4b56fa5fc9..a12eb2486d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -295,7 +295,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): filters = ( list(filter_results.values()) if isinstance(filter_results, dict) - else filter_results if isinstance(filter_results, list) else [] + else filter_results + if isinstance(filter_results, list) + else [] ) # Prefer sanitized text from deidentifyResult if present diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 04f88e1a3d..1794751a08 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -119,9 +119,7 @@ class NomaGuardrail(CustomGuardrail): self.api_base = api_base or os.environ.get( "NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE ) - self.application_id = application_id or os.environ.get( - "NOMA_APPLICATION_ID" - ) + self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID") self.default_application_id = "litellm" if monitor_mode is None: @@ -582,7 +580,6 @@ class NomaGuardrail(CustomGuardrail): data: dict, call_type: CallTypesLiteral, ) -> Optional[Union[Exception, str, dict]]: - verbose_proxy_logger.debug("Running Noma pre-call hook") if ( From c6d6ffd567055acde3247bcaa6d00a8adb59b9f8 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:28:16 +0900 Subject: [PATCH 4/4] fix: mypy error --- litellm/integrations/custom_guardrail.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7bc02bb3bb..6a76b57e7f 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -516,12 +516,13 @@ class CustomGuardrail(CustomLogger): from litellm.types.utils import GuardrailMode # Use event_type if provided, otherwise fall back to self.event_hook + guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): - guardrail_mode = GuardrailMode(**self.event_hook.model_dump()) + guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) # type: ignore[typeddict-item] else: - guardrail_mode = self.event_hook + guardrail_mode = self.event_hook # type: ignore[assignment] slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name,