scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix/bedrock-inconsistent-postcall-hook (#19151)

Browse source 
* fix/bedrock-inconsistent-postcall-hook

* Add condition check to avoid multiple validation
This commit is contained in:
Harshit Jain 2026-01-19 19:48:02 +05:30 • committed by GitHub
parent 1dc2d2ddac
commit fc9988b686
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 83 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

110
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -893,15 +893,53 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return
#########################################################
########## 1. Make parallel Bedrock API requests ##########
########## 1. Make Bedrock API requests ##########
#########################################################
# Import asyncio for parallel execution
import asyncio
# Determine if INPUT validation is needed in post_call
# Skip INPUT validation if pre_call or during_call is already enabled
# (to avoid redundant validation - those hooks would have already validated INPUT)
should_validate_input = not (
self._event_hook_is_event_type(GuardrailEventHooks.pre_call)
or self._event_hook_is_event_type(GuardrailEventHooks.during_call)
)
output_content_bedrock: Optional[Union[BedrockGuardrailResponse, str]] = None
try:
output_content_bedrock = await self.make_bedrock_api_request(
if should_validate_input:
# Prepare input messages (with optional filtering for latest role message)
input_filter = self._prepare_guardrail_messages_for_role(
messages=new_messages
)
input_messages = input_filter.payload_messages or new_messages
# Create tasks for parallel execution of both INPUT and OUTPUT validation
input_task = self.make_bedrock_api_request(
source="INPUT",
messages=input_messages,
request_data=data,
)
output_task = self.make_bedrock_api_request(
source="OUTPUT", response=response, request_data=data
)
except GuardrailInterventionNormalStringError as e:
output_content_bedrock = e.message
# Execute both requests in parallel
try:
_, output_content_bedrock = await asyncio.gather(
input_task, output_task
)
except GuardrailInterventionNormalStringError as e:
output_content_bedrock = e.message
else:
# 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
)
except GuardrailInterventionNormalStringError as e:
output_content_bedrock = e.message
#########################################################
########## 2. Apply masking to response with output guardrail response ##########
@ -910,7 +948,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
response = self.create_guardrail_blocked_response(
response=output_content_bedrock
)
else:
elif output_content_bedrock is not None:
self._apply_masking_to_response(
response=response,
bedrock_guardrail_response=output_content_bedrock,
@ -997,36 +1035,54 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
if isinstance(assembled_model_response, ModelResponse):
####################################################################
########## 1. Make parallel Bedrock Apply Guardrail API requests ##########
########## 1. Make Bedrock Apply Guardrail API requests ##########
# Bedrock will raise an exception if this violates the guardrail policy
###################################################################
# Create tasks for parallel execution
input_filter = self._prepare_guardrail_messages_for_role(
messages=request_data.get("messages")
# Determine if INPUT validation is needed in post_call
# Skip INPUT validation if pre_call or during_call is already enabled
# (to avoid redundant validation - those hooks would have already validated INPUT)
should_validate_input = not (
self._event_hook_is_event_type(GuardrailEventHooks.pre_call)
or self._event_hook_is_event_type(GuardrailEventHooks.during_call)
)
input_messages = input_filter.payload_messages or request_data.get(
"messages"
)
input_task = self.make_bedrock_api_request(
source="INPUT",
messages=input_messages,
request_data=request_data,
) # Only input messages
output_guardrail_response: Optional[
Union[BedrockGuardrailResponse, str]
] = None
output_task = self.make_bedrock_api_request(
source="OUTPUT", response=assembled_model_response
) # Only response
# Execute both requests in parallel
try:
_, output_guardrail_response = await asyncio.gather(
input_task, output_task
if should_validate_input:
# Create tasks for parallel execution
input_filter = self._prepare_guardrail_messages_for_role(
messages=request_data.get("messages")
)
except GuardrailInterventionNormalStringError as e:
output_guardrail_response = e.message
input_messages = input_filter.payload_messages or request_data.get(
"messages"
)
input_task = self.make_bedrock_api_request(
source="INPUT",
messages=input_messages,
request_data=request_data,
) # Only input messages
output_task = self.make_bedrock_api_request(
source="OUTPUT", response=assembled_model_response
) # Only response
# Execute both requests in parallel
try:
_, output_guardrail_response = await asyncio.gather(
input_task, output_task
)
except GuardrailInterventionNormalStringError as e:
output_guardrail_response = e.message
else:
# 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
)
except GuardrailInterventionNormalStringError as e:
output_guardrail_response = e.message
#########################################################################
########## 2. Apply masking to response with output guardrail response ##########