scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(bedrock): preserve redacted_thinking blocks on Converse replay (#43009)

Browse source
This commit is contained in:
Raj Ojha 2026-09-25 02:58:32 +05:30
parent dd63637322
commit 5a460a4441
4 changed files with 272 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -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 = (

2
litellm/types/llms/bedrock.py
View file

@ -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):

121
tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -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"

116
tests/unit/llms/bedrock/chat/test_converse_transformation.py
View file

@ -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"