mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +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 from c82e8244a5 (main).
This commit is contained in:
Mateo Wang
mateo-berri
parent
b6f1a033de
commit
c3d07e70d6
5 changed files with 721 additions and 25 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)
|
||||
|
|
@ -34,6 +34,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
|
||||
|
|
@ -343,6 +347,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:
|
||||
|
|
@ -400,6 +407,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(
|
||||
|
|
@ -463,7 +471,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")
|
||||
|
|
@ -557,6 +567,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:
|
||||
|
|
@ -564,26 +575,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(
|
||||
|
|
@ -1000,15 +1041,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):
|
||||
|
|
@ -1020,12 +1077,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 {}),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"<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", "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"<file {marker}-{index}>" in sent for index in range(_FILES)) == files_kept
|
||||
|
|
@ -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,9 @@ 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 itertools import accumulate, chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
|
@ -39,6 +40,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,
|
||||
|
|
@ -515,6 +517,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([deepcopy(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"""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue