From c82e8244a57862feecfb48218bdd20c72c2afe6d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:02:56 +0000 Subject: [PATCH] fix(count_tokens)!: count tools the way DeepSeek V4 and Qwen3.5 chat templates render them (#47173) * fix(count_tokens): count tools the way DeepSeek V4 and Qwen3.5 chat templates render them * refactor(count_tokens): hold the tool counting wiring for a follow-up commit * refactor(count_tokens): hold the chat template tool sections for a follow-up commit * fix(count_tokens): count tools the way DeepSeek V4 and Qwen3.5 chat templates render them * fix(count_tokens): skip tool_choice and non-function tools for chat template families, leave out Fireworks V4.1 spelling * fix(count_tokens): count Anthropic tools in the shape the messages bridge sends them A missing description is no longer rendered as null and strict is counted when set, matching what DeepSeek V4 and Qwen3.5 deployments receive. Adds tests that pin each template's schema rendering and its system turn. * test(count_tokens): check the prompt caching threshold agrees with token_counter for DeepSeek V4 and Qwen3.5 tools * fix(count_tokens): skip Anthropic web search tools and count unmapped Anthropic tool keys for chat template families * test(token_counter): cover chat-template tool counts across count routes, budgets, ITPM and context edits Integration cells check that DeepSeek V4 and Qwen3.5 tools are counted the way their chat templates render them on /v1/messages/count_tokens, /utils/token_counter, /v1/responses/input_tokens, strict budget reservation, project ITPM on chat, messages and responses, and the clear_tool_uses trigger, with grok controls keeping OpenAI framing --- .../prompt_templates/chat_template_tools.py | 37 ++ litellm/litellm_core_utils/token_counter.py | 117 ++++-- .../test_chat_template_tool_count_wire.py | 383 ++++++++++++++++++ ...ls_prompt_templates_chat_template_tools.py | 31 ++ .../litellm_core_utils/test_token_counter.py | 189 ++++++++- 5 files changed, 728 insertions(+), 29 deletions(-) create mode 100644 litellm/litellm_core_utils/prompt_templates/chat_template_tools.py create mode 100644 tests/integration/providers/test_chat_template_tool_count_wire.py create mode 100644 tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py new file mode 100644 index 00000000000..3c437ca7fbb --- /dev/null +++ b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py @@ -0,0 +1,37 @@ +import json +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final + + +@dataclass(frozen=True, slots=True) +class ChatTemplateToolSection: + preamble_tokens: int + implicit_system_turn_tokens: int + render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] + + +def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions) + + +def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "".join( + "\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions + ) + + +_TOOL_SECTIONS: Final = ( + ( + re.compile(r"deepseek-v4(?![.\d]|p\d)", re.IGNORECASE), + ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas), + ), + (re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)), +) + + +def chat_template_tool_section(model: str | None) -> ChatTemplateToolSection | None: + if not model: + return None + return next((section for pattern, section in _TOOL_SECTIONS if pattern.search(model)), None) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 7ac3e9c80c5..8aded69aed2 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -35,6 +35,10 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( + ChatTemplateToolSection, + chat_template_tool_section, +) from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding if TYPE_CHECKING: @@ -347,6 +351,9 @@ EXTRAPOLATION_SAMPLES: Final = 16 T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") _COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") +_ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE: Final = frozenset( + ("name", "input_schema", "description", "cache_control", "strict", "type", "eager_input_streaming") +) def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: @@ -404,6 +411,7 @@ class _MessageCountParams: self.tokens_per_message = 3 self.tokens_per_name = 1 self.count_function = _get_count_function(model, custom_tokenizer) + self.tool_section = chat_template_tool_section(model) def token_counter( @@ -467,7 +475,9 @@ def token_counter( num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count) if count_response_tokens is False: includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) - num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message) + num_tokens += _count_extra( + params.count_function, tools, tool_choice, includes_system_message, params.tool_section + ) else: raise ValueError("Either text or messages must be provided") @@ -501,7 +511,8 @@ def messages_reach_token_count( _count_messages(params, [message], use_default_image_token_count, None) for message in new_messages ) running_totals: Final = accumulate( - per_message_counts, initial=_count_extra(params.count_function, tools, None, includes_system_message) + per_message_counts, + initial=_count_extra(params.count_function, tools, None, includes_system_message, params.tool_section), ) return any(total >= threshold for total in running_totals) @@ -592,6 +603,7 @@ def _count_extra( tools: list[ChatCompletionToolParam] | None, tool_choice: ChatCompletionNamedToolChoiceParam | None, includes_system_message: bool, + tool_section: ChatTemplateToolSection | None, ) -> int: """Count extra tokens for function definitions and tool choices. Args: @@ -599,26 +611,56 @@ def _count_extra( tools (Optional[List[ChatCompletionToolParam]]): The available tools. tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. includes_system_message (bool): Whether the messages include a system message. + tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. """ + reply_primer_tokens: Final = 3 + tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice) + return ( + reply_primer_tokens + + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) + + tool_choice_tokens + ) - num_tokens = 3 # every reply is primed with <|start|>assistant<|message|> - if tools: - num_tokens += count_function(_format_function_definitions(tools)) - num_tokens += 9 # Additional tokens for function definition of tools - # If there's a system message and tools are present, subtract four tokens - if tools and includes_system_message: - num_tokens -= 4 - # If tool_choice is 'none', add one token. - # If it's an object, add 4 + the number of tokens in the function name. - # If it's undefined or 'auto', don't add anything. +def _count_tool_definitions( + count_function: TokenCounterFunction, + tools: Sequence[object] | None, + includes_system_message: bool, + tool_section: ChatTemplateToolSection | None, +) -> int: + if not tools: + return 0 + if tool_section is not None: + return _count_chat_template_tools(count_function, tools, includes_system_message, tool_section) + system_message_discount: Final = 4 if includes_system_message else 0 + return count_function(_format_function_definitions(tools)) + 9 - system_message_discount + + +def _count_tool_choice( + count_function: TokenCounterFunction, + tool_choice: ChatCompletionNamedToolChoiceParam | None, +) -> int: if tool_choice == "none": - num_tokens += 1 - elif isinstance(tool_choice, dict): - num_tokens += 7 - num_tokens += count_function(str(tool_choice["function"]["name"])) + return 1 + if isinstance(tool_choice, dict): + return 7 + count_function(str(tool_choice["function"]["name"])) + return 0 - return num_tokens + +def _count_chat_template_tools( + count_function: TokenCounterFunction, + tools: Sequence[object], + includes_system_message: bool, + tool_section: ChatTemplateToolSection, +) -> int: + functions: Final = tuple( + function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name") + ) + if not functions: + return 0 + system_turn_tokens: Final = 0 if includes_system_message else tool_section.implicit_system_turn_tokens + schemas: Final = tool_section.render_tool_schemas(functions) + return tool_section.preamble_tokens + system_turn_tokens + count_function(schemas) def _get_extrapolating_count_function( @@ -1034,15 +1076,31 @@ def _format_function_definitions(tools: Sequence[object]) -> str: lines: Final = [] lines.append("namespace functions {") lines.append("") - for tool in tools: - if not isinstance(tool, Mapping): - continue - for function in _function_definitions_for_tool(cast(Mapping[str, object], tool)): - lines.extend(_format_single_function_definition(function)) + for function in _function_definitions(tools): + lines.extend(_format_single_function_definition(function)) lines.append("} // namespace functions") return "\n".join(lines) +def _tools_sent_as_functions(tools: Sequence[object]) -> tuple[Mapping[str, object], ...]: + return tuple(tool for tool in _tool_mappings(tools) if not _is_anthropic_web_search_tool(tool)) + + +def _is_anthropic_web_search_tool(tool: Mapping[str, object]) -> bool: + return tool.get("name") == "web_search" and "input_schema" not in tool and tool.get("type") != "function" + + +def _tool_mappings(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + for tool in tools: + if isinstance(tool, Mapping): + yield cast(Mapping[str, object], tool) + + +def _function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + for tool in _tool_mappings(tools): + yield from _function_definitions_for_tool(tool) + + def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mapping[str, object]]: function: Final = tool.get("function") if isinstance(function, Mapping): @@ -1054,12 +1112,21 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi if isinstance(declaration, Mapping): yield declaration return - parameters: Final = tool.get("input_schema") or tool.get("parameters") or {} + input_schema: Final = tool.get("input_schema") + parameters: Final = input_schema or tool.get("parameters") or {} normalized_parameters: Final = parameters if isinstance(parameters, Mapping) else {} + unmapped_anthropic_keys: Final = ( + {key: value for key, value in tool.items() if key not in _ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE} + if input_schema is not None + else {} + ) + description: Final = tool.get("description") + strict: Final = tool.get("strict") yield { "name": tool.get("name"), - "description": tool.get("description"), - "parameters": normalized_parameters, + "parameters": {**normalized_parameters, **unmapped_anthropic_keys}, + **({"description": description} if description is not None else {}), + **({"strict": bool(strict)} if strict is not None else {}), } diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py new file mode 100644 index 00000000000..d2b85e7168f --- /dev/null +++ b/tests/integration/providers/test_chat_template_tool_count_wire.py @@ -0,0 +1,383 @@ +import json +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from litellm.litellm_core_utils.default_encoding import encoding as cl100k_base +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_API_KEY: Final = "scripted-provider-key" +_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" +_QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b" +_GROK: Final = "azure_ai/grok-4-1" +_DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222 +_QWEN3_5_TOOL_PREAMBLE: Final = 196 +_QWEN3_5_OPENED_SYSTEM_TURN: Final = 5 +_OPENAI_TOOL_OVERHEAD: Final = 9 +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], +} +_FUNCTION: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": _SCHEMA} +_OPENAI_NAMESPACE: Final = "namespace functions {\n\ntype lookup = (_: {\nquery: string,\n}) => any;\n\n} // namespace functions" +_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {"name": "lookup", "input_schema": _SCHEMA} +_OPENAI_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "function": _FUNCTION} +_RESPONSES_TOOL: Final[dict[str, JsonValue]] = {"type": "function", **_FUNCTION} +_NAMED_TOOL_CHOICE: Final[dict[str, JsonValue]] = {"type": "function", "function": {"name": "lookup"}} +_WEB_SEARCH_TOOL: Final[dict[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 3} +_PROMPT: Final = "Find record 14." +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _PROMPT}] +_INPUT_PRICE: Final = 0.001 +_OUTPUT_PRICE: Final = 0.002 +_TINY_BUDGET: Final = 0.000000001 +_ESTIMATED_COST: Final = re.compile(r"Estimated request cost: ([^\s,]+)") + + +def _tokens(text: str) -> int: + return len(cl100k_base.encode(text)) + + +def _deepseek_v4_tool_tokens(function: Mapping[str, JsonValue]) -> int: + return _DEEPSEEK_V4_TOOL_PREAMBLE + _tokens(json.dumps(function, ensure_ascii=False)) + + +def _qwen3_5_tool_tokens(function: Mapping[str, JsonValue], *, with_system: bool) -> int: + opened_system_turn: Final = 0 if with_system else _QWEN3_5_OPENED_SYSTEM_TURN + schema: Final = json.dumps({"type": "function", "function": function}, ensure_ascii=False) + return _QWEN3_5_TOOL_PREAMBLE + opened_system_turn + _tokens("\n" + schema) + + +def _completion(_request: Request) -> Reply: + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "scripted-model", "OK", stream=False) + + +def _deployment(scenario: Scenario, group: str, model: str, api_base: str) -> str: + name: Final = f"{group}-{uuid.uuid4().hex}" + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": _API_KEY}, + "model_info": {}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +@dataclass(frozen=True, slots=True) +class _CountRoute: + path: str + field: str + tool: dict[str, JsonValue] + + def body(self, model: str, tools: Sequence[JsonValue]) -> dict[str, JsonValue]: + prompt: Final[dict[str, JsonValue]] = ( + {"input": _PROMPT} if self.path == "/v1/responses/input_tokens" else {"messages": _MESSAGES} + ) + return {"model": model, **prompt, **({"tools": list(tools)} if tools else {})} + + +_COUNT_ROUTES: Final[Mapping[str, _CountRoute]] = { + "messages-count-tokens": _CountRoute("/v1/messages/count_tokens", "input_tokens", _ANTHROPIC_TOOL), + "utils-token-counter": _CountRoute("/utils/token_counter", "total_tokens", _OPENAI_TOOL), + "responses-input-tokens": _CountRoute("/v1/responses/input_tokens", "input_tokens", _RESPONSES_TOOL), +} +_MESSAGES_COUNT: Final = _COUNT_ROUTES["messages-count-tokens"] + + +def _count(gateway: Gateway, route: _CountRoute, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", route.path, body) + assert response.status_code == 200, response.text + counted: Final = JSON_OBJECT.validate_json(response.content)[route.field] + assert isinstance(counted, int), response.text + return counted + + +@pytest.mark.parametrize("route", sorted(_COUNT_ROUTES)) +def test_count_routes_count_deepseek_v4_tools_the_way_its_chat_template_renders_them( + gateway: Gateway, route: str +) -> None: + selected: Final = _COUNT_ROUTES[route] + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url) + without_tools: Final = _count(gateway, selected, selected.body(model, ())) + with_tool: Final = _count(gateway, selected, selected.body(model, (selected.tool,))) + upstream: Final = wire.drain() + assert upstream == () + assert with_tool - without_tools == _deepseek_v4_tool_tokens(_FUNCTION) + + +@pytest.mark.parametrize("with_system", (False, True), ids=("no-system", "system")) +def test_messages_count_tokens_counts_qwen3_5_tools_the_way_its_chat_template_renders_them( + gateway: Gateway, with_system: bool +) -> None: + system: Final[dict[str, JsonValue]] = {"system": "Be brief."} if with_system else {} + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _QWEN3_5, wire.url) + without_tools: Final = _count(gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, ()), **system}) + with_tool: Final = _count( + gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,)), **system} + ) + upstream: Final = wire.drain() + assert upstream == () + assert with_tool - without_tools == _qwen3_5_tool_tokens(_FUNCTION, with_system=with_system) + + +def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_for_deepseek_v4( + gateway: Gateway, +) -> None: + deferred_function: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": {**_SCHEMA, "defer_loading": True}} + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url) + plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) + with_web_search: Final = _count( + gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL, _WEB_SEARCH_TOOL)) + ) + deferred: Final = _count( + gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},)) + ) + upstream: Final = wire.drain() + assert upstream == () + assert with_web_search == plain + assert deferred - plain == _deepseek_v4_tool_tokens(deferred_function) - _deepseek_v4_tool_tokens(_FUNCTION) + + +def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _GROK, wire.url) + without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ())) + plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) + deferred: Final = _count( + gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},)) + ) + upstream: Final = wire.drain() + assert upstream == () + assert plain - without_tools == _OPENAI_TOOL_OVERHEAD + _tokens(_OPENAI_NAMESPACE) + assert deferred == plain + + +@dataclass(frozen=True, slots=True) +class _Endpoint: + path: str + body: dict[str, JsonValue] + tools: dict[str, JsonValue] + + +_ENDPOINTS: Final[Mapping[str, _Endpoint]] = { + "chat-completions": _Endpoint( + "/v1/chat/completions", + {"messages": _MESSAGES, "max_tokens": 1}, + {"tools": [_OPENAI_TOOL], "tool_choice": _NAMED_TOOL_CHOICE}, + ), + "messages": _Endpoint( + "/v1/messages", + {"messages": _MESSAGES, "max_tokens": 1}, + {"tools": [_ANTHROPIC_TOOL], "tool_choice": {"type": "auto"}}, + ), + "responses": _Endpoint( + "/v1/responses", + {"input": _PROMPT, "max_output_tokens": 16}, + {"tools": [_RESPONSES_TOOL]}, + ), +} + + +def _send(gateway: Gateway, endpoint: _Endpoint, group: str, key: str, *, tools: bool) -> httpx.Response: + return gateway.request( + "POST", + endpoint.path, + {"model": group, **endpoint.body, **(endpoint.tools if tools else {})}, + key=key, + ) + + +@dataclass(frozen=True, slots=True) +class _BudgetRig: + gateway: Gateway + wire: Wire + group: str + + +def _strict_budget_config(directory: Path, group: str, api_base: str) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "model_list": [ + { + "model_name": group, + "litellm_params": { + "model": _DEEPSEEK_V4, + "api_base": api_base, + "api_key": _API_KEY, + "input_cost_per_token": _INPUT_PRICE, + "output_cost_per_token": _OUTPUT_PRICE, + }, + } + ], + "general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True}, + } + path: Final = directory / f"chat-template-tool-budget-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def strict_budget(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_BudgetRig]: + directory: Final = tmp_path_factory.mktemp("chat-template-tool-budget") + group: Final = f"deepseek-v4-flash-{uuid.uuid4().hex}" + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + wire: Final = stack.enter_context(wire_server(_completion)) + owned: Final = stack.enter_context( + owned_proxy_process(rig_gateway, directory, {}, config=_strict_budget_config(directory, group, wire.url)) + ) + yield _BudgetRig(owned.gateway, wire, group) + + +def _estimated_cost(response: httpx.Response) -> float: + assert response.status_code in (400, 422, 429), f"{response.status_code} {response.text}" + error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"]) + estimate: Final = _ESTIMATED_COST.search(string_value(error["message"])) + assert estimate is not None, response.text + return float(estimate[1]) + + +@pytest.mark.parametrize("endpoint", ("chat-completions", "messages")) +def test_budget_reservation_estimates_deepseek_v4_tools_the_way_its_chat_template_renders_them( + strict_budget: _BudgetRig, endpoint: str +) -> None: + selected: Final = _ENDPOINTS[endpoint] + with strict_budget.gateway.scenario() as scenario: + key: Final = scenario.key(models=[strict_budget.group], max_budget=_TINY_BUDGET) + without_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=False)) + with_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=True)) + assert strict_budget.wire.drain() == () + assert with_tools - without_tools == pytest.approx(_deepseek_v4_tool_tokens(_FUNCTION) * _INPUT_PRICE) + + +def _project_key(scenario: Scenario, group: str) -> str: + team: Final = scenario.team() + project: Final = scenario.project(team, model_itpm_limit={group: _deepseek_v4_tool_tokens(_FUNCTION)}) + return scenario.key(team_id=team, project_id=project) + + +@pytest.mark.parametrize("endpoint", sorted(_ENDPOINTS)) +def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_renders_them( + gateway: Gateway, endpoint: str +) -> None: + with wire_server(_completion) as wire, gateway.scenario() as scenario: + group: Final = _deployment(scenario, "deepseek-v4-flash", _DEEPSEEK_V4, wire.url) + response: Final = _send(gateway, _ENDPOINTS[endpoint], group, _project_key(scenario, group), tools=True) + upstream: Final = wire.drain() + assert response.status_code == 429, response.text + assert "model_per_project_itpm" in response.text, response.text + assert upstream == () + + +def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: + chat: Final = _ENDPOINTS["chat-completions"] + with wire_server(_completion) as wire, gateway.scenario() as scenario: + group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url) + response: Final = gateway.request( + "POST", + chat.path, + {"model": group, **chat.body, "tools": [_OPENAI_TOOL]}, + key=_project_key(scenario, group), + ) + upstream: Final = wire.drain() + assert response.status_code == 200, response.text + assert len(upstream) == 1, upstream + + +_FILES: Final = 3 +_READ_FILE_FUNCTION: Final[dict[str, JsonValue]] = { + "name": "read_file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_READ_FILE_TOOL: Final[dict[str, JsonValue]] = { + "name": "read_file", + "input_schema": _READ_FILE_FUNCTION["parameters"], +} + + +def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{index}", + "content": f"short file body.", + } + ], + }, + ) + + +def _agent_body(group: str, marker: str) -> dict[str, JsonValue]: + rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) + return { + "model": group, + "max_tokens": 16, + "tools": [_READ_FILE_TOOL], + "messages": [*rounds, {"role": "user", "content": "Summarize what you read."}], + "context_management": { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": _deepseek_v4_tool_tokens(_READ_FILE_FUNCTION)}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + } + + +@pytest.mark.parametrize( + ("model", "files_kept"), + ((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))), + ids=("deepseek-v4", "grok-4-1"), +) +def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them( + gateway: Gateway, model: str, files_kept: tuple[bool, ...] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_completion) as wire, gateway.scenario() as scenario: + group: Final = _deployment(scenario, "context", model, wire.url) + response: Final = gateway.request("POST", "/v1/messages", _agent_body(group, marker)) + upstream: Final = wire.drain() + assert response.status_code == 200, response.text + assert len(upstream) == 1, upstream + sent: Final = upstream[0].body.decode() + assert tuple(f"" in sent for index in range(_FILES)) == files_kept diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py new file mode 100644 index 00000000000..51249e01d4d --- /dev/null +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py @@ -0,0 +1,31 @@ +from typing import Final + +import pytest + +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section + +_DESCRIBED_FUNCTION: Final = {"name": "f", "description": "é", "parameters": {}} +_BARE_FUNCTION: Final = {"name": "g", "parameters": {}} + + +# Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by +# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17 +@pytest.mark.parametrize( + ("model", "rendered"), + [ + ( + "DeepSeek-V4-Flash-0731", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), + ( + "databricks-qwen35-122b-a10b", + '\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}' + '\n{"type": "function", "function": {"name": "g", "parameters": {}}}', + ), + ], +) +def test_chat_template_tool_section_renders_schemas_like_the_models_template(model: str, rendered: str): + section: Final = chat_template_tool_section(model) + + assert section is not None + assert section.render_tool_schemas((_DESCRIBED_FUNCTION, _BARE_FUNCTION)) == rendered diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 84c139e867b..6126d919180 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -9,7 +9,7 @@ import subprocess import sys import threading import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping from concurrent.futures import Future, wait from itertools import accumulate, chain from pathlib import Path @@ -39,6 +39,7 @@ from litellm.litellm_core_utils.token_counter import ( offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -515,6 +516,183 @@ def test_token_counter_skips_non_mapping_tools(): assert mixed_tokens == valid_tokens +_WEATHER_FUNCTION: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string", "description": "City name"}}, + "required": ["location"], + }, +} +_TIME_FUNCTION: Final = { + "name": "get_time", + "description": "Get the local time for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, +} +_WEB_SEARCH_FUNCTION: Final = { + "name": "web_search", + "description": "Search the web", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, +} + + +def _openai_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"type": "function", "function": function} + + +def _anthropic_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"name": function["name"], "description": function["description"], "input_schema": function["parameters"]} + + +_USER_MESSAGES: Final = ({"role": "user", "content": "What is the weather in Paris today?"},) +_SYSTEM_AND_USER_MESSAGES: Final = ({"role": "system", "content": "You are a helpful assistant."}, *_USER_MESSAGES) + + +@pytest.mark.parametrize("tool_shape", [_openai_tool, _anthropic_tool]) +@pytest.mark.parametrize( + ("model", "messages", "functions", "billed_prompt_tokens"), + [ + # usage.prompt_tokens billed for this request by Azure AI Foundry DeepSeek-V4-Flash-0731, + # Databricks databricks-qwen35-122b-a10b and DashScope qwen3.5-plus on 2026-10-17 + ("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION,), 295), + ("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 348), + ("DeepSeek-V4-Flash-0731", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 300), + ("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION,), 289), + ("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 352), + ("databricks-qwen35-122b-a10b", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 296), + ("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION,), 284), + ("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 342), + ], +) +def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates( + tool_shape: Callable[[Mapping[str, object]], dict[str, object]], + model: str, + messages: tuple[dict[str, str], ...], + functions: tuple[Mapping[str, object], ...], + billed_prompt_tokens: int, +): + counted: Final = token_counter_new( + model=model, + messages=list(messages), + tools=[tool_shape(function) for function in functions], + ) + + assert abs(counted - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +@pytest.mark.parametrize( + ("model", "counted_like"), + [ + ("DeepSeek-V4-Pro", "DeepSeek-V4-Flash-0731"), + ("accounts/fireworks/models/deepseek-v4-flash", "DeepSeek-V4-Flash-0731"), + ("DeepSeek-V4.1-Flash", "deepseek-v3.2"), + ("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"), + ("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"), + ("qwen3-235b-a22b", "deepseek-v3.2"), + ], +) +def test_token_counter_counts_tools_by_the_models_chat_template_family(model: str, counted_like: str): + tools: Final = [_openai_tool(_WEATHER_FUNCTION), _openai_tool(_TIME_FUNCTION)] + + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) == token_counter_new( + model=counted_like, messages=list(_USER_MESSAGES), tools=tools + ) + + +@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) +@pytest.mark.parametrize( + "builtin_tool", + [ + {"type": "web_search_preview"}, + {"type": "web_search_20250305", "name": "web_search", "max_uses": 5}, + {"name": "web_search"}, + ], +) +def test_token_counter_skips_tools_not_sent_as_functions_for_chat_template_families( + model: str, builtin_tool: dict[str, object] +): + weather_tool: Final = _openai_tool(_WEATHER_FUNCTION) + + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool]) == token_counter_new( + model=model, messages=list(_USER_MESSAGES) + ) + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool, weather_tool] + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[weather_tool]) + + +def _responses_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"type": "function", **function} + + +def _gemini_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"function_declarations": [function]} + + +@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) +@pytest.mark.parametrize("function", [_WEATHER_FUNCTION, _WEB_SEARCH_FUNCTION]) +@pytest.mark.parametrize("tool_shape", [_responses_tool, _gemini_tool]) +def test_token_counter_counts_function_tools_like_chat_completions_tools_for_chat_template_families( + model: str, function: Mapping[str, object], tool_shape: Callable[[Mapping[str, object]], dict[str, object]] +): + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[tool_shape(function)] + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool(function)]) + + +# Azure AI Foundry DeepSeek-V4-Flash-0731 billed 295 prompt tokens for the weather tool with each of +# these tool_choice values on 2026-10-17, the same as without one +@pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) +def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: object): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + + assert token_counter_new( + model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice + ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) + + +@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) +@pytest.mark.parametrize( + "anthropic_tool", + [ + {"name": "get_weather", "input_schema": _WEATHER_FUNCTION["parameters"]}, + _anthropic_tool(_WEATHER_FUNCTION), + {**_anthropic_tool(_WEATHER_FUNCTION), "strict": True}, + {**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]}, + _anthropic_tool(_WEB_SEARCH_FUNCTION), + ], +) +def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families( + model: str, anthropic_tool: dict[str, object] +): + sent_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([anthropic_tool]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict + + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[anthropic_tool]) == token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=sent_tools + ) + + +# Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a +# request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn +# either way (both read on Hugging Face on 2026-10-17) +@pytest.mark.parametrize( + ("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)] +) +def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message( + model: str, implicit_system_turn_tokens: int +): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + system_cost_with_tools: Final = token_counter_new( + model=model, messages=list(_SYSTEM_AND_USER_MESSAGES), tools=tools + ) - token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + system_cost_without_tools: Final = token_counter_new( + model=model, messages=list(_SYSTEM_AND_USER_MESSAGES) + ) - token_counter_new(model=model, messages=list(_USER_MESSAGES)) + + assert system_cost_with_tools == system_cost_without_tools - implicit_system_turn_tokens + + class NeedsToleranceUpdateError(Exception): """Custom exception to mark tests that have improved""" @@ -1770,14 +1948,17 @@ _THRESHOLD_TEST_TOOLS: Final = [ ] -def test_messages_reach_token_count_agrees_with_token_counter_at_every_threshold() -> None: +@pytest.mark.parametrize( + "model", ["claude-3-5-sonnet-20240620", "DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"] +) +def test_messages_reach_token_count_agrees_with_token_counter_at_every_threshold(model: str) -> None: """The threshold check is the same arithmetic as token_counter(...) >= threshold, including the tools and system-message adjustments, so the boundary values must agree exactly.""" from litellm.litellm_core_utils.token_counter import messages_reach_token_count messages = _threshold_test_messages(turns=12) total = token_counter_new( - model="claude-3-5-sonnet-20240620", + model=model, messages=messages, tools=_THRESHOLD_TEST_TOOLS, use_default_image_token_count=True, @@ -1785,7 +1966,7 @@ def test_messages_reach_token_count_agrees_with_token_counter_at_every_threshold assert total > 100 for threshold in (0, 1, total - 1, total, total + 1, 10 * total): assert messages_reach_token_count( - model="claude-3-5-sonnet-20240620", + model=model, messages=messages, threshold=threshold, tools=_THRESHOLD_TEST_TOOLS,