Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 37 additions & 4 deletions litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
CallTypesLiteral,
Choices,
GuardrailStatus,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
Expand Down Expand Up @@ -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,
Comment thread
shivamrawat1 marked this conversation as resolved.
)

# Apply any masking that was applied by the guardrail
output_list = bedrock_response.get("output")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
BedrockGuardrail,
_redact_pii_matches,
)
from litellm.types.utils import ModelResponse


@pytest.mark.asyncio
Expand Down Expand Up @@ -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
Expand Down
Loading