From b0754dbcb0f17ebbb32d4fc3676efebc45a29c9a Mon Sep 17 00:00:00 2001 From: Filippo Menghi <113345637+Cyberfilo@users.noreply.github.com> Date: Wed, 10 Jun 2026 12:51:57 +0200 Subject: [PATCH] fix(bedrock): route Anthropic-shape count_tokens to InvokeModel and base64-encode the body (#30102) * fix(bedrock): route Anthropic-shape count_tokens to InvokeModel POST /v1/messages/count_tokens with Anthropic content blocks ({"type": "text"|"tool_use"|...}) was routed to the Converse input of the Bedrock CountTokens API. The Converse transform copies list content through verbatim, so Bedrock rejected the request with a 400 and the caller silently fell back to the local tokenizer, returning counts that can be off by ~50% on tool-heavy payloads. _detect_input_type now routes messages whose content blocks carry a "type" key (Anthropic shape) to the invokeModel input, which forwards the body verbatim. The invokeModel body is now base64-encoded as the CountTokens API requires (InvokeModelTokensRequest.body is a base64-encoded blob), and Anthropic Messages bodies get the anthropic_version and max_tokens fields Bedrock validates against. Fixes #27632. * refactor(bedrock): name the CountTokens max_tokens placeholder Replace the magic 1024 with a module-level DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS constant so the intent is explicit and there is a single place to update if Bedrock's InvokeModel schema ever changes. Module-local rather than litellm/constants.py because the value is only a schema-validation placeholder for token counting, not a user-tunable generation default. --- .../bedrock/count_tokens/transformation.py | 36 ++++++++-- ...est_bedrock_count_tokens_transformation.py | 72 ++++++++++++++++++- 2 files changed, 103 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index c967fd334bc..bdef3349e00 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -11,6 +11,11 @@ from typing import Any, Dict, List, Optional from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model +# Placeholder satisfying the Anthropic InvokeModel schema's required +# max_tokens field; CountTokens only counts input, so it has no effect +# on any generation. +DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS = 1024 + class BedrockCountTokensConfig(BaseAWSLLM): """ @@ -32,8 +37,20 @@ class BedrockCountTokensConfig(BaseAWSLLM): Returns: 'converse' or 'invokeModel' """ - # If the request has messages in the expected Anthropic format, use converse - if "messages" in request_data and isinstance(request_data["messages"], list): + messages = request_data.get("messages") + if isinstance(messages, list): + # Anthropic content blocks carry a "type" key ({"type": "text", ...}); + # Converse blocks don't ({"text": ...}, {"toolUse": ...}). Converse + # rejects Anthropic-shape blocks, so route those to invokeModel, + # which forwards the body verbatim. + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, list) and any( + isinstance(block, dict) and "type" in block for block in content + ): + return "invokeModel" return "converse" # For raw text or other formats, use invokeModel @@ -68,7 +85,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): { "input": { "invokeModel": { - "body": "{...raw model input...}" + "body": "" } } } @@ -168,13 +185,24 @@ class BedrockCountTokensConfig(BaseAWSLLM): self, request_data: Dict[str, Any] ) -> Dict[str, Any]: """Transform to InvokeModel input format.""" + import base64 import json # For InvokeModel, we need to provide the raw body that would be sent to the model # Remove the 'model' field from the body as it's not part of the model input body_data = {k: v for k, v in request_data.items() if k != "model"} - return {"input": {"invokeModel": {"body": json.dumps(body_data)}}} + if "messages" in body_data: + # Bedrock validates the body against the model's InvokeModel schema; + # Anthropic Messages bodies require these fields. + body_data.setdefault("anthropic_version", "bedrock-2023-05-31") + body_data.setdefault( + "max_tokens", DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS + ) + + # The CountTokens API expects invokeModel.body as a base64-encoded blob + encoded_body = base64.b64encode(json.dumps(body_data).encode()).decode() + return {"input": {"invokeModel": {"body": encoded_body}}} def get_bedrock_count_tokens_endpoint( self, diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index 0de3f833a37..6812f40829a 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -1,10 +1,15 @@ +import base64 +import json import os import sys sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig +from litellm.llms.bedrock.count_tokens.transformation import ( + DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS, + BedrockCountTokensConfig, +) def test_detect_input_type(): @@ -20,6 +25,71 @@ def test_detect_input_type(): assert config._detect_input_type(request_with_text) == "invokeModel" +def test_detect_input_type_anthropic_blocks_route_to_invoke_model(): + """Anthropic-shape content blocks must not go through the Converse path, + which Bedrock rejects with a 400 (and the caller then silently falls back + to the local tokenizer).""" + config = BedrockCountTokensConfig() + + request = { + "messages": [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Reading the file."}, + { + "type": "tool_use", + "id": "toolu_01", + "name": "read_file", + "input": {"path": "main.py"}, + }, + ], + }, + ], + } + assert config._detect_input_type(request) == "invokeModel" + + +def test_detect_input_type_converse_blocks_route_to_converse(): + """Converse-shape blocks (no "type" key) keep using the converse input.""" + config = BedrockCountTokensConfig() + + request = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} + assert config._detect_input_type(request) == "converse" + + +def test_transform_to_invoke_model_format_base64_encodes_body(): + """The CountTokens API expects invokeModel.body as a base64-encoded blob; + Anthropic Messages bodies additionally need anthropic_version/max_tokens + to pass Bedrock's InvokeModel schema validation.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body["messages"] == request["messages"] + assert "model" not in body + assert body["anthropic_version"] == "bedrock-2023-05-31" + assert body["max_tokens"] == DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS + + +def test_transform_to_invoke_model_format_raw_body_unchanged(): + """Non-messages bodies (e.g. Titan inputText) must not get Anthropic fields.""" + config = BedrockCountTokensConfig() + + result = config.transform_anthropic_to_bedrock_count_tokens( + {"model": "amazon.titan-text-express-v1", "inputText": "hello"} + ) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body == {"inputText": "hello"} + + def test_transform_anthropic_to_bedrock_request(): """Test basic request transformation""" config = BedrockCountTokensConfig()