Skip to content
Closed
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
30 changes: 28 additions & 2 deletions src/strands/event_loop/event_loop.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,10 +219,36 @@ async def event_loop_cycle(
)
structured_output_context.set_forced_mode()
logger.debug("Forcing structured output tool")
await agent._append_messages(
{"role": "user", "content": [{"text": "You must format the previous response as structured output."}]}

# Use guardContent for Bedrock models with guardrails to avoid prompt attack filter
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will always apply the guardContent when a guardrail is added for bedrock. I dont think that should be the case. Instead, can we pass in a structured_output_retry_message so the user can configure the below message on their own?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure this would solve the issue. Regardless of what message the user passes, it gets injected during execution (not as an initial system prompt), which will still trigger the guardrail. I tested this by changing the message text, and it still triggered.

has_bedrock_guardrail = (
hasattr(agent.model, "config") and agent.model.config.get("guardrail_id") is not None
)

if has_bedrock_guardrail:
await agent._append_messages(
{
"role": "user",
"content": [
{
"guardContent": {
"text": {
"text": "You must format the previous response as structured output.",
"qualifiers": [],
}
}
}
],
}
)
else:
await agent._append_messages(
{
"role": "user",
"content": [{"text": "You must format the previous response as structured output."}],
}
)

events = recurse_event_loop(
agent=agent, invocation_state=invocation_state, structured_output_context=structured_output_context
)
Expand Down
92 changes: 92 additions & 0 deletions tests/strands/event_loop/test_streaming_structured_output.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
"""Tests for streaming.py with structured output support."""

import unittest.mock
from unittest.mock import AsyncMock, patch

import pytest
from pydantic import BaseModel

import strands.event_loop.streaming
from strands import Agent
from strands.event_loop.event_loop import event_loop_cycle
from strands.tools.structured_output._structured_output_context import StructuredOutputContext
from strands.tools.structured_output.structured_output_tool import StructuredOutputTool
from strands.types._events import TypedEvent

Expand Down Expand Up @@ -161,3 +165,91 @@ async def test_stream_messages_with_forced_structured_output(agenerator, alist):
assert tool_use_content is not None
assert tool_use_content["name"] == "SampleModel"
assert tool_use_content["input"] == {"name": "Alice", "age": 30}


@pytest.mark.asyncio
async def test_structured_output_system_prompt_with_or_without_guardrail(agenerator, alist):
"""Test that structured output appends correct message format based on guardrail presence."""

# Test with model WITH guardrail_id
model_with_guardrail = unittest.mock.MagicMock()
model_with_guardrail.config = {"guardrail_id": "test-guardrail-123"}

agent_with_guardrail = Agent(model=model_with_guardrail)
agent_with_guardrail.event_loop_metrics.reset_usage_metrics()
original_append = agent_with_guardrail._append_messages
agent_with_guardrail._append_messages = AsyncMock(side_effect=original_append)
structured_output_context = StructuredOutputContext(structured_output_model=SampleModel)

with patch("strands.event_loop.event_loop.recurse_event_loop") as mock_recurse:
mock_recurse.return_value = agenerator([])
model_with_guardrail.stream.return_value = agenerator(
[
{"messageStart": {"role": "assistant"}},
{"contentBlockStart": {"start": {}}},
{"contentBlockDelta": {"delta": {"text": "Response without tool use"}}},
{"contentBlockStop": {}},
{"messageStop": {"stopReason": "end_turn"}},
{
"metadata": {
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15},
"metrics": {"latencyMs": 100},
}
},
]
)

events = event_loop_cycle(
agent=agent_with_guardrail,
invocation_state={},
structured_output_context=structured_output_context,
)
await alist(events)

# Verify guardContent was appended
agent_with_guardrail._append_messages.assert_called_once()
call_args = agent_with_guardrail._append_messages.call_args[0][0]
assert call_args["role"] == "user"
assert "guardContent" in call_args["content"][0]

# Test with model WITHOUT guardrail_id
model_without_guardrail = unittest.mock.MagicMock()
model_without_guardrail.config = {}

