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.

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": "<base64-encoded raw model input>"
                 }
             }
         }
@@ -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,

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()