fix(token-counter): enhance handling of Anthropic image blocks in token counting

litellm/litellm_core_utils/token_counter.py

@@ -3,7 +3,7 @@
 import base64
 import io
 import struct
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from typing import Any, Final, Literal, cast
 
 import tiktoken
@@ -25,6 +25,10 @@ from litellm.litellm_core_utils.default_encoding import encoding as default_enco
 from litellm.litellm_core_utils.url_utils import safe_get
 from litellm.llms.custom_httpx.http_handler import _get_httpx_client
 from litellm.types.llms.anthropic import (
+    AnthropicContentParamSource,
+    AnthropicContentParamSourceFileId,
+    AnthropicContentParamSourceUrl,
+    AnthropicMessagesImageParam,
     AnthropicMessagesToolResultParam,
     AnthropicMessagesToolUseParam,
 )
@@ -32,7 +36,7 @@ from litellm.types.llms.openai import (
     AllMessageValues,
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionToolParam,
-    OpenAIMessageContent,
+    OpenAIMessageContentListBlock,
 )
 from litellm.types.utils import Message, SelectTokenizerResponse
 
@@ -646,20 +650,21 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type:
     return expected_cls
 
 
-def _anthropic_image_source_data(source: Mapping[str, str]) -> str:
+def _anthropic_image_source_data(
+    source: AnthropicContentParamSource | AnthropicContentParamSourceUrl | AnthropicContentParamSourceFileId,
+) -> str:
     """
     Resolve an Anthropic image `source` to the data string `calculate_img_tokens` prices.
 
     Returns "" for a `file` source, whose bytes the proxy cannot resolve locally.
     """
-    source_type: Final = source.get("type")
-    if source_type == "base64":
+    if source["type"] == "base64":
         data: Final = source.get("data")
         if not data:
             return ""
         media_type: Final = source.get("media_type") or "image/png"
         return f"data:{media_type};base64,{data}"
-    if source_type == "url":
+    if source["type"] == "url":
         return source.get("url") or ""
     return ""
 
@@ -715,12 +720,16 @@ def _count_anthropic_content(
 
 def _count_content_list(
     count_function: TokenCounterFunction,
-    content_list: OpenAIMessageContent,
+    content_list: str | Iterable[OpenAIMessageContentListBlock | AnthropicMessagesImageParam],
     use_default_image_token_count: bool,
     default_token_count: int | None,
 ) -> int:
     """
     Recursively count tokens from a list of content blocks.
+
+    The block union is wider than OpenAI's: the proxy's Anthropic endpoints count
+    their native blocks through this same helper, so an `image` block is as much
+    an input here as OpenAI's `image_url`.
     """
     try:
         num_tokens = 0
@@ -733,9 +742,8 @@ def _count_content_list(
                 image_url = c.get("image_url")
                 num_tokens += _count_image_tokens(image_url, use_default_image_token_count)
             elif c["type"] == "image":
-                source = c.get("source")
                 num_tokens += calculate_img_tokens(
-                    data=_anthropic_image_source_data(source) if isinstance(source, dict) else "",
+                    data=_anthropic_image_source_data(c["source"]),
                     mode="auto",
                     use_default_image_token_count=use_default_image_token_count,
                 )

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1280,3 +1280,50 @@ def test_anthropic_image_block_nested_in_tool_result():
         use_default_image_token_count=True,
     )
     assert tokens > 0
+
+
+def test_anthropic_image_block_with_empty_base64_data():
+    """
+    A base64 source carrying no bytes must still price as an image rather than
+    raise: the block is well-formed enough to count, and an empty `data` only
+    means there is nothing to measure the dimensions from.
+    """
+    from litellm.litellm_core_utils.token_counter import _count_content_list
+
+    tokens = _count_content_list(
+        count_function=len,
+        content_list=[
+            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}}
+        ],
+        use_default_image_token_count=False,
+        default_token_count=None,
+    )
+    assert tokens > 0
+
+
+def test_anthropic_image_block_without_source_raises():
+    """
+    An `image` block with no `source` is malformed, and must fail the same way
+    the OpenAI `image_url` block with no `url` does - a ValueError the caller
+    can turn into a 400 - instead of being silently counted as a valid image.
+    """
+    from litellm.litellm_core_utils.token_counter import _count_content_list
+
+    with pytest.raises(ValueError):
+        _count_content_list(
+            count_function=len,
+            content_list=[{"type": "image"}],
+            use_default_image_token_count=False,
+            default_token_count=None,
+        )
+
+    # ... and `default_token_count`, the caller's opt-out from raising, still wins.
+    assert (
+        _count_content_list(
+            count_function=len,
+            content_list=[{"type": "image"}],
+            use_default_image_token_count=False,
+            default_token_count=7,
+        )
+        == 7
+    )