agent_without_guardrail = Agent(model=model_without_guardrail)
agent_without_guardrail.event_loop_metrics.reset_usage_metrics()
original_append2 = agent_without_guardrail._append_messages
agent_without_guardrail._append_messages = AsyncMock(side_effect=original_append2)
structured_output_context_2 = StructuredOutputContext(structured_output_model=SampleModel)

with patch("strands.event_loop.event_loop.recurse_event_loop") as mock_recurse:
mock_recurse.return_value = agenerator([])
model_without_guardrail.stream.return_value = agenerator(
[
{"messageStart": {"role": "assistant"}},
{"contentBlockStart": {"start": {}}},
{"contentBlockDelta": {"delta": {"text": "Response without tool use"}}},
{"contentBlockStop": {}},
{"messageStop": {"stopReason": "end_turn"}},
{
"metadata": {
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15},
"metrics": {"latencyMs": 100},
}
},
]
)

events = event_loop_cycle(
agent=agent_without_guardrail,
invocation_state={},
structured_output_context=structured_output_context_2,
)
await alist(events)

# Verify regular text was appended
agent_without_guardrail._append_messages.assert_called_once()
call_args = agent_without_guardrail._append_messages.call_args[0][0]
assert call_args["role"] == "user"
assert "text" in call_args["content"][0]
assert call_args["content"][0]["text"] == "You must format the previous response as structured output."
69 changes: 69 additions & 0 deletions tests_integ/test_bedrock_guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,45 @@ def boto_session():
return boto3.Session(region_name="us-east-1")


@pytest.fixture(scope="module")
def bedrock_prompt_attack_guardrail(boto_session):
"""
Fixture that creates a guardrail with prompt attack filter (medium strength).
"""
client = boto_session.client("bedrock")

guardrail_name = "test-guardrail-prompt-attack-medium"
guardrail_id = get_guardrail_id(client, guardrail_name)

if guardrail_id:
print(f"Guardrail {guardrail_name} already exists with ID: {guardrail_id}")
else:
print(f"Creating guardrail {guardrail_name}")
response = client.create_guardrail(
name=guardrail_name,
description="Testing Guardrail with Prompt Attack Filter",
contentPolicyConfig={
"filtersConfig": [
{
"type": "PROMPT_ATTACK",
"inputStrength": "MEDIUM",
"outputStrength": "NONE",
"inputAction": "BLOCK",
"outputAction": "NONE",
"inputEnabled": True,
"outputEnabled": False,
}
]
},
blockedInputMessaging=BLOCKED_INPUT,
blockedOutputsMessaging=BLOCKED_OUTPUT,
)
guardrail_id = response.get("guardrailId")
print(f"Created test guardrail with ID: {guardrail_id}")
wait_for_guardrail_active(client, guardrail_id)
return guardrail_id


@pytest.fixture(scope="module")
def bedrock_guardrail(boto_session):
"""
Expand Down Expand Up @@ -378,3 +417,33 @@ def test_guardrail_input_intervention_properly_redacts_in_session(boto_session,

# Assert that the restored agent redacted message is equal to the original agent
assert agent.messages[0] == agent_2.messages[0]


def test_structured_output_guardrail_not_intervened(bedrock_prompt_attack_guardrail, boto_session):
"""Test that structured output works with prompt attack guardrail enabled."""
from pydantic import BaseModel

class Output(BaseModel):
age: str
user_message: str

bedrock_model = BedrockModel(
guardrail_id=bedrock_prompt_attack_guardrail,
guardrail_version="DRAFT",
boto_session=boto_session,
guardrail_trace="enabled",
)

agent = Agent(
model=bedrock_model,
system_prompt="You are an friendly assistant",
structured_output_model=Output,
)

response = agent("What are you doing?")

# Verify structured output succeeded without guardrail intervention
assert response.stop_reason == "tool_use", f"Expected tool_use, got {response.stop_reason}"
assert isinstance(response.structured_output, Output), "Should return structured output"
assert hasattr(response.structured_output, "age"), "Output should have age field"
assert hasattr(response.structured_output, "user_message"), "Output should have user_message field"
Loading