diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 25b46cf36412..77b2f466f2aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -62,6 +62,7 @@ CallTypesLiteral, Choices, GuardrailStatus, + Message, ModelResponse, ModelResponseStream, StreamingChoices, @@ -1563,11 +1564,43 @@ async def apply_guardrail( # 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") 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 1d46012382fb..7d454eb6fe83 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 @@ -17,6 +17,7 @@ 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