mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
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:
parent
a5d05355b6
commit
f307eccf44
2 changed files with 188 additions and 10 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"}]},
|
||||
],
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue