diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7b12d1e939f..983cf9bfa1b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -35,7 +35,9 @@ from litellm.types.llms.openai import ( ChatCompletionFunctionMessage, ChatCompletionImageObject, ChatCompletionImageUrlObject, + ChatCompletionRedactedThinkingBlock, ChatCompletionTextObject, + ChatCompletionThinkingBlock, ChatCompletionToolCallFunctionChunk, ChatCompletionToolMessage, ChatCompletionUserMessage, @@ -4501,7 +4503,7 @@ class BedrockConverseMessagesProcessor: ) _assistant_content = assistant_message_block.get("content", None) thinking_blocks = cast( - list[ChatCompletionThinkingBlock] | None, + list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, assistant_message_block.get("thinking_blocks"), ) @@ -4520,9 +4522,11 @@ class BedrockConverseMessagesProcessor: assistants_parts: list[BedrockContentBlock] = [] for element in _assistant_content: if isinstance(element, dict): - if element["type"] == "thinking": + if element["type"] in ("thinking", "redacted_thinking"): thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[cast(ChatCompletionThinkingBlock, element)] + thinking_blocks=[ + cast(ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, element) + ] ) assistants_parts = ( BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -4586,22 +4590,33 @@ class BedrockConverseMessagesProcessor: @staticmethod def translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks: list[ChatCompletionThinkingBlock], + thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], ) -> list[BedrockContentBlock]: reasoning_content_blocks: Final[list[BedrockContentBlock]] = [] for thinking_block in thinking_blocks: - reasoning_text = thinking_block.get("thinking") - reasoning_signature = thinking_block.get("signature") - text_block = BedrockConverseReasoningTextBlock( - text=reasoning_text or "", - ) - if reasoning_signature is not None: - text_block["signature"] = reasoning_signature - reasoning_content_block = BedrockConverseReasoningContentBlock( - reasoningText=text_block, - ) - bedrock_content_block = BedrockContentBlock(reasoningContent=reasoning_content_block) - reasoning_content_blocks.append(bedrock_content_block) + if thinking_block.get("type") == "redacted_thinking" or "data" in thinking_block: + redacted_data = thinking_block.get("data") + if redacted_data is not None: + if isinstance(redacted_data, bytes): + redacted_data = base64.b64encode(redacted_data).decode("utf-8") + reasoning_content_block = BedrockConverseReasoningContentBlock( + redactedContent=redacted_data, + ) + bedrock_content_block = BedrockContentBlock(reasoningContent=reasoning_content_block) + reasoning_content_blocks.append(bedrock_content_block) + else: + reasoning_text = thinking_block.get("thinking") + reasoning_signature = thinking_block.get("signature") + text_block = BedrockConverseReasoningTextBlock( + text=reasoning_text or "", + ) + if reasoning_signature is not None: + text_block["signature"] = reasoning_signature + reasoning_content_block = BedrockConverseReasoningContentBlock( + reasoningText=text_block, + ) + bedrock_content_block = BedrockContentBlock(reasoningContent=reasoning_content_block) + reasoning_content_blocks.append(bedrock_content_block) return reasoning_content_blocks @staticmethod @@ -4876,7 +4891,7 @@ def _bedrock_converse_messages_pt( ) _assistant_content = assistant_message_block.get("content", None) thinking_blocks = cast( - list[ChatCompletionThinkingBlock] | None, + list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, assistant_message_block.get("thinking_blocks"), ) @@ -4895,10 +4910,12 @@ def _bedrock_converse_messages_pt( assistants_parts: list[BedrockContentBlock] = [] for element in _assistant_content: if isinstance(element, dict): - if element["type"] == "thinking": + if element["type"] in ("thinking", "redacted_thinking"): thinking_block = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[cast(ChatCompletionThinkingBlock, element)] + thinking_blocks=[ + cast(ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, element) + ] ) ) assistants_parts = ( diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index e4c41c3ee5b..1c07627f69e 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -89,7 +89,7 @@ class BedrockConverseReasoningTextBlock(TypedDict, total=False): class BedrockConverseReasoningContentBlock(TypedDict, total=False): reasoningText: BedrockConverseReasoningTextBlock - redactedContent: str + redactedContent: str | bytes class BedrockConverseReasoningContentBlockDelta(TypedDict, total=False): diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 26124ac24de..0a365a1b431 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2642,9 +2642,11 @@ def test_bedrock_tool_call_invoke_multiple_normal_tools(): assert result[1]["toolUse"]["toolUseId"] == "call_2" -# ======================================================================== +# ================================================================= + # Tool result deduplication tests (Case D in sanitize_messages_for_tool_calling) -# ======================================================================== +# ================================================================= + def test_sanitize_messages_deduplicates_tool_results(): @@ -3895,3 +3897,118 @@ def test_anthropic_messages_pt_drops_a_system_message_with_no_text(): result = anthropic_messages_pt(messages=messages, model="claude-opus-4-8", llm_provider="anthropic") assert [m["role"] for m in result] == ["user", "assistant"] + + +def test_bedrock_converse_messages_pt_preserves_redacted_thinking_blocks(): + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt + + # Case 1: redacted_thinking in assistant message 'thinking_blocks' + messages_1 = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi there!", + "thinking_blocks": [{"type": "redacted_thinking", "data": "redacted_secret_data_123"}], + }, + {"role": "user", "content": "Followup"}, + ] + result_1 = _bedrock_converse_messages_pt( + messages=messages_1, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks_1 = result_1[1]["content"] + assert any( + block.get("reasoningContent", {}).get("redactedContent") == "redacted_secret_data_123" + for block in assistant_blocks_1 + ), f"Expected redactedContent in assistant blocks, got: {assistant_blocks_1}" + + # Case 2: redacted_thinking in assistant message 'content' list + messages_2 = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "redacted_secret_data_456"}, + {"type": "text", "text": "Hi there!"}, + ], + }, + {"role": "user", "content": "Followup"}, + ] + result_2 = _bedrock_converse_messages_pt( + messages=messages_2, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks_2 = result_2[1]["content"] + assert any( + block.get("reasoningContent", {}).get("redactedContent") == "redacted_secret_data_456" + for block in assistant_blocks_2 + ), f"Expected redactedContent in assistant blocks, got: {assistant_blocks_2}" + + +def test_bedrock_converse_messages_pt_handles_bytes_redacted_thinking(): + import base64 + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt + + raw_bytes = b"binary_redacted_payload" + messages = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": raw_bytes}, + {"type": "text", "text": "Hi there!"}, + ], + }, + ] + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks = result[1]["content"] + expected_b64 = base64.b64encode(raw_bytes).decode("utf-8") + assert any( + block.get("reasoningContent", {}).get("redactedContent") == expected_b64 for block in assistant_blocks + ), f"Expected base64-encoded redactedContent, got: {assistant_blocks}" + + +def test_bedrock_converse_redacted_thinking_with_tool_calls_ordering(): + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt + + messages = [ + {"role": "user", "content": "Calculate something"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "redacted_jwt_tokens"}, + {"type": "text", "text": "Calling tool now"}, + ], + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "call_abc123", "content": "42"}], + }, + ] + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks = result[1]["content"] + + assert "reasoningContent" in assistant_blocks[0] + assert assistant_blocks[0]["reasoningContent"]["redactedContent"] == "redacted_jwt_tokens" + assert "text" in assistant_blocks[1] + assert assistant_blocks[1]["text"] == "Calling tool now" + assert "toolUse" in assistant_blocks[2] + assert assistant_blocks[2]["toolUse"]["toolUseId"] == "call_abc123" + diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 499096621c5..793937d2c15 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -7913,3 +7913,119 @@ def test_supports_sampling_params_prefixed_and_anthropic_fallback(monkeypatch: p ) assert AmazonConverseConfig._supports_sampling_params("custom-test-reasoning-model") is False assert AmazonConverseConfig._supports_sampling_params("anthropic.claude-custom-unregistered") is True + + +def test_bedrock_converse_messages_pt_preserves_redacted_thinking_blocks(): + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt + + # Case 1: redacted_thinking in assistant message 'thinking_blocks' + messages_1 = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi there!", + "thinking_blocks": [ + {"type": "redacted_thinking", "data": "redacted_secret_data_123"} + ], + }, + {"role": "user", "content": "Followup"}, + ] + result_1 = _bedrock_converse_messages_pt( + messages=messages_1, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks_1 = result_1[1]["content"] + assert any( + block.get("reasoningContent", {}).get("redactedContent") == "redacted_secret_data_123" + for block in assistant_blocks_1 + ), f"Expected redactedContent in assistant blocks, got: {assistant_blocks_1}" + + # Case 2: redacted_thinking in assistant message 'content' list + messages_2 = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "redacted_secret_data_456"}, + {"type": "text", "text": "Hi there!"}, + ], + }, + {"role": "user", "content": "Followup"}, + ] + result_2 = _bedrock_converse_messages_pt( + messages=messages_2, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks_2 = result_2[1]["content"] + assert any( + block.get("reasoningContent", {}).get("redactedContent") == "redacted_secret_data_456" + for block in assistant_blocks_2 + ), f"Expected redactedContent in assistant blocks, got: {assistant_blocks_2}" + + +@pytest.mark.asyncio +async def test_bedrock_converse_messages_pt_async_preserves_redacted_thinking_blocks(): + from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor + + messages = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "async_redacted_data_789"}, + {"type": "text", "text": "Async response"}, + ], + }, + ] + result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks = result[1]["content"] + assert any( + block.get("reasoningContent", {}).get("redactedContent") == "async_redacted_data_789" + for block in assistant_blocks + ), f"Expected redactedContent in async assistant blocks, got: {assistant_blocks}" + + +def test_bedrock_converse_redacted_thinking_with_tool_calls_ordering(): + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt + + messages = [ + {"role": "user", "content": "Calculate something"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "redacted_jwt_tokens"}, + {"type": "text", "text": "Calling tool now"}, + ], + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "call_abc123", "content": "42"}], + }, + ] + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + llm_provider="bedrock", + ) + assistant_blocks = result[1]["content"] + + # Verify sort order: reasoningContent must be first, followed by text, followed by toolUse + assert "reasoningContent" in assistant_blocks[0] + assert assistant_blocks[0]["reasoningContent"]["redactedContent"] == "redacted_jwt_tokens" + assert "text" in assistant_blocks[1] + assert assistant_blocks[1]["text"] == "Calling tool now" + assert "toolUse" in assistant_blocks[2] + assert assistant_blocks[2]["toolUse"]["toolUseId"] == "call_abc123"