From 8bf0b9e19bb79c57b784cf74b88de5b26d1c28ad Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:16:30 +0900 Subject: [PATCH] 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 62c997659bd..c89decfd01c 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 6915286a2d7..59381149809 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 55fa17c21e7..18228ae4b48 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 6d4ed089818..953275acf14 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 6d98866eadf..732331349e0 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 51136c29eca..4b56fa5fc90 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 a0ea90ccf21..04f88e1a3da 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 88145ae9e47..02e481acddd 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