From c770756cf3c82292e2bc561caf27342b6287c9a8 Mon Sep 17 00:00:00 2001
From: shivam <shivam@berri.ai>
Date: Mon, 20 Apr 2026 19:42:51 -0700
Subject: [PATCH] fix(bedrock_guardrails): route apply_guardrail to OUTPUT for
 response scans

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
---
 .../guardrail_hooks/bedrock_guardrails.py     | 41 +++++++++++++++++--
 1 file changed, 37 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
index 25b46cf364..77b2f466f2 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
@@ -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")