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,