From c770756cf3c82292e2bc561caf27342b6287c9a8 Mon Sep 17 00:00:00 2001 From: shivam Date: Mon, 20 Apr 2026 19:42:51 -0700 Subject: [PATCH 1/2] 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 25b46cf3641..77b2f466f2a 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") From 6beba97d2019c35129405e262b75d73d8276a809 Mon Sep 17 00:00:00 2001 From: shivam Date: Mon, 20 Apr 2026 19:53:49 -0700 Subject: [PATCH 2/2] test(bedrock_guardrails): assert apply_guardrail maps response to OUTPUT source Add regression tests that mock make_bedrock_api_request and verify input_type=request uses source=INPUT with user messages, and input_type=response uses source=OUTPUT with synthetic ModelResponse. Made-with: Cursor --- .../test_bedrock_guardrails.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) 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 1d46012382f..7d454eb6fe8 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 @@ 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