Merge pull request #26144 from BerriAI/litellm_post_call_non_streaming

fix(bedrock_guardrails): use Bedrock OUTPUT source for apply_guardrail when scanning model responses

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")

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -17,6 +17,7 @@ from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
     BedrockGuardrail,
     _redact_pii_matches,
 )
+from litellm.types.utils import ModelResponse
 
 
 @pytest.mark.asyncio
@@ -1113,6 +1114,72 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response():
         print("✅ apply_guardrail with tool_calls test passed - no API call made")
 
 
+@pytest.mark.asyncio
+async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source():
+    """input_type='response' must call Bedrock with source=OUTPUT and assistant content.
+
+    Regression: apply_guardrail used to always use source=INPUT. Output-only Bedrock
+    policies (e.g. PII on model output) then returned action=NONE for non-streaming
+    completions that go through unified_guardrail -> process_output_response.
+    """
+    guardrail = BedrockGuardrail(
+        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
+    )
+    bedrock_none = {"action": "NONE", "output": [], "outputs": []}
+
+    with patch.object(
+        guardrail, "make_bedrock_api_request", new_callable=AsyncMock
+    ) as mock_api:
+        mock_api.return_value = bedrock_none
+
+        await guardrail.apply_guardrail(
+            inputs={"texts": ["first line", "second line"]},
+            request_data={"model": "gpt-4o"},
+            input_type="response",
+        )
+
+        mock_api.assert_called_once()
+        kwargs = mock_api.call_args.kwargs
+        assert kwargs["source"] == "OUTPUT"
+        assert kwargs["request_data"] == {"model": "gpt-4o"}
+        synthetic = kwargs["response"]
+        assert isinstance(synthetic, ModelResponse)
+        assert len(synthetic.choices) == 2
+        assert synthetic.choices[0].message.content == "first line"
+        assert synthetic.choices[0].message.role == "assistant"
+        assert synthetic.choices[1].message.content == "second line"
+        assert synthetic.choices[1].message.role == "assistant"
+
+
+@pytest.mark.asyncio
+async def test_bedrock_apply_guardrail_request_uses_INPUT_source():
+    """input_type='request' must call Bedrock with source=INPUT and user messages."""
+    guardrail = BedrockGuardrail(
+        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
+    )
+    bedrock_none = {"action": "NONE", "output": [], "outputs": []}
+
+    with patch.object(
+        guardrail, "make_bedrock_api_request", new_callable=AsyncMock
+    ) as mock_api:
+        mock_api.return_value = bedrock_none
+
+        await guardrail.apply_guardrail(
+            inputs={"texts": ["user prompt"]},
+            request_data={},
+            input_type="request",
+        )
+
+        mock_api.assert_called_once()
+        kwargs = mock_api.call_args.kwargs
+        assert kwargs["source"] == "INPUT"
+        assert kwargs["messages"] is not None
+        assert len(kwargs["messages"]) == 1
+        assert kwargs["messages"][0]["role"] == "user"
+        assert kwargs["messages"][0]["content"] == "user prompt"
+        assert kwargs.get("response") is None
+
+
 @pytest.mark.asyncio
 async def test_bedrock_guardrail_blocked_content_with_masking_enabled():
     """Test that BLOCKED content raises exception even when masking is enabled