diff --git a/src/strands/event_loop/streaming.py b/src/strands/event_loop/streaming.py index 183fe1ec8..f24bd2a76 100644 --- a/src/strands/event_loop/streaming.py +++ b/src/strands/event_loop/streaming.py @@ -10,6 +10,7 @@ ModelStopReason, ModelStreamChunkEvent, ModelStreamEvent, + ReasoningRedactedContentStreamEvent, ReasoningSignatureStreamEvent, ReasoningTextStreamEvent, TextStreamEvent, @@ -170,6 +171,10 @@ def handle_content_block_delta( delta=delta_content, ) + elif redacted_content := delta_content["reasoningContent"].get("redactedContent"): + state["redactedContent"] = state.get("redactedContent", b"") + redacted_content + typed_event = ReasoningRedactedContentStreamEvent(redacted_content=redacted_content, delta=delta_content) + return state, typed_event @@ -188,6 +193,7 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]: text = state["text"] reasoning_text = state["reasoningText"] citations_content = state["citationsContent"] + redacted_content = state.get("redactedContent") if current_tool_use: if "input" not in current_tool_use: @@ -231,6 +237,9 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]: content.append(content_block) state["reasoningText"] = "" + elif redacted_content: + content.append({"reasoningContent": {"redactedContent": redacted_content}}) + state["redactedContent"] = b"" return state diff --git a/src/strands/types/_events.py b/src/strands/types/_events.py index ccdab1846..3d0f1d0f0 100644 --- a/src/strands/types/_events.py +++ b/src/strands/types/_events.py @@ -169,6 +169,14 @@ def __init__(self, delta: ContentBlockDelta, reasoning_text: str | None) -> None super().__init__({"reasoningText": reasoning_text, "delta": delta, "reasoning": True}) +class ReasoningRedactedContentStreamEvent(ModelStreamEvent): + """Event emitted during redacted content streaming.""" + + def __init__(self, delta: ContentBlockDelta, redacted_content: bytes | None) -> None: + """Initialize with delta and redacted content.""" + super().__init__({"reasoningRedactedContent": redacted_content, "delta": delta, "reasoning": True}) + + class ReasoningSignatureStreamEvent(ModelStreamEvent): """Event emitted during reasoning signature streaming.""" diff --git a/tests/fixtures/mocked_model_provider.py b/tests/fixtures/mocked_model_provider.py index 2a397bb18..c05089f34 100644 --- a/tests/fixtures/mocked_model_provider.py +++ b/tests/fixtures/mocked_model_provider.py @@ -72,6 +72,10 @@ def map_agent_message_to_events(self, agent_message: Union[Message, RedactionMes stop_reason = "guardrail_intervened" else: for content in agent_message["content"]: + if "reasoningContent" in content: + yield {"contentBlockStart": {"start": {}}} + yield {"contentBlockDelta": {"delta": {"reasoningContent": content["reasoningContent"]}}} + yield {"contentBlockStop": {}} if "text" in content: yield {"contentBlockStart": {"start": {}}} yield {"contentBlockDelta": {"delta": {"text": content["text"]}}} diff --git a/tests/strands/agent/hooks/test_agent_events.py b/tests/strands/agent/hooks/test_agent_events.py index 01bfc5409..9b3646144 100644 --- a/tests/strands/agent/hooks/test_agent_events.py +++ b/tests/strands/agent/hooks/test_agent_events.py @@ -387,6 +387,84 @@ async def test_stream_e2e_throttle_and_redact(alist, mock_sleep): assert typed_events == [] +@pytest.mark.asyncio +async def test_stream_e2e_reasoning_redacted_content(alist): + mock_provider = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"reasoningContent": {"redactedContent": b"test_redacted_data"}}, + {"text": "Response with redacted reasoning"}, + ], + }, + ] + ) + + mock_callback = unittest.mock.Mock() + agent = Agent(model=mock_provider, callback_handler=mock_callback) + + stream = agent.stream_async("Test redacted content") + + tru_events = await alist(stream) + exp_events = [ + {"init_event_loop": True}, + {"start": True}, + {"start_event_loop": True}, + {"event": {"messageStart": {"role": "assistant"}}}, + {"event": {"contentBlockStart": {"start": {}}}}, + {"event": {"contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"test_redacted_data"}}}}}, + { + **any_props, + "reasoningRedactedContent": b"test_redacted_data", + "delta": {"reasoningContent": {"redactedContent": b"test_redacted_data"}}, + "reasoning": True, + }, + {"event": {"contentBlockStop": {}}}, + {"event": {"contentBlockStart": {"start": {}}}}, + {"event": {"contentBlockDelta": {"delta": {"text": "Response with redacted reasoning"}}}}, + { + **any_props, + "data": "Response with redacted reasoning", + "delta": {"text": "Response with redacted reasoning"}, + }, + {"event": {"contentBlockStop": {}}}, + {"event": {"messageStop": {"stopReason": "end_turn"}}}, + { + "message": { + "content": [ + {"reasoningContent": {"redactedContent": b"test_redacted_data"}}, + {"text": "Response with redacted reasoning"}, + ], + "role": "assistant", + } + }, + { + "result": AgentResult( + stop_reason="end_turn", + message={ + "content": [ + {"reasoningContent": {"redactedContent": b"test_redacted_data"}}, + {"text": "Response with redacted reasoning"}, + ], + "role": "assistant", + }, + metrics=ANY, + state={}, + ) + }, + ] + assert tru_events == exp_events + + exp_calls = [call(**event) for event in exp_events] + act_calls = mock_callback.call_args_list + assert act_calls == exp_calls + + # Ensure that all events coming out of the agent are *not* typed events + typed_events = [event for event in tru_events if isinstance(event, TypedEvent)] + assert typed_events == [] + + @pytest.mark.asyncio async def test_event_loop_cycle_text_response_throttling_early_end( agenerator, diff --git a/tests/strands/event_loop/test_streaming.py b/tests/strands/event_loop/test_streaming.py index 32d1889e5..1de957619 100644 --- a/tests/strands/event_loop/test_streaming.py +++ b/tests/strands/event_loop/test_streaming.py @@ -131,6 +131,20 @@ def test_handle_content_block_start(chunk: ContentBlockStartEvent, exp_tool_use) {"signature": "val"}, {"reasoning_signature": "val", "reasoning": True}, ), + # Reasoning - redactedContent - New + pytest.param( + {"delta": {"reasoningContent": {"redactedContent": b"encoded"}}}, + {}, + {"redactedContent": b"encoded"}, + {"reasoningRedactedContent": b"encoded", "reasoning": True}, + ), + # Reasoning - redactedContent - Existing + pytest.param( + {"delta": {"reasoningContent": {"redactedContent": b"data"}}}, + {"redactedContent": b"encoded_"}, + {"redactedContent": b"encoded_data"}, + {"reasoningRedactedContent": b"data", "reasoning": True}, + ), # Reasoning - Empty ( {"delta": {"reasoningContent": {}}}, @@ -167,6 +181,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, { "content": [{"toolUse": {"toolUseId": "123", "name": "test", "input": {"key": "value"}}}], @@ -174,6 +189,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, ), # Tool Use - Missing input @@ -184,6 +200,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, { "content": [{"toolUse": {"toolUseId": "123", "name": "test", "input": {}}}], @@ -191,6 +208,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, ), # Text @@ -201,6 +219,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "test", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, { "content": [{"text": "test"}], @@ -208,6 +227,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, ), # Citations @@ -218,6 +238,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [{"citations": [{"text": "test", "source": "test"}]}], + "redactedContent": b"", }, { "content": [], @@ -225,6 +246,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [{"citations": [{"text": "test", "source": "test"}]}], + "redactedContent": b"", }, ), # Reasoning @@ -236,6 +258,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "reasoningText": "test", "signature": "123", "citationsContent": [], + "redactedContent": b"", }, { "content": [{"reasoningContent": {"reasoningText": {"text": "test", "signature": "123"}}}], @@ -244,6 +267,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "reasoningText": "", "signature": "123", "citationsContent": [], + "redactedContent": b"", }, ), # Reasoning without signature @@ -254,6 +278,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "test", "citationsContent": [], + "redactedContent": b"", }, { "content": [{"reasoningContent": {"reasoningText": {"text": "test"}}}], @@ -261,6 +286,26 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", + }, + ), + # redactedContent + ( + { + "content": [], + "current_tool_use": {}, + "text": "", + "reasoningText": "", + "redactedContent": b"encoded_data", + "citationsContent": [], + }, + { + "content": [{"reasoningContent": {"redactedContent": b"encoded_data"}}], + "current_tool_use": {}, + "text": "", + "reasoningText": "", + "redactedContent": b"", + "citationsContent": [], }, ), # Empty @@ -271,6 +316,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, { "content": [], @@ -278,6 +324,7 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, state, exp_up "text": "", "reasoningText": "", "citationsContent": [], + "redactedContent": b"", }, ), ], @@ -449,6 +496,23 @@ def test_extract_usage_metrics_with_cache_tokens(): }, ], ), + ], +) +@pytest.mark.asyncio +async def test_process_stream(response, exp_events, agenerator, alist): + stream = strands.event_loop.streaming.process_stream(agenerator(response)) + + tru_events = await alist(stream) + assert tru_events == exp_events + + # Ensure that we're getting typed events coming out of process_stream + non_typed_events = [event for event in tru_events if not isinstance(event, TypedEvent)] + assert non_typed_events == [] + + +@pytest.mark.parametrize( + ("response", "exp_events"), + [ # Redacted Message ( [ @@ -471,92 +535,116 @@ def test_extract_usage_metrics_with_cache_tokens(): }, { "metadata": { - "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1}, + "usage": { + "inputTokens": 1, + "outputTokens": 1, + "totalTokens": 1, + }, "metrics": {"latencyMs": 1}, } }, ], [ + {"event": {"messageStart": {"role": "assistant"}}}, + {"event": {"contentBlockStart": {"start": {}}}}, + {"event": {"contentBlockDelta": {"delta": {"text": "Hello!"}}}}, + {"data": "Hello!", "delta": {"text": "Hello!"}}, + {"event": {"contentBlockStop": {}}}, + {"event": {"messageStop": {"stopReason": "guardrail_intervened"}}}, { "event": { - "messageStart": { - "role": "assistant", - }, - }, + "redactContent": { + "redactUserContentMessage": "REDACTED", + "redactAssistantContentMessage": "REDACTED.", + } + } }, { "event": { - "contentBlockStart": { - "start": {}, - }, - }, + "metadata": { + "usage": { + "inputTokens": 1, + "outputTokens": 1, + "totalTokens": 1, + }, + "metrics": {"latencyMs": 1}, + } + } }, { - "event": { - "contentBlockDelta": { - "delta": { - "text": "Hello!", - }, - }, - }, + "stop": ( + "guardrail_intervened", + {"role": "assistant", "content": [{"text": "REDACTED."}]}, + {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1}, + {"latencyMs": 1}, + ) }, + ], + ), + ( + [ + {"messageStart": {"role": "assistant"}}, { - "data": "Hello!", - "delta": { - "text": "Hello!", - }, + "contentBlockStart": {"start": {}}, }, { - "event": { - "contentBlockStop": {}, - }, + "contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encoded_data"}}}, }, + {"contentBlockStop": {}}, { - "event": { - "messageStop": { - "stopReason": "guardrail_intervened", - }, - }, + "messageStop": {"stopReason": "end_turn"}, }, { - "event": { - "redactContent": { - "redactAssistantContentMessage": "REDACTED.", - "redactUserContentMessage": "REDACTED", + "metadata": { + "usage": { + "inputTokens": 1, + "outputTokens": 1, + "totalTokens": 1, }, - }, + "metrics": {"latencyMs": 1}, + } + }, + ], + [ + {"event": {"messageStart": {"role": "assistant"}}}, + {"event": {"contentBlockStart": {"start": {}}}}, + {"event": {"contentBlockDelta": {"delta": {"reasoningContent": {"redactedContent": b"encoded_data"}}}}}, + { + "reasoningRedactedContent": b"encoded_data", + "delta": {"reasoningContent": {"redactedContent": b"encoded_data"}}, + "reasoning": True, }, + {"event": {"contentBlockStop": {}}}, + {"event": {"messageStop": {"stopReason": "end_turn"}}}, { "event": { "metadata": { - "metrics": { - "latencyMs": 1, - }, "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 1, }, - }, - }, + "metrics": {"latencyMs": 1}, + } + } }, { "stop": ( - "guardrail_intervened", + "end_turn", { "role": "assistant", - "content": [{"text": "REDACTED."}], + "content": [{"reasoningContent": {"redactedContent": b"encoded_data"}}], }, {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1}, {"latencyMs": 1}, - ), + ) }, ], ), ], ) @pytest.mark.asyncio -async def test_process_stream(response, exp_events, agenerator, alist): +async def test_process_stream_redacted(response, exp_events, agenerator, alist): stream = strands.event_loop.streaming.process_stream(agenerator(response)) tru_events = await alist(stream) diff --git a/tests_integ/models/test_model_bedrock.py b/tests_integ/models/test_model_bedrock.py index 00107411a..9dff66fde 100644 --- a/tests_integ/models/test_model_bedrock.py +++ b/tests_integ/models/test_model_bedrock.py @@ -244,3 +244,26 @@ def test_structured_output_multi_modal_input(streaming_agent, yellow_img, yellow tru_color = streaming_agent.structured_output(type(yellow_color), content) exp_color = yellow_color assert tru_color == exp_color + + +def test_redacted_content_handling(): + """Test redactedContent handling with thinking mode.""" + bedrock_model = BedrockModel( + model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + additional_request_fields={ + "thinking": { + "type": "enabled", + "budget_tokens": 2000, + } + }, + ) + + agent = Agent(name="test_redact", model=bedrock_model) + # https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#example-working-with-redacted-thinking-blocks + result = agent( + "ANTHROPIC_MAGIC_STRING_TRIGGER_REDACTED_THINKING_46C9A13E193C177646C7398A98432ECCCE4C1253D5E2D82641AC0E52CC2876CB" + ) + + assert "reasoningContent" in result.message["content"][0] + assert "redactedContent" in result.message["content"][0]["reasoningContent"] + assert isinstance(result.message["content"][0]["reasoningContent"]["redactedContent"], bytes)