fix(anthropic): translate chat tool turns and function tools for count_tokens

litellm.acount_tokens and /utils/token_counter take chat-completions
messages and tools, but the Anthropic count_tokens transform only lifted
a leading system run and converted image_url blocks. Assistant tool_calls,
role "tool" results and type "function" tools reached count_tokens as
sent, Anthropic answered 400, and the count silently fell back to the
local tokenizer. An assistant tool_calls turn with no content key, which
OpenAI allows, failed request validation before that and fell back the
same way

Messages with an assistant tool_calls turn now go through
anthropic_messages_pt, and each type "function" tool goes through
AnthropicConfig.map_tools, the helpers the chat path also uses. Other
tools are sent as before, and so are messages without a tool_calls turn

Fixes #26324
This commit is contained in:
Armen Petrosyan 2026-10-13 11:54:08 +03:00
parent a5d05355b6
commit f307eccf44
2 changed files with 188 additions and 10 deletions

View file

@ -6,7 +6,10 @@ This module handles the transformation of requests to Anthropic's CountTokens AP
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from typing import (
Final,
cast, # noqa: TID251 # validating as AllMessageValues turns Iterable content into a one-shot iterator
)
from pydantic import JsonValue, TypeAdapter
@ -17,10 +20,10 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
)
from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers
from litellm.llms.anthropic.wif import resolve_anthropic_base
from litellm.types.llms.openai import ChatCompletionImageObject
from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
_SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue])
_JSON_LIST: Final = TypeAdapter(list[JsonValue])
_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject)
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
@ -53,6 +56,31 @@ def _count_content(content: JsonValue) -> JsonValue:
return [_count_block(block) for block in content] if isinstance(content, list) else content
def _count_messages(model: str, messages: Sequence[Mapping[str, JsonValue]]) -> JsonValue:
if not any("tool_calls" in message for message in messages):
return [{**message, "content": _count_content(message["content"])} for message in messages]
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
return _JSON_LIST.validate_python(
anthropic_messages_pt(
messages=cast( # cast-ok: validating as AllMessageValues turns Iterable content into a one-shot iterator
list[AllMessageValues], [dict(message) for message in messages]
),
model=model,
llm_provider="anthropic",
)
)
def _count_tool(tool: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
if "function" not in tool:
return tool
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
[anthropic_tool], _ = AnthropicConfig().map_tools([dict[str, object](tool)])
return _COUNT_REQUEST.validate_python(anthropic_tool)
def _lift_leading_system(
messages: Sequence[Mapping[str, JsonValue]], system: JsonValue
) -> tuple[tuple[Mapping[str, JsonValue], ...], JsonValue]:
@ -66,7 +94,7 @@ def _lift_leading_system(
leading, conversation = split_leading_system_run(messages)
if not leading or not (system is None or isinstance(system, (str, list))):
return tuple(messages), system
lifted: Final = _SYSTEM_BLOCKS.validate_python(list(anthropic_system_blocks(leading)))
lifted: Final = _JSON_LIST.validate_python(list(anthropic_system_blocks(leading)))
if isinstance(system, list):
return conversation, [*system, *lifted]
if isinstance(system, str) and system:
@ -112,7 +140,8 @@ class AnthropicCountTokensConfig:
Transform request to Anthropic CountTokens format.
Includes optional system and tools fields for accurate token counting; a leading run of
system-role messages is counted through ``system``, the only place count_tokens accepts it.
system-role messages is counted through ``system``, the only place count_tokens accepts it,
and chat tool turns and ``function`` tools get the chat path's Anthropic translation.
"""
options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({})
counted_messages, counted_system = _lift_leading_system(messages, system)
@ -120,13 +149,14 @@ class AnthropicCountTokensConfig:
MappingProxyType(
{
"model": model,
"messages": [
{**message, "content": _count_content(message["content"])} for message in counted_messages
],
"messages": _count_messages(model, counted_messages),
**MappingProxyType(
{
key: value
for key, value in (("system", counted_system), ("tools", tools))
for key, value in (
("system", counted_system),
("tools", None if tools is None else [_count_tool(tool) for tool in tools]),
)
if value is not None
}
),
@ -185,5 +215,5 @@ class AnthropicCountTokensConfig:
if "role" not in message:
raise ValueError(f"Message {i} must have a 'role' field")
if "content" not in message:
if "content" not in message and "tool_calls" not in message:
raise ValueError(f"Message {i} must have a 'content' field")

View file

@ -441,3 +441,151 @@ async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(h
"system": [{"type": "text", "text": "Be terse"}],
"messages": [{"role": "user", "content": "Hello, how are you?"}],
}
@pytest.mark.parametrize("config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig))
def test_count_translates_chat_tool_turns_and_function_tools(
config_type: type[AnthropicCountTokensConfig],
) -> None:
"""``acount_tokens`` takes chat-completions messages and tools. count_tokens answers 400 on
``role: "tool"`` and ``type: "function"`` tools, so chat tool turns and function tools get the chat
path's translation, and an Anthropic server tool next to them is sent as is."""
cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"}
web_search: Final[dict[str, JsonValue]] = {"type": "web_search_20260209", "name": "web_search", "max_uses": 2}
messages: Final[list[dict[str, JsonValue]]] = [
{"role": "user", "content": [{"type": "text", "text": "Weather in SF?", "cache_control": cache_control}]},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call-1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}}
],
},
{"role": "tool", "tool_call_id": "call-1", "content": '{"temp_f": 72}'},
]
tools: Final[list[dict[str, JsonValue]]] = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Current weather",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
},
},
web_search,
]
original: Final = deepcopy((messages, tools))
result: Final = config_type().transform_request_to_count_tokens(
model="claude-opus-5-5", messages=messages, tools=tools
)
assert result == {
"model": "claude-opus-5-5",
"messages": [
{"role": "user", "content": [{"type": "text", "text": "Weather in SF?", "cache_control": cache_control}]},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "call-1", "name": "get_weather", "input": {"city": "SF"}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "call-1", "content": '{"temp_f": 72}'}],
},
],
"tools": [
{
"type": "custom",
"name": "get_weather",
"description": "Current weather",
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}},
},
web_search,
],
}
assert (messages, tools) == original
def test_count_maps_only_function_tools_and_sends_the_rest_as_sent() -> None:
"""Only a chat ``function`` tool goes through the chat tool mapping, so a native computer tool next to
it, which that mapping only reads from a ``function`` wrapper, still reaches count_tokens as sent."""
messages: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Open the settings page"}]
native_tools: Final[list[dict[str, JsonValue]]] = [
{"name": "read_file", "input_schema": {"type": "object"}},
{"type": "computer_20250124", "name": "computer", "display_width_px": 1280, "display_height_px": 800},
]
function_tool: Final[dict[str, JsonValue]] = {
"type": "function",
"function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}},
}
result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens(
model="claude-opus-5-5", messages=messages, tools=[native_tools[0], function_tool, native_tools[1]]
)
assert result == {
"model": "claude-opus-5-5",
"messages": messages,
"tools": [
native_tools[0],
{"type": "custom", "name": "get_weather", "input_schema": {"type": "object", "properties": {}}},
native_tools[1],
],
}
def test_count_translates_a_dumped_sdk_message_without_tool_calls() -> None:
"""An OpenAI SDK message dumped into the history carries ``tool_calls: None`` and other chat-only keys
on a plain assistant turn, so it is translated too and only role and content reach count_tokens."""
result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens(
model="claude-opus-5-5",
messages=[
{"role": "user", "content": "Weather in SF?"},
{"role": "assistant", "content": "72F", "tool_calls": None, "function_call": None, "refusal": None},
],
)
assert result == {
"model": "claude-opus-5-5",
"messages": [
{"role": "user", "content": [{"type": "text", "text": "Weather in SF?"}]},
{"role": "assistant", "content": [{"type": "text", "text": "72F"}]},
],
}
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport_clients")
async def test_handler_counts_a_tool_call_turn_without_content() -> None:
"""OpenAI makes ``content`` optional on an assistant ``tool_calls`` turn and SDK dumps that drop None
leave it out, so the request still reaches count_tokens instead of failing validation and falling back."""
with respx.mock:
route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 42})
)
result: Final = await AnthropicCountTokensHandler().handle_count_tokens_request(
model="claude-opus-5-5",
messages=[
{"role": "user", "content": "Weather in SF?"},
{
"role": "assistant",
"tool_calls": [
{"id": "call-1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}
],
},
{"role": "tool", "tool_call_id": "call-1", "content": "72F"},
],
auth_header={"x-api-key": "sk-ant-api03-test-key"},
api_base="https://gateway.example",
)
assert result == {"input_tokens": 42}
assert TypeAdapter(dict[str, JsonValue]).validate_json(route.calls.last.request.content) == {
"model": "claude-opus-5-5",
"messages": [
{"role": "user", "content": [{"type": "text", "text": "Weather in SF?"}]},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "call-1", "name": "get_weather", "input": {}}],
},
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call-1", "content": "72F"}]},
],
}