mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(count_tokens)!: count tools the way DeepSeek V4 and Qwen3.5 chat templates render them (#47173)
Backport of #47173 to stable/1.104.x. Cherry-picked fromc82e8244a5(main), applied through the line's earlier resolutionc3d07e70d6. Adapted to the line: the default tool count keeps the line's formatter, since #43417 is not on the line, and only the chat-template path gets the function-definition helpers. Adapted to the stable compatibility rule: the chat-template tool count is opt-in through token_counter(chat_template_tools=True), which only the proxy count routes pass, so ITPM, budget reservation and context-management estimates keep the line's OpenAI tool framing and a request they admitted on v1.104.4 is still admitted. The tests that asserted the new framing on those paths now assert the line's framing is kept.
This commit is contained in:
Mateo Wang
mateo-berri
parent
5346a40c37
commit
6ce123ea21
7 changed files with 738 additions and 18 deletions
|
|
@ -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)
|
||||
|
|
@ -29,6 +29,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 Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding
|
||||
from litellm.litellm_core_utils.url_utils import safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
|
|
@ -338,6 +342,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:
|
||||
|
|
@ -380,6 +387,7 @@ class _MessageCountParams:
|
|||
self,
|
||||
model: str,
|
||||
custom_tokenizer: dict | SelectTokenizerResponse | None,
|
||||
chat_template_tools: bool,
|
||||
):
|
||||
from litellm.utils import print_verbose
|
||||
|
||||
|
|
@ -395,6 +403,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) if chat_template_tools else None
|
||||
|
||||
|
||||
def token_counter(
|
||||
|
|
@ -407,6 +416,7 @@ def token_counter(
|
|||
tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
|
||||
use_default_image_token_count: bool | None = False,
|
||||
default_token_count: int | None = None,
|
||||
chat_template_tools: bool = False,
|
||||
) -> int:
|
||||
"""
|
||||
Count the number of tokens in a given text using a specified model.
|
||||
|
|
@ -421,6 +431,7 @@ def token_counter(
|
|||
tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. Default is None.
|
||||
use_default_image_token_count (Optional[bool]): When True, will NOT make a GET request to the image URL and instead return the default image dimensions. Default is False.
|
||||
default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs. Default is None.
|
||||
chat_template_tools (bool): Count tools the way the model's chat template renders them, when it is known. Default is False.
|
||||
|
||||
Returns:
|
||||
int: The number of tokens in the text.
|
||||
|
|
@ -454,11 +465,13 @@ def token_counter(
|
|||
|
||||
elif messages is not None:
|
||||
new_messages: Final = cast(list[AllMessageValues], convert_list_message_to_dict(messages))
|
||||
params: Final = _MessageCountParams(model, custom_tokenizer)
|
||||
params: Final = _MessageCountParams(model, custom_tokenizer, chat_template_tools)
|
||||
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")
|
||||
|
|
@ -552,6 +565,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:
|
||||
|
|
@ -559,26 +573,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(
|
||||
|
|
@ -994,6 +1038,54 @@ def _format_function_definitions(tools):
|
|||
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):
|
||||
yield function
|
||||
return
|
||||
declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations")
|
||||
if isinstance(declarations, list):
|
||||
for declaration in declarations:
|
||||
if isinstance(declaration, Mapping):
|
||||
yield declaration
|
||||
return
|
||||
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"),
|
||||
"parameters": {**normalized_parameters, **unmapped_anthropic_keys},
|
||||
**({"description": description} if description is not None else {}),
|
||||
**({"strict": bool(strict)} if strict is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
def _format_object_parameters(parameters, indent):
|
||||
properties: Final = parameters.get("properties")
|
||||
if not properties:
|
||||
|
|
|
|||
|
|
@ -13922,6 +13922,7 @@ async def _local_token_estimate(
|
|||
messages=messages,
|
||||
tools=tools,
|
||||
custom_tokenizer=tokenizer,
|
||||
chat_template_tools=True,
|
||||
)
|
||||
total_tokens: Final = await (
|
||||
count if provider_failure is None else _estimate_or_provider_error(count, provider_failure)
|
||||
|
|
|
|||
|
|
@ -2546,6 +2546,7 @@ def token_counter(
|
|||
tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
|
||||
use_default_image_token_count: bool | None = False,
|
||||
default_token_count: int | None = None,
|
||||
chat_template_tools: bool = False,
|
||||
) -> int:
|
||||
"""
|
||||
The same as `litellm.litellm_core_utils.token_counter`.
|
||||
|
|
@ -2572,6 +2573,7 @@ def token_counter(
|
|||
tool_choice,
|
||||
use_default_image_token_count,
|
||||
default_token_count,
|
||||
chat_template_tools,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,365 @@
|
|||
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
|
||||
grok_group: str
|
||||
|
||||
|
||||
def _priced_deployment(group: str, model: str, api_base: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": group,
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_base": api_base,
|
||||
"api_key": _API_KEY,
|
||||
"input_cost_per_token": _INPUT_PRICE,
|
||||
"output_cost_per_token": _OUTPUT_PRICE,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _strict_budget_config(directory: Path, groups: Mapping[str, 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": [_priced_deployment(group, model, api_base) for group, model in groups.items()],
|
||||
"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}"
|
||||
grok_group: Final = f"grok-4-1-{uuid.uuid4().hex}"
|
||||
with ExitStack() as stack:
|
||||
rig_gateway: Final = stack.enter_context(gateway_from_environment())
|
||||
wire: Final = stack.enter_context(wire_server(_completion))
|
||||
config: Final = _strict_budget_config(directory, {group: _DEEPSEEK_V4, grok_group: _GROK}, wire.url)
|
||||
owned: Final = stack.enter_context(owned_proxy_process(rig_gateway, directory, {}, config=config))
|
||||
yield _BudgetRig(owned.gateway, wire, group, grok_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_counts_deepseek_v4_tools_like_models_without_a_chat_template(
|
||||
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, strict_budget.grok_group], max_budget=_TINY_BUDGET)
|
||||
deepseek: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=True))
|
||||
grok: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.grok_group, key, tools=True))
|
||||
assert strict_budget.wire.drain() == ()
|
||||
assert deepseek == grok
|
||||
|
||||
|
||||
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(("alias", "model"), (("deepseek-v4-flash", _DEEPSEEK_V4), ("grok-4-1", _GROK)))
|
||||
@pytest.mark.parametrize("endpoint", sorted(_ENDPOINTS))
|
||||
def test_project_itpm_admits_tools_under_a_limit_only_chat_template_framing_exceeds(
|
||||
gateway: Gateway, endpoint: str, alias: str, model: str
|
||||
) -> None:
|
||||
with wire_server(_completion) as wire, gateway.scenario() as scenario:
|
||||
group: Final = _deployment(scenario, alias, model, wire.url)
|
||||
response: Final = _send(gateway, _ENDPOINTS[endpoint], group, _project_key(scenario, group), tools=True)
|
||||
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"<file {marker}-{index}>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", (_DEEPSEEK_V4, _GROK), ids=("deepseek-v4", "grok-4-1"))
|
||||
def test_messages_context_management_trigger_keeps_openai_tool_framing(gateway: Gateway, model: str) -> 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"<file {marker}-{index}>" in sent for index in range(_FILES)) == (True,) * _FILES
|
||||
|
|
@ -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
|
||||
|
|
@ -9,8 +9,10 @@ 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 copy import deepcopy
|
||||
from functools import partial
|
||||
from itertools import accumulate, chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
|
@ -39,6 +41,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.experimental_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,
|
||||
|
|
@ -443,6 +446,195 @@ def test_token_counter_with_tools(message_count_pair):
|
|||
), f"Expected {expected_tokens} tokens, got {counted_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)
|
||||
|
||||
|
||||
_chat_template_count: Final = partial(token_counter_new, chat_template_tools=True)
|
||||
|
||||
|
||||
@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 = _chat_template_count(
|
||||
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 _chat_template_count(model=model, messages=list(_USER_MESSAGES), tools=tools) == _chat_template_count(
|
||||
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 _chat_template_count(
|
||||
model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool]
|
||||
) == _chat_template_count(model=model, messages=list(_USER_MESSAGES))
|
||||
assert _chat_template_count(
|
||||
model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool, weather_tool]
|
||||
) == _chat_template_count(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 _chat_template_count(
|
||||
model=model, messages=list(_USER_MESSAGES), tools=[tool_shape(function)]
|
||||
) == _chat_template_count(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 _chat_template_count(
|
||||
model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice
|
||||
) == _chat_template_count(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([deepcopy(anthropic_tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict
|
||||
|
||||
assert _chat_template_count(
|
||||
model=model, messages=list(_USER_MESSAGES), tools=[anthropic_tool]
|
||||
) == _chat_template_count(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 = _chat_template_count(
|
||||
model=model, messages=list(_SYSTEM_AND_USER_MESSAGES), tools=tools
|
||||
) - _chat_template_count(model=model, messages=list(_USER_MESSAGES), tools=tools)
|
||||
system_cost_without_tools: Final = _chat_template_count(
|
||||
model=model, messages=list(_SYSTEM_AND_USER_MESSAGES)
|
||||
) - _chat_template_count(model=model, messages=list(_USER_MESSAGES))
|
||||
|
||||
assert system_cost_with_tools == system_cost_without_tools - implicit_system_turn_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b", "qwen3.5-plus"])
|
||||
def test_token_counter_keeps_openai_tool_framing_for_chat_template_families_unless_asked(model: 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="mistral-large-latest", messages=list(_USER_MESSAGES), tools=tools
|
||||
)
|
||||
|
||||
|
||||
class NeedsToleranceUpdateError(Exception):
|
||||
"""Custom exception to mark tests that have improved"""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue