diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 6046f1bb581..4bc7da48911 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -2,6 +2,7 @@ from datetime import datetime from typing import ( TYPE_CHECKING, Any, + ClassVar, Dict, List, Literal, @@ -81,6 +82,9 @@ class ModifyResponseException(Exception): class CustomGuardrail(CustomLogger): + # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. + use_native_during_call_hook: ClassVar[bool] = False + def __init__( self, guardrail_name: Optional[str] = None, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index b4b8e681133..ba04f92cffd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -18,6 +18,7 @@ from typing import ( TYPE_CHECKING, Any, AsyncGenerator, + ClassVar, Dict, List, Literal, @@ -123,6 +124,10 @@ def _redact_pii_matches(response_json: dict) -> dict: class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): + # During-call must use async_moderation_hook (not unified apply_guardrail), otherwise + # OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL. + use_native_during_call_hook: ClassVar[bool] = True + def __init__( self, guardrailIdentifier: Optional[str] = None, @@ -413,6 +418,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): messages: Optional[List[AllMessageValues]] = None, response: Optional[Union[Any, litellm.ModelResponse]] = None, request_data: Optional[dict] = None, + logging_event_type: Optional[GuardrailEventHooks] = None, ) -> BedrockGuardrailResponse: from datetime import datetime @@ -450,11 +456,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - event_type = ( - GuardrailEventHooks.pre_call - if source == "INPUT" - else GuardrailEventHooks.post_call - ) + # UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API + # body, which must not be confused with the proxy hook (pre_call / during_call / + # post_call). When omitted, keep legacy mapping for backward compatibility. + if logging_event_type is not None: + event_type = logging_event_type + else: + event_type = ( + GuardrailEventHooks.pre_call + if source == "INPUT" + else GuardrailEventHooks.post_call + ) try: httpx_response = await self.async_handler.post( @@ -944,7 +956,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( - source="INPUT", messages=filtered_messages, request_data=data + source="INPUT", + messages=filtered_messages, + request_data=data, + logging_event_type=GuardrailEventHooks.pre_call, ) except GuardrailInterventionNormalStringError as e: bedrock_guardrail_response = e.message @@ -1016,7 +1031,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( - source="INPUT", messages=filtered_messages, request_data=data + source="INPUT", + messages=filtered_messages, + request_data=data, + logging_event_type=GuardrailEventHooks.during_call, ) except GuardrailInterventionNormalStringError as e: bedrock_guardrail_response = e.message @@ -1120,9 +1138,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source="INPUT", messages=input_messages, request_data=data, + logging_event_type=GuardrailEventHooks.post_call, ) output_task = self.make_bedrock_api_request( - source="OUTPUT", response=response, request_data=data + source="OUTPUT", + response=response, + request_data=data, + logging_event_type=GuardrailEventHooks.post_call, ) # Execute both requests in parallel @@ -1136,7 +1158,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Only run OUTPUT validation (INPUT was already validated in pre_call or during_call) try: output_content_bedrock = await self.make_bedrock_api_request( - source="OUTPUT", response=response, request_data=data + source="OUTPUT", + response=response, + request_data=data, + logging_event_type=GuardrailEventHooks.post_call, ) except GuardrailInterventionNormalStringError as e: output_content_bedrock = e.message @@ -1263,9 +1288,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source="INPUT", messages=input_messages, request_data=request_data, + logging_event_type=GuardrailEventHooks.post_call, ) # Only input messages output_task = self.make_bedrock_api_request( - source="OUTPUT", response=assembled_model_response + source="OUTPUT", + response=assembled_model_response, + logging_event_type=GuardrailEventHooks.post_call, ) # Only response # Execute both requests in parallel @@ -1279,7 +1307,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Only run OUTPUT validation (INPUT was already validated in pre_call or during_call) try: output_guardrail_response = await self.make_bedrock_api_request( - source="OUTPUT", response=assembled_model_response + source="OUTPUT", + response=assembled_model_response, + logging_event_type=GuardrailEventHooks.post_call, ) except GuardrailInterventionNormalStringError as e: output_guardrail_response = e.message @@ -1554,10 +1584,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Bedrock will throw an error if there is no text to process if filtered_messages: + _log_hook = ( + GuardrailEventHooks.pre_call + if input_type == "request" + else GuardrailEventHooks.post_call + ) bedrock_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=request_data, + logging_event_type=_log_hook, ) # Apply any masking that was applied by the guardrail diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e15b48577de..f137b0bb3ba 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -940,7 +940,11 @@ class ProxyLogging: Result from the guardrail execution """ # Use unified_guardrail if callback has apply_guardrail method - use_unified = "apply_guardrail" in type(callback).__dict__ + has_apply_guardrail = "apply_guardrail" in type(callback).__dict__ + use_unified = has_apply_guardrail and not ( + hook_type == "during_call" + and getattr(callback, "use_native_during_call_hook", False) + ) if use_unified: data["guardrail_to_apply"] = callback @@ -1537,10 +1541,12 @@ class ProxyLogging: else: user_api_key_auth_dict = user_api_key_dict # Add task to list for parallel execution - if ( + use_unified_during = ( "apply_guardrail" in type(callback).__dict__ and user_api_key_dict is not None - ): + and not getattr(callback, "use_native_during_call_hook", False) + ) + if use_unified_during: data["guardrail_to_apply"] = callback guardrail_task = self._run_guardrail_task_with_enrichment( callback, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 010ead425ca..95b5f4dd15c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1204,6 +1204,11 @@ def _make_guardrail() -> BedrockGuardrail: ) +def test_bedrock_guardrail_uses_native_during_call_hook(): + """during_call must use async_moderation_hook, not unified apply_guardrail(input=request).""" + assert BedrockGuardrail.use_native_during_call_hook is True + + def test_extract_blocked_assessments_pii_entity(): """L3: PII entity match (BLOCKED) is surfaced with category, type, and matched term.""" g = _make_guardrail()