From 8858f9da78f6d70c4f475df86ffc1ceec182d10a Mon Sep 17 00:00:00 2001 From: Yuki Matsuda <13781813+mazyu36@users.noreply.github.com> Date: Wed, 28 May 2025 22:09:03 +0900 Subject: [PATCH 1/4] support redactedContent --- src/strands/event_loop/streaming.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/strands/event_loop/streaming.py b/src/strands/event_loop/streaming.py index 6e8a806fd..f0b6ef743 100644 --- a/src/strands/event_loop/streaming.py +++ b/src/strands/event_loop/streaming.py @@ -152,6 +152,18 @@ def handle_content_block_delta( reasoning=True, **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 @@ -207,6 +219,16 @@ def handle_content_block_stop(state: Dict[str, Any]) -> Dict[str, Any]: } ) state["reasoningText"] = "" + state["signature"] = "" + elif "redactedContent" in state and state["redactedContent"]: + content.append( + { + "reasoningContent": { + "redactedContent": state["redactedContent"] + } + } + ) + state["redactedContent"] = b"" return state @@ -279,6 +301,7 @@ def process_stream( "current_tool_use": {}, "reasoningText": "", "signature": "", + "redactedContent": b"", } state["content"] = state["message"]["content"] From 4a19f45e50c17cbb2f0bb974fc4018e85e002b42 Mon Sep 17 00:00:00 2001 From: Yuki Matsuda <13781813+mazyu36@users.noreply.github.com> Date: Thu, 29 May 2025 00:15:06 +0900 Subject: [PATCH 2/4] test --- src/strands/event_loop/streaming.py | 11 +--- tests/strands/event_loop/test_streaming.py | 61 ++++++++++++++++++++++ 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/src/strands/event_loop/streaming.py b/src/strands/event_loop/streaming.py index f0b6ef743..df0a9308f 100644 --- a/src/strands/event_loop/streaming.py +++ b/src/strands/event_loop/streaming.py @@ -152,7 +152,7 @@ def handle_content_block_delta( reasoning=True, **kwargs, ) - + elif "redactedContent" in delta_content["reasoningContent"]: if "redactedContent" not in state: state["redactedContent"] = b"" @@ -219,15 +219,8 @@ def handle_content_block_stop(state: Dict[str, Any]) -> Dict[str, Any]: } ) state["reasoningText"] = "" - state["signature"] = "" elif "redactedContent" in state and state["redactedContent"]: - content.append( - { - "reasoningContent": { - "redactedContent": state["redactedContent"] - } - } - ) + content.append({"reasoningContent": {"redactedContent": state["redactedContent"]}}) state["redactedContent"] = b"" return state diff --git a/tests/strands/event_loop/test_streaming.py b/tests/strands/event_loop/test_streaming.py index c24e7e48a..c5f6ec2b7 100644 --- a/tests/strands/event_loop/test_streaming.py +++ b/tests/strands/event_loop/test_streaming.py @@ -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"encrypted"}}}, + {}, + {"redactedContent": b"encrypted"}, + {"redactedContent": b"encrypted", "reasoning": True}, + ), + # Reasoning - redactedContent - Existing + ( + {"delta": {"reasoningContent": {"redactedContent": b"data"}}}, + {"redactedContent": b"encrypted_"}, + {"redactedContent": b"encrypted_data"}, + {"redactedContent": b"data", "reasoning": True}, + ), # Reasoning - Empty ( {"delta": {"reasoningContent": {}}}, @@ -230,6 +244,23 @@ def callback_handler(**kwargs): "signature": "123", }, ), + # redactedContent + ( + { + "content": [], + "current_tool_use": {}, + "text": "", + "reasoningText": "", + "redactedContent": b"encrypted_data", + }, + { + "content": [{"reasoningContent": {"redactedContent": b"encrypted_data"}}], + "current_tool_use": {}, + "text": "", + "reasoningText": "", + "redactedContent": b"", + }, + ), # Empty ( { @@ -355,6 +386,36 @@ def test_extract_usage_metrics(): {"calls": 1}, [{"role": "user", "content": [{"text": "REDACTED"}]}], ), + ( + [ + {"messageStart": {"role": "assistant"}}, + { + "contentBlockStart": {"start": {}}, + }, + { + "contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encrypted_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"encrypted_data"}}], + }, + {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1}, + {"latencyMs": 1}, + {"calls": 1}, + [{"role": "user", "content": [{"text": "Some input!"}]}], + ), ], ) def test_process_stream( From 530e2ce5b498e52845a776a9c5c98bb66347f5b5 Mon Sep 17 00:00:00 2001 From: Yuki Matsuda <13781813+mazyu36@users.noreply.github.com> Date: Wed, 4 Jun 2025 12:58:07 +0900 Subject: [PATCH 3/4] fix --- src/strands/event_loop/streaming.py | 5 ++-- tests/strands/event_loop/test_streaming.py | 28 +++++++++++++++------- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/src/strands/event_loop/streaming.py b/src/strands/event_loop/streaming.py index df0a9308f..9ee6a0a61 100644 --- a/src/strands/event_loop/streaming.py +++ b/src/strands/event_loop/streaming.py @@ -182,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: @@ -219,8 +220,8 @@ def handle_content_block_stop(state: Dict[str, Any]) -> Dict[str, Any]: } ) state["reasoningText"] = "" - elif "redactedContent" in state and state["redactedContent"]: - content.append({"reasoningContent": {"redactedContent": state["redactedContent"]}}) + elif redacted_content: + content.append({"reasoningContent": {"redactedContent": redacted_content}}) state["redactedContent"] = b"" return state diff --git a/tests/strands/event_loop/test_streaming.py b/tests/strands/event_loop/test_streaming.py index c5f6ec2b7..00a8765ae 100644 --- a/tests/strands/event_loop/test_streaming.py +++ b/tests/strands/event_loop/test_streaming.py @@ -134,16 +134,16 @@ def test_handle_content_block_start(chunk: ContentBlockStartEvent, exp_tool_use) ), # Reasoning - redactedContent - New ( - {"delta": {"reasoningContent": {"redactedContent": b"encrypted"}}}, + {"delta": {"reasoningContent": {"redactedContent": b"encoded"}}}, {}, - {"redactedContent": b"encrypted"}, - {"redactedContent": b"encrypted", "reasoning": True}, + {"redactedContent": b"encoded"}, + {"redactedContent": b"encoded", "reasoning": True}, ), # Reasoning - redactedContent - Existing ( {"delta": {"reasoningContent": {"redactedContent": b"data"}}}, - {"redactedContent": b"encrypted_"}, - {"redactedContent": b"encrypted_data"}, + {"redactedContent": b"encoded_"}, + {"redactedContent": b"encoded_data"}, {"redactedContent": b"data", "reasoning": True}, ), # Reasoning - Empty @@ -189,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 @@ -204,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 @@ -219,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 @@ -235,6 +241,7 @@ def callback_handler(**kwargs): "text": "", "reasoningText": "test", "signature": "123", + "redactedContent": b"", }, { "content": [{"reasoningContent": {"reasoningText": {"text": "test", "signature": "123"}}}], @@ -242,6 +249,7 @@ def callback_handler(**kwargs): "text": "", "reasoningText": "", "signature": "123", + "redactedContent": b"", }, ), # redactedContent @@ -251,10 +259,10 @@ def callback_handler(**kwargs): "current_tool_use": {}, "text": "", "reasoningText": "", - "redactedContent": b"encrypted_data", + "redactedContent": b"encoded_data", }, { - "content": [{"reasoningContent": {"redactedContent": b"encrypted_data"}}], + "content": [{"reasoningContent": {"redactedContent": b"encoded_data"}}], "current_tool_use": {}, "text": "", "reasoningText": "", @@ -268,12 +276,14 @@ def callback_handler(**kwargs): "current_tool_use": {}, "text": "", "reasoningText": "", + "redactedContent": b"", }, { "content": [], "current_tool_use": {}, "text": "", "reasoningText": "", + "redactedContent": b"", }, ), ], @@ -393,7 +403,7 @@ def test_extract_usage_metrics(): "contentBlockStart": {"start": {}}, }, { - "contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encrypted_data"}}}, + "contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encoded_data"}}}, }, {"contentBlockStop": {}}, { @@ -409,7 +419,7 @@ def test_extract_usage_metrics(): "end_turn", { "role": "assistant", - "content": [{"reasoningContent": {"redactedContent": b"encrypted_data"}}], + "content": [{"reasoningContent": {"redactedContent": b"encoded_data"}}], }, {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1}, {"latencyMs": 1}, From 3310114c2d65a8084e855ac58f4f8a716b86331e Mon Sep 17 00:00:00 2001 From: Yuki Matsuda <13781813+mazyu36@users.noreply.github.com> Date: Wed, 4 Jun 2025 16:03:52 +0900 Subject: [PATCH 4/4] add callback test --- tests/strands/event_loop/test_streaming.py | 31 ++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/strands/event_loop/test_streaming.py b/tests/strands/event_loop/test_streaming.py index 00a8765ae..1d739bf97 100644 --- a/tests/strands/event_loop/test_streaming.py +++ b/tests/strands/event_loop/test_streaming.py @@ -493,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" + 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