fix: support Anthropic image and document in token counter (#26704)

Squash-merged by litellm-agent from michelligabriele's PR.

litellm/litellm_core_utils/token_counter.py

@@ -623,6 +623,94 @@ def _count_image_tokens(
         )
 
 
+def _count_anthropic_image(
+    content: Mapping[str, Any],
+    use_default_image_token_count: bool,
+) -> int:
+    """
+    Count tokens for an Anthropic-native image content block.
+
+    Anthropic shape: {"type": "image", "source": {...}}
+    Source variants:
+        - {"type": "base64", "media_type": ..., "data": ...}
+        - {"type": "url", "url": ...}
+        - {"type": "file", "file_id": ...}
+
+    For base64/url sources we delegate to calculate_img_tokens (same path as
+    the OpenAI image_url branch). For file sources, and unknown future source
+    types, we fall back to DEFAULT_IMAGE_TOKEN_COUNT — there is no fetchable
+    payload to dimension and we don't want to re-introduce the log-spam bug
+    fixed here.
+    """
+    source = content.get("source")
+    if not isinstance(source, dict):
+        raise ValueError("Anthropic image block missing required 'source' dict")
+
+    source_type = source.get("type")
+    if source_type == "base64":
+        media_type = source.get("media_type", "")
+        data = source.get("data", "")
+        data_url = f"data:{media_type};base64,{data}"
+        return calculate_img_tokens(
+            data=data_url,
+            mode="auto",
+            use_default_image_token_count=use_default_image_token_count,
+        )
+    if source_type == "url":
+        url = source.get("url")
+        if not url:
+            raise ValueError(
+                "Anthropic image block with source.type='url' missing 'url'"
+            )
+        return calculate_img_tokens(
+            data=url,
+            mode="auto",
+            use_default_image_token_count=use_default_image_token_count,
+        )
+    # source.type == "file" or unknown — no fetchable image data.
+    return DEFAULT_IMAGE_TOKEN_COUNT
+
+
+def _count_anthropic_document(
+    content: Mapping[str, Any],
+    count_function: TokenCounterFunction,
+) -> int:
+    """
+    Count tokens for an Anthropic-native document content block.
+
+    Anthropic shape:
+        {"type": "document", "source": {...}, "title": ..., "context": ..., "citations": ...}
+
+    Source variants (per the Anthropic Citations API):
+        - {"type": "text", "media_type": "text/plain", "data": "<inline text>"}
+        - {"type": "base64", "media_type": ..., "data": ...}
+        - {"type": "url", "url": ...}
+        - {"type": "file", "file_id": ...}
+
+    For text sources we tokenize source["data"] with the model tokenizer
+    (it's plain prompt content). For base64 / url / file sources the
+    payload is opaque (typically a PDF — we cannot dimension it) and we
+    add DEFAULT_IMAGE_TOKEN_COUNT as the fallback.
+
+    Title and context are always counted. Skipped: "type", "cache_control",
+    "citations" (configuration metadata, not prompt content).
+    """
+    tokens = 0
+    for field in ("title", "context"):
+        value = content.get(field)
+        if isinstance(value, str) and value:
+            tokens += count_function(value)
+    source = content.get("source")
+    if isinstance(source, dict) and source.get("type") == "text":
+        data = source.get("data")
+        if isinstance(data, str) and data:
+            tokens += count_function(data)
+    elif source is not None:
+        # base64 / url / file / unknown — opaque payload, use the fallback.
+        tokens += DEFAULT_IMAGE_TOKEN_COUNT
+    return tokens
+
+
 def _validate_anthropic_content(content: Mapping[str, Any]) -> type:
     """
     Validate and determine which Anthropic TypedDict applies.
@@ -736,6 +824,10 @@ def _count_content_list(
                 thinking_text = str(c.get("thinking", ""))
                 if thinking_text:
                     num_tokens += count_function(thinking_text)
+            elif c["type"] == "image":
+                num_tokens += _count_anthropic_image(c, use_default_image_token_count)
+            elif c["type"] == "document":
+                num_tokens += _count_anthropic_document(c, count_function)
             else:
                 content_type = (
                     c.get("type", type(c).__name__)
@@ -744,7 +836,8 @@ def _count_content_list(
                 )
                 raise ValueError(
                     f"Invalid content item type: {content_type}. "
-                    f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking)."
+                    f"Expected str or dict with 'type' field "
+                    f"(text, image_url, image, document, tool_use, tool_result, thinking)."
                 )
         return num_tokens
     except Exception as e:

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -980,3 +980,245 @@ def test_token_counter_with_thinking_content():
     assert (
         tokens_no_thinking < 15
     ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}"
+
+
+# 1×1 transparent PNG, base64-encoded — small enough to inline in tests.
+_TEST_PNG_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
+
+
+def test_token_counter_with_anthropic_image_base64():
+    """
+    Anthropic-native image block with source.type='base64' is counted, not rejected.
+
+    Regression test for https://github.com/BerriAI/litellm/issues/20367 —
+    previously raised ValueError("Invalid content item type: image. ...").
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What's in this image?"},
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/png",
+                        "data": _TEST_PNG_BASE64,
+                    },
+                },
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+    )
+    # User text + image base tokens (calculate_img_tokens auto-mode = 85) + overhead.
+    assert tokens > 85, f"Expected image block to contribute tokens, got {tokens}"
+
+
+def test_token_counter_with_anthropic_image_url():
+    """Anthropic-native image block with source.type='url' is counted."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "url",
+                        "url": "https://example.com/image.png",
+                    },
+                }
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="anthropic/claude-sonnet-4-5-20250929",
+        messages=messages,
+        use_default_image_token_count=True,  # avoid an actual HTTP fetch
+    )
+    assert tokens > 0, f"Expected positive token count, got {tokens}"
+
+
+def test_token_counter_with_anthropic_image_file():
+    """
+    Anthropic-native image block with source.type='file' is counted.
+
+    file_id refers to an Anthropic Files API upload; there is no fetchable
+    image payload, so the counter falls back to DEFAULT_IMAGE_TOKEN_COUNT.
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "file",
+                        "file_id": "file_011CQc1DRcaXZyDtCkP9NbBz",
+                    },
+                }
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+    )
+    assert tokens > 0, f"Expected positive token count, got {tokens}"
+
+
+def test_token_counter_with_anthropic_document_base64():
+    """
+    Anthropic-native document block with title/context text fields is counted.
+
+    Validates that:
+    - 'title' and 'context' string fields are counted via the model tokenizer
+    - source payload contributes the default image-token fallback
+    - 'citations' and 'type' metadata are not counted as text
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "document",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "application/pdf",
+                        "data": "JVBERi0xLjQKJeLjz9MK",  # truncated, content not parsed
+                    },
+                    "title": "Q4 financial report",
+                    "context": "Internal quarterly summary for the finance team.",
+                    "citations": {"enabled": True},
+                }
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+    )
+    # title + context together are ~15 tokens, plus DEFAULT_IMAGE_TOKEN_COUNT (250) for the source.
+    assert (
+        tokens > 250
+    ), f"Expected document tokens to include source fallback, got {tokens}"
+
+
+def test_token_counter_with_anthropic_document_file():
+    """Anthropic-native document block with source.type='file' is counted."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "document",
+                    "source": {
+                        "type": "file",
+                        "file_id": "file_011CQc1DRcaXZyDtCkP9NbBz",
+                    },
+                }
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+    )
+    assert tokens > 0, f"Expected positive token count, got {tokens}"
+
+
+def test_token_counter_with_anthropic_document_text():
+    """
+    Anthropic-native document block with source.type='text' tokenizes the
+    inline text in source['data'], not the flat DEFAULT_IMAGE_TOKEN_COUNT
+    fallback used for opaque (base64/url/file) payloads.
+    """
+    inline_text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. " * 200
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "document",
+                    "source": {
+                        "type": "text",
+                        "media_type": "text/plain",
+                        "data": inline_text,
+                    },
+                    "title": "Long inline doc",
+                }
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+    )
+    # The inline text alone is well over a thousand tokens — substantially
+    # more than the 250-token fallback used for opaque sources.
+    assert (
+        tokens > 500
+    ), f"Expected text-source document to tokenize source['data'], got {tokens}"
+
+
+def test_token_counter_with_image_inside_tool_result():
+    """
+    Image block nested inside tool_result.content is counted via recursion.
+
+    _count_anthropic_content iterates tool_result fields; the 'content' field
+    being a list re-enters _count_content_list, where the new 'image' branch
+    fires.
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": "toolu_01234567890",
+                    "content": [
+                        {"type": "text", "text": "Here is the rendered chart:"},
+                        {
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": "image/png",
+                                "data": _TEST_PNG_BASE64,
+                            },
+                        },
+                    ],
+                }
+            ],
+        }
+    ]
+
+    tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+    assert tokens > 85, f"Expected nested image to contribute tokens, got {tokens}"
+
+
+def test_token_counter_with_malformed_image_uses_default_token_count():
+    """
+    Malformed Anthropic image block (no 'source') doesn't crash when
+    default_token_count is provided — _count_content_list's outer try/except
+    swallows the inner ValueError and returns the fallback.
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},  # missing 'source' — would raise without fallback
+            ],
+        }
+    ]
+
+    tokens = token_counter(
+        model="gpt-3.5-turbo",
+        messages=messages,
+        default_token_count=42,
+    )
+    # The malformed-block fallback returns 42 for the content-list count;
+    # outer message bookkeeping (tokens_per_message etc.) adds a few more.
+    assert tokens >= 42, f"Expected default_token_count fallback, got {tokens}"