From eed6b2bb751188f8ce3e436fc7b4e467ee7cf3b4 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Thu, 24 Sep 2026 19:59:29 +0000 Subject: [PATCH] fix(bedrock): keep count_tokens body key order and narrow the thinking display warning Rebuild the InvokeModel count body in place so keys stay in request order and only messages is rewritten. Gate the thinking display coercion on the dict check before calling the normalizer. Pin the already-empty content list next to an offender, malformed message entries next to an offender, and the count body key order in the mapped tests --- .../bedrock/count_tokens/transformation.py | 3 +- .../anthropic_claude3_transformation.py | 21 +++++++----- .../test_anthropic_claude3_transformation.py | 12 +++++-- .../llms/bedrock/test_bedrock_common_utils.py | 32 +++++++++++++++++++ ...est_bedrock_count_tokens_transformation.py | 5 +-- 5 files changed, 59 insertions(+), 14 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index f1cc1a33f94..df042a4da73 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -208,8 +208,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): else messages ) body_data: Final = { # mutable-ok: outbound JSON body, defaults are set below like before - **{k: v for k, v in request_data.items() if k not in ("model", "messages")}, # mutable-ok: spread source - **({"messages": sanitized_messages} if "messages" in request_data else {}), # mutable-ok: spread source + k: (sanitized_messages if k == "messages" else v) for k, v in request_data.items() if k != "model" } if "messages" in body_data: diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index d687fa79234..b7ba135feca 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -652,18 +652,23 @@ class AmazonAnthropicClaudeMessagesConfig( sanitized.offenders, model, ) - anthropic_messages_request["messages"] = list(sanitized.messages) # mutable-ok: outbound JSON array + anthropic_messages_request["messages"] = list( # rebind-ok: out-param # mutable-ok: json + sanitized.messages + ) thinking: Final = anthropic_messages_request.get("thinking") + if not isinstance(thinking, dict): + return normalized_thinking: Final = normalize_bedrock_invoke_thinking_display( thinking, supported_display_values=self.BEDROCK_INVOKE_SUPPORTED_THINKING_DISPLAY_VALUES ) - if normalized_thinking is not thinking: - verbose_logger.warning( - "Bedrock Invoke: mapping unsupported thinking.display %r to 'summarized' for model=%s", - thinking.get("display") if isinstance(thinking, dict) else thinking, - model, - ) - anthropic_messages_request["thinking"] = normalized_thinking + if normalized_thinking is thinking: + return + verbose_logger.warning( + "Bedrock Invoke: mapping unsupported thinking.display %r to 'summarized' for model=%s", + thinking.get("display"), + model, + ) + anthropic_messages_request["thinking"] = normalized_thinking # rebind-ok: out-param def _strip_unsupported_bedrock_invoke_fields( self, diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 0a03fa0e157..7630e760f8c 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3558,9 +3558,17 @@ def test_bedrock_invoke_rejects_message_emptied_by_stripping(): assert "messages[1]" in str(exc.value) already_empty: Final = _transform_for_bedrock_invoke( - [{"role": "user", "content": "hi"}, {"role": "assistant", "content": []}, {"role": "user", "content": "go"}] + [ + {"role": "user", "content": "hi", "output_config": {"effort": "low"}}, + {"role": "assistant", "content": []}, + {"role": "user", "content": [_TOOL_ADDITION_BLOCK, {"type": "text", "text": "go"}]}, + ] ) - assert already_empty["messages"][1]["content"] == [] + assert already_empty["messages"] == [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": []}, + {"role": "user", "content": [{"type": "text", "text": "go"}]}, + ] def test_bedrock_invoke_maps_thinking_display_updates_to_summarized(): diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 117814a41ff..86e5f01e5b3 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -550,6 +550,38 @@ def test_strip_unsupported_output_config_keeps_format_drops_effort(local_model_c assert body["output_config"] == {"format": schema_format} +def test_sanitize_bedrock_invoke_messages_forwards_malformed_entries_next_to_an_offender(): + from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_UNSUPPORTED_CONTENT_BLOCK_TYPES, + BEDROCK_INVOKE_UNSUPPORTED_MESSAGE_KEYS, + sanitize_bedrock_invoke_messages, + ) + + raw_string_block = "not a block dict" + non_dict_message = "not a message dict" + tool_addition = {"type": "tool_addition", "tool_reference": {"type": "tool_reference", "tool_name": "Read"}} + messages = [ + {"role": "user", "content": [raw_string_block, {"type": "text", "text": "hi"}]}, + non_dict_message, + {"role": "assistant", "content": [tool_addition, {"type": "text", "text": "ok"}]}, + ] + + sanitized = sanitize_bedrock_invoke_messages( + messages, + unsupported_keys=BEDROCK_INVOKE_UNSUPPORTED_MESSAGE_KEYS, + unsupported_block_types=BEDROCK_INVOKE_UNSUPPORTED_CONTENT_BLOCK_TYPES, + ) + + assert sanitized.offenders == ("messages[2].content[0] (type 'tool_addition')",) + assert sanitized.emptied == () + assert list(sanitized.messages) == [ + {"role": "user", "content": [raw_string_block, {"type": "text", "text": "hi"}]}, + non_dict_message, + {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}, + ] + assert sanitized.messages[1] is non_dict_message + + def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_map): """The legacy ``output_format`` wins over ``output_config.format`` when a request carries both, matching the pre-existing precedence.""" diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index fde18b51a50..f0153ab995b 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -302,21 +302,22 @@ def test_transform_to_invoke_model_format_leaves_clean_anthropic_body_unchanged( config = BedrockCountTokensConfig() request = { "model": "us.anthropic.claude-sonnet-4-6", - "system": "be brief", "messages": [ {"role": "user", "content": [{"type": "text", "text": "hi"}]}, {"role": "assistant", "content": [{"type": "text", "text": "hello"}]}, ], + "system": "be brief", } body = _decoded_invoke_body(config.transform_anthropic_to_bedrock_count_tokens(request)) assert body == { - "system": "be brief", "messages": request["messages"], + "system": "be brief", "anthropic_version": "bedrock-2023-05-31", "max_tokens": DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS, } + assert list(body) == ["messages", "system", "anthropic_version", "max_tokens"] def test_transform_anthropic_to_bedrock_request_string_content_still_uses_converse():