scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

Merge pull request #38657 from BerriAI/litellm_count_tokens_fallback_tools_system

Browse source 
fix(proxy): count tools, system, and Anthropic image and document blocks in the count_tokens fallback (internal copy of #36671)
This commit is contained in:
Mateo Wang 2026-08-28 09:52:14 -07:00 • committed by GitHub
commit ecf84a2c2a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 402 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
litellm/litellm_core_utils/token_counter.py
View file

@ -3,7 +3,7 @@
import base64
import io
import struct
from collections.abc import Callable, Mapping
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Any, Final, Literal, cast
import tiktoken
@ -25,14 +25,21 @@ 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,
AnthropicMessagesDocumentParam,
AnthropicMessagesImageParam,
AnthropicMessagesTextParam,
AnthropicMessagesToolResultParam,
AnthropicMessagesToolUseParam,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionDocumentObject,
ChatCompletionNamedToolChoiceParam,
ChatCompletionToolParam,
OpenAIMessageContent,
OpenAIMessageContentListBlock,
)
from litellm.types.utils import Message, SelectTokenizerResponse
@ -346,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,
@ -646,6 +653,46 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type:
return expected_cls
def _anthropic_image_source_data(
source: AnthropicContentParamSource | AnthropicContentParamSourceUrl | AnthropicContentParamSourceFileId,
) -> str:
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":
return source.get("url") or ""
return ""
def _count_document_tokens(
document: ChatCompletionDocumentObject | AnthropicMessagesDocumentParam,
count_function: TokenCounterFunction,
use_default_image_token_count: bool,
default_token_count: int | None,
) -> int:
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,
@ -697,13 +744,17 @@ def _count_anthropic_content(
def _count_content_list(
count_function: TokenCounterFunction,
content_list: OpenAIMessageContent,
content_list: str
| Iterable[
OpenAIMessageContentListBlock
| AnthropicMessagesTextParam
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
],
use_default_image_token_count: bool,
default_token_count: int | None,
) -> int:
"""
Recursively count tokens from a list of content blocks.
"""
"""Recursively count tokens from a list of content blocks."""
try:
num_tokens = 0
for c in content_list:
@ -714,6 +765,19 @@ def _count_content_list(
elif c["type"] == "image_url":
image_url = c.get("image_url")
num_tokens += _count_image_tokens(image_url, use_default_image_token_count)
elif c["type"] == "image":
num_tokens += calculate_img_tokens(
data=_anthropic_image_source_data(c["source"]),
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,
@ -742,7 +806,8 @@ def _count_content_list(
content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__
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, tool_reference)."
f"Expected str or dict with 'type' field "
f"(text, image_url, image, document, tool_use, tool_result, thinking, tool_reference)."
)
return num_tokens
except Exception as e:

27
litellm/proxy/proxy_server.py
View file

@ -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 if counted_messages is not None else None
)
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(

21
litellm/types/llms/anthropic.py
View file

@ -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

4
litellm/utils.py
View file

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

217
tests/test_litellm/litellm_core_utils/test_token_counter.py
View file

@ -1160,3 +1160,220 @@ def test_count_content_list_rejects_unknown_type():
message = str(exc_info.value)
assert "Invalid content item type: totally_unknown_block" in message
assert "tool_reference" in message
@pytest.mark.parametrize(
"source",
[
{"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="},
{"type": "url", "url": "https://example.com/image.png"},
{"type": "file", "file_id": "file-abc123"},
],
ids=["base64", "url", "file"],
)
def test_token_counter_with_anthropic_image_block(source: dict[str, str]):
"""Anthropic `image` blocks must count for every source variant, not raise `Invalid content item type` (which the router's context-window pre-call check swallows into an unfiltered dispatch)."""
from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is in this image?"},
{"type": "image", "source": source},
],
}
]
tokens = token_counter(
model="anthropic/claude-sonnet-4-5-20250929",
messages=messages,
use_default_image_token_count=True,
)
assert tokens > DEFAULT_IMAGE_TOKEN_COUNT, (
f"Expected the image block to contribute tokens, got {tokens}"
)
def test_anthropic_image_block_matches_equivalent_image_url():
"""An Anthropic `image` block prices identically to the OpenAI `image_url` carrying the same bytes."""
anthropic_messages = [
{
"role": "user",
"content": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": "iVBORw0KGgo=",
},
}
],
}
]
openai_messages = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
}
],
}
]
anthropic_tokens = token_counter(
model="anthropic/claude-sonnet-4-5-20250929", messages=anthropic_messages
)
openai_tokens = token_counter(
model="anthropic/claude-sonnet-4-5-20250929", messages=openai_messages
)
assert anthropic_tokens == openai_tokens
def test_anthropic_image_block_nested_in_tool_result():
"""An `image` block nested in a `tool_result.content` list is counted through the same recursion."""
messages = [
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_01",
"content": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": "iVBORw0KGgo=",
},
}
],
}
],
}
]
tokens = token_counter(
model="anthropic/claude-sonnet-4-5-20250929",
messages=messages,
use_default_image_token_count=True,
)
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):
"""base64 sources become a data URI, url sources pass through, file sources resolve to an empty string."""
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 with empty `data` prices as an image rather than raising."""
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` raises, matching the OpenAI `image_url`-without-`url` behavior."""
from litellm.litellm_core_utils.token_counter import _count_content_list
with pytest.raises(ValueError, match="Error getting number of tokens from content list"):
_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
)
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 can't be tokenized locally is priced like an `image`, not raised on."""
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},
]
)

66
tests/test_litellm/proxy/proxy_server/test_routes_utils.py
View file

@ -184,3 +184,69 @@ 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):
"""The ``litellm.token_counter`` fallback counts the request's tools and system prompt, and Anthropic ``image``/``document`` blocks, instead of 500ing."""
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
def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch):
"""Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path)."""
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
prompt = "count the tokens in this sentence please"
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Look up the current weather for a city",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
},
}
]
with auth_as():
response = client.post(
"/utils/token_counter", json={"model": "claude-fable-5", "prompt": prompt, "tools": tools}
)
assert response.status_code == 200, response.text
assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt)