tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-18 02:17:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(bedrock_guardrails): route apply_guardrail to OUTPUT for response scans

Browse Source 
BedrockGuardrail.apply_guardrail hardcoded source="INPUT" regardless of the
input_type parameter. On the non-streaming post-call path (unified_guardrail
-> OpenAIChatCompletionsHandler.process_output_response -> apply_guardrail),
the model response text was sent to Bedrock as INPUT, so guardrail policies
configured for Output (e.g. PII/NAME blocking) returned action=NONE and the
response passed through unblocked. The streaming path was unaffected because
it calls make_bedrock_api_request(source="OUTPUT", ...) directly.

Map input_type to the correct Bedrock source ("request" -> INPUT,
"response" -> OUTPUT) and build a synthetic ModelResponse for the OUTPUT
path so _create_bedrock_output_content_request produces the correct payload.

Made-with: Cursor
This commit is contained in:
shivam
2026-04-20 19:42:51 -07:00
parent 26fcbc93e5
commit c770756cf3
1 changed files with 37 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+37 -4
View File
@@ -62,6 +62,7 @@ from litellm.types.utils import (
CallTypesLiteral,
Choices,
GuardrailStatus,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
@@ -1563,11 +1564,43 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# Bedrock will throw an error if there is no text to process
if filtered_messages:
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=request_data,
# Map the abstract input_type to the Bedrock source parameter.
# "request" -> INPUT (scan user-supplied content)
# "response" -> OUTPUT (scan model-generated content)
# Bedrock guardrail policies are often configured differently
# for Input vs Output (e.g. PII blocking only on Output), so
# the source MUST match where the text originated.
bedrock_source: Literal["INPUT", "OUTPUT"] = (
"OUTPUT" if input_type == "response" else "INPUT"
)
if bedrock_source == "OUTPUT":
# Build a synthetic ModelResponse whose choices carry the
# text(s) to scan, so _create_bedrock_output_content_request
# can produce the correct Bedrock OUTPUT payload.
synthetic_response = ModelResponse(
choices=[
Choices(
index=_idx,
message=Message(
role="assistant",
content=str(_msg.get("content") or ""),
),
finish_reason="stop",
)
for _idx, _msg in enumerate(filtered_messages)
]
)
bedrock_response = await self.make_bedrock_api_request(
source="OUTPUT",
response=synthetic_response,
request_data=request_data,
)
else:
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=request_data,
)
# Apply any masking that was applied by the guardrail
output_list = bedrock_response.get("output")