Skip to content

feat(bedrock): stream redactedContent events #229

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

feat(bedrock): stream redactedContent events #229

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
8858f9d
support redactedContent
mazyu36 May 28, 2025
4a19f45
test
mazyu36 May 28, 2025
530e2ce
fix
mazyu36 Jun 4, 2025
3af994e
Merge remote-tracking branch 'upstream/main' into feature/redactedCon…
mazyu36 Jun 4, 2025
3310114
add callback test
mazyu36 Jun 4, 2025
797c643
Merge remote-tracking branch 'strands/main' into feature/redactedContent
awsarron Jun 10, 2025
4eec5f1
format
awsarron Jun 16, 2025
1d462e7
remove unused ignore type
awsarron Jun 16, 2025
ed10796
Merge branch 'main' into feature/redactedContent
awsarron Jun 16, 2025
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
17 changes: 17 additions & 0 deletions src/strands/event_loop/streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,18 @@ def handle_content_block_delta(
**kwargs,
)

elif "redactedContent" in delta_content["reasoningContent"]:
if "redactedContent" not in state:
state["redactedContent"] = b""

state["redactedContent"] += delta_content["reasoningContent"]["redactedContent"]
callback_handler(
redactedContent=delta_content["reasoningContent"]["redactedContent"],
delta=delta_content,
reasoning=True,
**kwargs,
)

return state


Expand All @@ -170,6 +182,7 @@ def handle_content_block_stop(state: Dict[str, Any]) -> Dict[str, Any]:
current_tool_use = state["current_tool_use"]
text = state["text"]
reasoning_text = state["reasoningText"]
redacted_content = state["redactedContent"]

if current_tool_use:
if "input" not in current_tool_use:
Expand Down Expand Up @@ -207,6 +220,9 @@ def handle_content_block_stop(state: Dict[str, Any]) -> Dict[str, Any]:
}
)
state["reasoningText"] = ""
elif redacted_content:
content.append({"reasoningContent": {"redactedContent": redacted_content}})
state["redactedContent"] = b""

return state

Expand Down Expand Up @@ -279,6 +295,7 @@ def process_stream(
"current_tool_use": {},
"reasoningText": "",
"signature": "",
"redactedContent": b"",
}
state["content"] = state["message"]["content"]

Expand Down
102 changes: 102 additions & 0 deletions tests/strands/event_loop/test_streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,20 @@ def test_handle_content_block_start(chunk: ContentBlockStartEvent, exp_tool_use)
{"signature": "val"},
{"reasoning_signature": "val", "reasoning": True},
),
# Reasoning - redactedContent - New
(
{"delta": {"reasoningContent": {"redactedContent": b"encoded"}}},
{},
{"redactedContent": b"encoded"},
{"redactedContent": b"encoded", "reasoning": True},
),
# Reasoning - redactedContent - Existing
(
{"delta": {"reasoningContent": {"redactedContent": b"data"}}},
{"redactedContent": b"encoded_"},
{"redactedContent": b"encoded_data"},
{"redactedContent": b"data", "reasoning": True},
),
# Reasoning - Empty
(
{"delta": {"reasoningContent": {}}},
Expand Down Expand Up @@ -175,12 +189,14 @@ def callback_handler(**kwargs):
"current_tool_use": {"toolUseId": "123", "name": "test", "input": '{"key": "value"}'},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
{
"content": [{"toolUse": {"toolUseId": "123", "name": "test", "input": {"key": "value"}}}],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
),
# Tool Use - Missing input
Expand All @@ -190,12 +206,14 @@ def callback_handler(**kwargs):
"current_tool_use": {"toolUseId": "123", "name": "test"},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
{
"content": [{"toolUse": {"toolUseId": "123", "name": "test", "input": {}}}],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
),
# Text
Expand All @@ -205,12 +223,14 @@ def callback_handler(**kwargs):
"current_tool_use": {},
"text": "test",
"reasoningText": "",
"redactedContent": b"",
},
{
"content": [{"text": "test"}],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
),
# Reasoning
Expand All @@ -221,13 +241,32 @@ def callback_handler(**kwargs):
"text": "",
"reasoningText": "test",
"signature": "123",
"redactedContent": b"",
},
{
"content": [{"reasoningContent": {"reasoningText": {"text": "test", "signature": "123"}}}],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"signature": "123",
"redactedContent": b"",
},
),
# redactedContent
(
{
"content": [],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"encoded_data",
},
{
"content": [{"reasoningContent": {"redactedContent": b"encoded_data"}}],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
),
# Empty
Expand All @@ -237,12 +276,14 @@ def callback_handler(**kwargs):
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
{
"content": [],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"redactedContent": b"",
},
),
],
Expand Down Expand Up @@ -355,6 +396,36 @@ def test_extract_usage_metrics():
{"calls": 1},
[{"role": "user", "content": [{"text": "REDACTED"}]}],
),
(
[
{"messageStart": {"role": "assistant"}},
{
"contentBlockStart": {"start": {}},
},
{
"contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encoded_data"}}},
},
{"contentBlockStop": {}},
{
"messageStop": {"stopReason": "end_turn"},
},
{
"metadata": {
"usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1},
"metrics": {"latencyMs": 1},
}
},
],
"end_turn",
{
"role": "assistant",
"content": [{"reasoningContent": {"redactedContent": b"encoded_data"}}],
},
{"inputTokens": 1, "outputTokens": 1, "totalTokens": 1},
{"latencyMs": 1},
{"calls": 1},
[{"role": "user", "content": [{"text": "Some input!"}]}],
),
],
)
def test_process_stream(
Expand Down Expand Up @@ -422,3 +493,34 @@ def callback_handler(**kwargs):
None,
"test prompt",
)


def test_process_stream_redacted_content_callback():
callback_args = []

def callback_handler(**kwargs):
callback_args.append(kwargs)

response = [
{"messageStart": {"role": "assistant"}},
{"contentBlockStart": {"start": {}}},
{"contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encoded_data_1"}}}},
{"contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encoded_data_2"}}}},
{"contentBlockStop": {}},
{"messageStop": {"stopReason": "end_turn"}},
]

messages = [{"role": "user", "content": [{"text": "Some input!"}]}]

strands.event_loop.streaming.process_stream(response, callback_handler, messages)

redacted_callbacks = [args for args in callback_args if "redactedContent" in args]
assert len(redacted_callbacks) == 2

assert redacted_callbacks[0]["redactedContent"] == b"encoded_data_1"
Copy link
Member

Choose a reason for hiding this comment

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

Can we do an assert redacted_callbacks[0] == {... entire object } to make this more readable?

assert redacted_callbacks[0]["delta"] == {"reasoningContent": {"redactedContent": b"encoded_data_1"}}
assert redacted_callbacks[0]["reasoning"] is True

assert redacted_callbacks[1]["redactedContent"] == b"encoded_data_2"
assert redacted_callbacks[1]["delta"] == {"reasoningContent": {"redactedContent": b"encoded_data_2"}}
assert redacted_callbacks[1]["reasoning"] is True