fix(proxy): count tools, system, and Anthropic document blocks in the count_tokens fallback

litellm/litellm_core_utils/token_counter.py

@@ -3,7 +3,7 @@
 import base64
 import io
 import struct
-from collections.abc import Callable, Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from typing import Any, Final, Literal, cast
 
 import tiktoken
@@ -28,12 +28,15 @@ from litellm.types.llms.anthropic import (
     AnthropicContentParamSource,
     AnthropicContentParamSourceFileId,
     AnthropicContentParamSourceUrl,
+    AnthropicMessagesDocumentParam,
     AnthropicMessagesImageParam,
+    AnthropicMessagesTextParam,
     AnthropicMessagesToolResultParam,
     AnthropicMessagesToolUseParam,
 )
 from litellm.types.llms.openai import (
     AllMessageValues,
+    ChatCompletionDocumentObject,
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionToolParam,
     OpenAIMessageContentListBlock,
@@ -350,7 +353,7 @@ def token_counter(
     model="",
     custom_tokenizer: dict | SelectTokenizerResponse | None = None,
     text: str | list[str] | None = None,
-    messages: list[AllMessageValues | Message] | None = None,
+    messages: Sequence[AllMessageValues | Message] | None = None,
     count_response_tokens: bool | None = False,
     tools: list[ChatCompletionToolParam] | None = None,
     tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
@@ -669,6 +672,38 @@ def _anthropic_image_source_data(
     return ""
 
 
+def _count_document_tokens(
+    document: ChatCompletionDocumentObject | AnthropicMessagesDocumentParam,
+    count_function: TokenCounterFunction,
+    use_default_image_token_count: bool,
+    default_token_count: int | None,
+) -> int:
+    """
+    Count an Anthropic `document` block: its title and context text, plus the source itself.
+
+    Text-bearing sources (`text`, `content`) count their text; opaque ones (`base64`, `url`,
+    `file`) are priced like an image, since their bytes cannot be tokenized locally.
+    """
+    source: Final = document["source"]
+    metadata_tokens: Final = sum(
+        count_function(text) for text in (document.get("title"), document.get("context")) if text
+    )
+    if source["type"] == "text":
+        return metadata_tokens + count_function(source["data"])
+    if source["type"] == "content":
+        content: Final = source["content"]
+        if isinstance(content, str):
+            return metadata_tokens + count_function(content)
+        return metadata_tokens + _count_content_list(
+            count_function, content, use_default_image_token_count, default_token_count
+        )
+    return metadata_tokens + calculate_img_tokens(
+        data=_anthropic_image_source_data(source),
+        mode="auto",
+        use_default_image_token_count=use_default_image_token_count,
+    )
+
+
 def _count_anthropic_content(
     content: Mapping[str, Any],
     count_function: TokenCounterFunction,
@@ -720,7 +755,13 @@ def _count_anthropic_content(
 
 def _count_content_list(
     count_function: TokenCounterFunction,
-    content_list: str | Iterable[OpenAIMessageContentListBlock | AnthropicMessagesImageParam],
+    content_list: str
+    | Iterable[
+        OpenAIMessageContentListBlock
+        | AnthropicMessagesTextParam
+        | AnthropicMessagesImageParam
+        | AnthropicMessagesDocumentParam
+    ],
     use_default_image_token_count: bool,
     default_token_count: int | None,
 ) -> int:
@@ -747,6 +788,13 @@ def _count_content_list(
                     mode="auto",
                     use_default_image_token_count=use_default_image_token_count,
                 )
+            elif c["type"] == "document":
+                num_tokens += _count_document_tokens(
+                    c,
+                    count_function,
+                    use_default_image_token_count,
+                    default_token_count,
+                )
             elif c["type"] in ("tool_use", "tool_result"):
                 num_tokens += _count_anthropic_content(
                     c,
@@ -776,7 +824,7 @@ def _count_content_list(
                 raise ValueError(
                     f"Invalid content item type: {content_type}. "
                     f"Expected str or dict with 'type' field "
-                    f"(text, image_url, image, tool_use, tool_result, thinking, tool_reference)."
+                    f"(text, image_url, image, document, tool_use, tool_result, thinking, tool_reference)."
                 )
         return num_tokens
     except Exception as e:

litellm/proxy/proxy_server.py

@@ -672,7 +672,12 @@ from litellm.types.llms.anthropic import (
     AnthropicResponseContentBlockText,
     AnthropicResponseUsageBlock,
 )
-from litellm.types.llms.openai import HttpxBinaryResponseContent
+from litellm.types.llms.openai import (
+    AllMessageValues,
+    ChatCompletionSystemMessage,
+    ChatCompletionToolParam,
+    HttpxBinaryResponseContent,
+)
 from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
 from litellm.types.proxy.management_endpoints.model_management_endpoints import (
     ModelGroupInfoProxy,
@@ -12126,6 +12131,13 @@ async def _try_provider_token_count(
     return result
 
 
+def _system_message(system: object) -> ChatCompletionSystemMessage | None:
+    if not isinstance(system, (str, list)) or not system:
+        return None
+    message: Final[ChatCompletionSystemMessage] = {"role": "system", "content": system}
+    return message
+
+
 @router.post(
     "/utils/token_counter",
     tags=["llm utils"],
@@ -12224,10 +12236,21 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer)
 
     tokenizer_used: Final = str(_tokenizer_used["type"])
+    system_message: Final = _system_message(system)
+    typed_messages: Final = cast(  # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes
+        Sequence[AllMessageValues] | None, messages
+    )
+    counted_messages: Final = (
+        typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
+    )
+    counted_tools: Final = cast(  # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
+        list[ChatCompletionToolParam] | None, tools
+    )
     total_tokens: Final = await asyncify(litellm.token_counter)(
         model=model_to_use,
         text=prompt,
-        messages=messages,
+        messages=counted_messages,
+        tools=counted_tools,
         custom_tokenizer=_tokenizer_used,
     )
     return TokenCountResponse(

litellm/types/llms/anthropic.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from enum import Enum
 from typing import Any, Final, Literal, TypeAlias
 
@@ -254,6 +254,17 @@ class AnthropicContentParamSourceFileId(TypedDict):
     file_id: str
 
 
+class AnthropicContentParamSourceText(TypedDict):
+    type: ReadOnly[Literal["text"]]
+    media_type: ReadOnly[Literal["text/plain"]]
+    data: ReadOnly[str]
+
+
+class AnthropicContentParamSourceContent(TypedDict):
+    type: ReadOnly[Literal["content"]]
+    content: ReadOnly[str | Sequence["AnthropicMessagesTextParam | AnthropicMessagesImageParam"]]
+
+
 class AnthropicMessagesContainerUploadParam(TypedDict, total=False):
     type: Required[Literal["container_upload"]]
     file_id: str
@@ -305,7 +316,13 @@ AnthropicCitation = AnthropicCitationPageLocation | AnthropicCitationCharLocatio
 
 class AnthropicMessagesDocumentParam(TypedDict, total=False):
     type: Required[Literal["document"]]
-    source: Required[AnthropicContentParamSource | AnthropicContentParamSourceFileId | AnthropicContentParamSourceUrl]
+    source: Required[
+        AnthropicContentParamSource
+        | AnthropicContentParamSourceFileId
+        | AnthropicContentParamSourceUrl
+        | AnthropicContentParamSourceText
+        | AnthropicContentParamSourceContent
+    ]
     cache_control: dict | ChatCompletionCachedContent | None
     title: str
     context: str

litellm/utils.py

@@ -2302,7 +2302,7 @@ def token_counter(
     model="",
     custom_tokenizer: dict | SelectTokenizerResponse | None = None,
     text: str | list[str] | None = None,
-    messages: list | None = None,
+    messages: Sequence | None = None,
     count_response_tokens: bool | None = False,
     tools: list[ChatCompletionToolParam] | None = None,
     tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
@@ -7741,7 +7741,7 @@ def convert_to_dict(message: BaseModel | dict) -> dict:
         raise TypeError(f"Invalid message type: {type(message)}. Expected dict or Pydantic model.")
 
 
-def convert_list_message_to_dict(messages: list):
+def convert_list_message_to_dict(messages: Sequence):
     new_messages: Final = []
     for message in messages:
         convert_msg_to_dict = cast(AllMessageValues, convert_to_dict(message))

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1282,6 +1282,26 @@ def test_anthropic_image_block_nested_in_tool_result():
     assert tokens > 0
 
 
+@pytest.mark.parametrize(
+    ("source", "expected"),
+    [
+        ({"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQ"}, "data:image/jpeg;base64,/9j/4AAQ"),
+        ({"type": "url", "url": "https://example.com/image.png"}, "https://example.com/image.png"),
+        ({"type": "file", "file_id": "file-abc123"}, ""),
+    ],
+    ids=["base64", "url", "file"],
+)
+def test_anthropic_image_source_resolves_to_what_the_image_pricer_reads(source: dict[str, str], expected: str):
+    """
+    The image pricer reads either a data URI or a fetchable URL: a base64 source keeps its
+    media type inside the URI, a url source passes through untouched, and a file source has
+    no bytes the proxy can measure locally.
+    """
+    from litellm.litellm_core_utils.token_counter import _anthropic_image_source_data
+
+    assert _anthropic_image_source_data(source) == expected
+
+
 def test_anthropic_image_block_with_empty_base64_data():
     """
     A base64 source carrying no bytes must still price as an image rather than
@@ -1309,7 +1329,7 @@ def test_anthropic_image_block_without_source_raises():
     """
     from litellm.litellm_core_utils.token_counter import _count_content_list
 
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="Error getting number of tokens from content list"):
         _count_content_list(
             count_function=len,
             content_list=[{"type": "image"}],
@@ -1327,3 +1347,65 @@ def test_anthropic_image_block_without_source_raises():
         )
         == 7
     )
+
+
+def _count_user_content(content: list[dict]) -> int:
+    from litellm.litellm_core_utils.token_counter import token_counter
+
+    return token_counter(
+        model="anthropic/claude-fable-5",
+        messages=[{"role": "user", "content": content}],
+        use_default_image_token_count=True,
+    )
+
+
+@pytest.mark.parametrize(
+    "source",
+    [
+        {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
+        {"type": "url", "url": "https://example.com/report.pdf"},
+        {"type": "file", "file_id": "file-abc123"},
+    ],
+    ids=["base64", "url", "file"],
+)
+def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]):
+    """
+    A `document` whose bytes cannot be tokenized locally must not raise (it 500ed
+    /v1/messages/count_tokens before) and is priced exactly like an `image` block.
+    """
+    prompt = {"type": "text", "text": "Summarize this file."}
+
+    assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content(
+        [prompt, {"type": "image", "source": source}]
+    )
+
+
+def test_anthropic_document_block_text_sources_count_their_text():
+    """`text` and `content` document sources count the text they carry, as inline text blocks would."""
+    prompt = {"type": "text", "text": "Summarize this file."}
+    body = {"type": "text", "text": "Revenue grew eleven percent while churn fell to two percent."}
+    picture = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}}
+
+    text_source = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": body["text"]}}
+    assert _count_user_content([prompt, text_source]) == _count_user_content([prompt, body])
+
+    string_content = {"type": "document", "source": {"type": "content", "content": body["text"]}}
+    assert _count_user_content([prompt, string_content]) == _count_user_content([prompt, body])
+
+    block_content = {"type": "document", "source": {"type": "content", "content": [body, picture]}}
+    assert _count_user_content([prompt, block_content]) == _count_user_content([prompt, body, picture])
+
+
+def test_anthropic_document_title_and_context_add_their_tokens():
+    prompt = {"type": "text", "text": "Summarize this file."}
+    source = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}
+    described = {"type": "document", "source": source, "title": "Q3 board packet", "context": "Shared by finance"}
+
+    assert _count_user_content([prompt, described]) == _count_user_content(
+        [
+            prompt,
+            {"type": "text", "text": "Q3 board packet"},
+            {"type": "text", "text": "Shared by finance"},
+            {"type": "document", "source": source},
+        ]
+    )

tests/test_litellm/proxy/proxy_server/test_routes_utils.py

@@ -184,3 +184,48 @@ def test_transform_request_unsafe_body(client, auth_as, monkeypatch):
         response = client.post("/utils/transform_request", json=payload)
     assert response.status_code == 400
     assert "unsafe" in response.text or "error" in response.text
+
+
+def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, auth_as, monkeypatch):
+    """
+    Without a provider counter the route falls back to ``litellm.token_counter``. That count
+    must include the request's tools and system prompt, and Anthropic ``image`` and ``document``
+    blocks must be counted instead of turning the whole request into a 500.
+    """
+    monkeypatch.setattr(proxy_server, "llm_router", None)
+    monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
+    system = [{"type": "text", "text": "You are a terse assistant. Answer in one sentence."}]
+    tools = [
+        {
+            "name": "get_weather",
+            "description": "Look up the current weather for a city",
+            "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+        }
+    ]
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What is in this file?"},
+                {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}},
+                {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}},
+            ],
+        }
+    ]
+
+    def count(payload: dict) -> int:
+        with auth_as():
+            response = client.post("/utils/token_counter", json={"model": "claude-fable-5", **payload})
+        assert response.status_code == 200, response.text
+        return response.json()["total_tokens"]
+
+    bare = count({"messages": messages})
+    full = count({"messages": messages, "tools": tools, "system": system})
+
+    assert bare == litellm.token_counter(model="claude-fable-5", messages=messages)
+    assert full == litellm.token_counter(
+        model="claude-fable-5",
+        messages=[{"role": "system", "content": system}, *messages],
+        tools=tools,
+    )
+    assert full > bare