mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
Merge pull request #47685 from BerriAI/litellm_backport_count_tokens_1_104_5_stable_1_104_x
fix(count_tokens): backport #47173 and #47353 to stable/1.104.x for v1.104.5
This commit is contained in:
yuneng-jiang
GitHub
commit
9372fc0b86
13 changed files with 1329 additions and 33 deletions
|
|
@ -0,0 +1,44 @@
|
|||
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}, # mutable-ok: serialized at once
|
||||
ensure_ascii=False,
|
||||
)
|
||||
for function in functions
|
||||
)
|
||||
|
||||
|
||||
_TOOL_SECTIONS: Final = (
|
||||
(
|
||||
re.compile(
|
||||
r"deepseek-v4(?![.\d]|p\d)|^(?:deepseek/)?deepseek-(?:chat|reasoner|flash)(?::free)?$", 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,56 @@ 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) # cast-ok: tool payloads are JSON objects with str keys
|
||||
|
||||
|
||||
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 {} # mutable-ok: read-only default
|
||||
normalized_parameters: Final = parameters if isinstance(parameters, Mapping) else {} # mutable-ok: never mutated
|
||||
unmapped_anthropic_keys: Final = (
|
||||
{ # mutable-ok: never mutated
|
||||
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 {} # mutable-ok: read-only default
|
||||
)
|
||||
description: Final = tool.get("description")
|
||||
strict: Final = tool.get("strict")
|
||||
yield { # mutable-ok: yielded as a read-only Mapping
|
||||
"name": tool.get("name"),
|
||||
"parameters": {**normalized_parameters, **unmapped_anthropic_keys}, # mutable-ok: never mutated
|
||||
**({"description": description} if description is not None else {}), # mutable-ok: spread at once
|
||||
**({"strict": bool(strict)} if strict is not None else {}), # mutable-ok: spread at once
|
||||
}
|
||||
|
||||
|
||||
def _format_object_parameters(parameters, indent):
|
||||
properties: Final = parameters.get("properties")
|
||||
if not properties:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,21 @@
|
|||
MIT License
|
||||
|
||||
Copyright (c) 2023 DeepSeek
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
|
@ -13793,7 +13793,7 @@ async def _loaded_tokenizer(
|
|||
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
|
||||
) -> SelectTokenizerResponse | BaseException:
|
||||
try:
|
||||
return await asyncify(litellm.utils._select_tokenizer, limiter=threads)(
|
||||
return await asyncify(litellm.utils.select_count_route_tokenizer, limiter=threads)(
|
||||
model=model, custom_tokenizer=custom_tokenizer
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
|
|
@ -13917,7 +13917,7 @@ async def _local_token_estimate(
|
|||
provider_failure: ProviderCountFailure | None,
|
||||
) -> _LocalTokenEstimate:
|
||||
tokenizer: Final = (
|
||||
await asyncify(litellm.utils._select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer)
|
||||
await asyncify(litellm.utils.select_count_route_tokenizer)(model=model, custom_tokenizer=custom_tokenizer)
|
||||
if provider_failure is None
|
||||
else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure)
|
||||
)
|
||||
|
|
@ -13927,6 +13927,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)
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import textwrap
|
|||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass, field
|
||||
from functools import lru_cache, wraps
|
||||
from importlib import resources
|
||||
|
|
@ -2386,12 +2387,12 @@ def _select_custom_tokenizer_helper(
|
|||
return create_pretrained_tokenizer(identifier=identifier, revision=revision, auth_token=auth_token)
|
||||
|
||||
|
||||
def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse:
|
||||
def _select_tokenizer_helper(model: str, *, count_route: bool = False) -> SelectTokenizerResponse:
|
||||
if litellm.disable_hf_tokenizer_download is True:
|
||||
return _return_openai_tokenizer(model)
|
||||
|
||||
try:
|
||||
result: Final = _return_huggingface_tokenizer(model)
|
||||
result: Final = _return_huggingface_tokenizer(model, count_route=count_route)
|
||||
if result is not None:
|
||||
return result
|
||||
except Exception as e:
|
||||
|
|
@ -2413,7 +2414,15 @@ def uses_anthropic_tokenizer(model: str) -> bool:
|
|||
return model in litellm.anthropic_models and "claude-3" not in model
|
||||
|
||||
|
||||
HuggingFaceTokenizerKind = Literal["cohere", "anthropic", "llama2", "llama3"]
|
||||
HuggingFaceTokenizerKind = Literal["cohere", "anthropic", "deepseek", "llama2", "llama3"]
|
||||
|
||||
_DEEPSEEK_TOKENIZER_MODEL: Final = re.compile(
|
||||
r"deepseek[-_.](?:chat[-_.])?v[34]|(?:^|/)deepseek-(?:chat|reasoner|flash)(?::free)?$", re.IGNORECASE
|
||||
)
|
||||
|
||||
_BUNDLED_TOKENIZER_LOAD_LOCKS: Final[Mapping[HuggingFaceTokenizerKind, threading.Lock]] = MappingProxyType(
|
||||
{"anthropic": threading.Lock(), "deepseek": threading.Lock()}
|
||||
)
|
||||
|
||||
|
||||
def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
|
||||
|
|
@ -2429,14 +2438,27 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
|
|||
return None
|
||||
|
||||
|
||||
def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
|
||||
kind: Final = huggingface_tokenizer_kind(model)
|
||||
def _return_huggingface_tokenizer(model: str, *, count_route: bool = False) -> SelectTokenizerResponse | None:
|
||||
kind: Final = (
|
||||
"deepseek" if count_route and _DEEPSEEK_TOKENIZER_MODEL.search(model) else huggingface_tokenizer_kind(model)
|
||||
)
|
||||
if kind is None:
|
||||
return None
|
||||
return {
|
||||
"type": "huggingface_tokenizer",
|
||||
"tokenizer": _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend()),
|
||||
}
|
||||
return _huggingface_tokenizer(kind)
|
||||
|
||||
|
||||
def _huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> SelectTokenizerResponse:
|
||||
with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()):
|
||||
tokenizer: Final = _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend())
|
||||
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} # mutable-ok: SelectTokenizerResponse shape
|
||||
|
||||
|
||||
def select_count_route_tokenizer(
|
||||
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None = None
|
||||
) -> SelectTokenizerResponse:
|
||||
if custom_tokenizer is None and _DEEPSEEK_TOKENIZER_MODEL.search(model):
|
||||
return _select_tokenizer_helper(model=model, count_route=True)
|
||||
return _select_tokenizer(model=model, custom_tokenizer=custom_tokenizer)
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
|
|
@ -2447,6 +2469,12 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio
|
|||
return tokenizer_dispatch.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer")
|
||||
case "anthropic":
|
||||
return tokenizer_dispatch.anthropic()
|
||||
case "deepseek":
|
||||
return tokenizer_dispatch.from_str(
|
||||
resources.files("litellm.litellm_core_utils.tokenizers")
|
||||
.joinpath("deepseek_tokenizer.json")
|
||||
.read_text(encoding="utf-8")
|
||||
)
|
||||
case "llama2":
|
||||
return tokenizer_dispatch.from_pretrained("hf-internal-testing/llama-tokenizer")
|
||||
case "llama3":
|
||||
|
|
@ -2543,6 +2571,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`.
|
||||
|
|
@ -2569,6 +2598,7 @@ def token_counter(
|
|||
tool_choice,
|
||||
use_default_image_token_count,
|
||||
default_token_count,
|
||||
chat_template_tools,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
35
tests/integration/_support/deepseek_tokens.py
Normal file
35
tests/integration/_support/deepseek_tokens.py
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
from collections.abc import Callable
|
||||
from functools import cache
|
||||
from importlib import resources
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.default_encoding import encoding as cl100k_base
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
CHINESE_TEXT: Final = "深度求索的分词器会把这段中文切成更少的词元。" * 20
|
||||
SHORT_TEXT: Final = "Find record 14."
|
||||
|
||||
|
||||
@cache
|
||||
def _deepseek_tokenizer() -> Tokenizer:
|
||||
bundled: Final = resources.files("litellm.litellm_core_utils.tokenizers").joinpath("deepseek_tokenizer.json")
|
||||
return Tokenizer.from_str(bundled.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def deepseek_tokens(text: str) -> int:
|
||||
return len(_deepseek_tokenizer().encode(text).ids)
|
||||
|
||||
|
||||
def cl100k_tokens(text: str) -> int:
|
||||
return len(cl100k_base.encode(text))
|
||||
|
||||
|
||||
def chinese_minus_short(tokens: Callable[[str], int]) -> int:
|
||||
return tokens(CHINESE_TEXT) - tokens(SHORT_TEXT)
|
||||
|
||||
|
||||
def assert_counted_with_deepseek(counted: int, expected: Callable[[Callable[[str], int]], int]) -> None:
|
||||
cl100k_count: Final = expected(cl100k_tokens)
|
||||
assert counted != cl100k_count, f"counted {counted}, the cl100k_base count, not DeepSeek's tokenizer"
|
||||
deepseek_count: Final = expected(deepseek_tokens)
|
||||
assert counted == deepseek_count, f"counted {counted}, DeepSeek's tokenizer counts {deepseek_count}"
|
||||
126
tests/integration/_support/openai_wire.py
Normal file
126
tests/integration/_support/openai_wire.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
from integration._support.wire import Reply, Request, Wire
|
||||
from pydantic import JsonValue
|
||||
|
||||
from tests.integration._support.provider import MODEL_DISCOVERY
|
||||
|
||||
_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
|
||||
|
||||
|
||||
def answering_model_discovery(respond: Callable[[Request], Reply]) -> Callable[[Request], Reply]:
|
||||
def guarded(request: Request) -> Reply:
|
||||
if (request.method, request.target) == MODEL_DISCOVERY:
|
||||
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
|
||||
return respond(request)
|
||||
|
||||
return guarded
|
||||
|
||||
|
||||
def posted_targets(wire: Wire) -> tuple[str, ...]:
|
||||
return tuple(request.target for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
def openai_error(status: int) -> Reply:
|
||||
return Reply(
|
||||
status=status,
|
||||
body=json.dumps({"error": {"message": f"scripted {status}", "type": "server_error", "code": None}}).encode(),
|
||||
)
|
||||
|
||||
|
||||
def _data_frame(frame: dict[str, JsonValue]) -> bytes:
|
||||
return b"data: " + json.dumps(frame).encode() + b"\n\n"
|
||||
|
||||
|
||||
def _typed_frame(event: dict[str, JsonValue]) -> bytes:
|
||||
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
|
||||
|
||||
|
||||
def chat_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]:
|
||||
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": model}
|
||||
role_only: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}]})
|
||||
content: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"content": text}}]})
|
||||
finish: Final = _data_frame(
|
||||
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": _USAGE}
|
||||
)
|
||||
return (role_only, content, finish + b"data: [DONE]\n\n")
|
||||
|
||||
|
||||
def chat_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply:
|
||||
if stream:
|
||||
return Reply(content_type="text/event-stream", chunks=chat_stream(identity, model, text))
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
|
||||
"usage": _USAGE,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _response_object(identity: str, model: str, text: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": model,
|
||||
"output": [_message_item(identity, text, "completed")],
|
||||
"usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8},
|
||||
}
|
||||
|
||||
|
||||
def _message_item(identity: str, text: str, status: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"msg_{identity}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": status,
|
||||
"content": [{"type": "output_text", "text": text, "annotations": []}] if status == "completed" else [],
|
||||
}
|
||||
|
||||
|
||||
def responses_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]:
|
||||
response: Final = _response_object(identity, model, text)
|
||||
item: Final = _message_item(identity, text, "in_progress")
|
||||
opened: Final = _typed_frame(
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**response, "status": "in_progress", "output": []},
|
||||
}
|
||||
) + _typed_frame({"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": item})
|
||||
delta: Final = _typed_frame(
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 2,
|
||||
"item_id": f"msg_{identity}",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
}
|
||||
)
|
||||
closed: Final = _typed_frame(
|
||||
{
|
||||
"type": "response.output_item.done",
|
||||
"sequence_number": 3,
|
||||
"output_index": 0,
|
||||
"item": _message_item(identity, text, "completed"),
|
||||
}
|
||||
) + _typed_frame({"type": "response.completed", "sequence_number": 4, "response": response})
|
||||
return (opened, delta, closed)
|
||||
|
||||
|
||||
def responses_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply:
|
||||
if stream:
|
||||
return Reply(content_type="text/event-stream", chunks=responses_stream(identity, model, text))
|
||||
return Reply(body=json.dumps(_response_object(identity, model, text)).encode())
|
||||
|
|
@ -79,7 +79,9 @@ def wire_server(
|
|||
self.send_header("connection", "close")
|
||||
self.end_headers()
|
||||
try:
|
||||
if reply.chunks is None:
|
||||
if self.command == "HEAD":
|
||||
self.wfile.flush()
|
||||
elif reply.chunks is None:
|
||||
self.wfile.write(reply.body)
|
||||
else:
|
||||
for index, chunk in enumerate(reply.chunks):
|
||||
|
|
@ -104,6 +106,8 @@ def wire_server(
|
|||
do_PUT = respond
|
||||
do_GET = respond
|
||||
do_DELETE = respond
|
||||
do_PATCH = respond
|
||||
do_HEAD = respond
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -0,0 +1,528 @@
|
|||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
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.deepseek_tokens import (
|
||||
CHINESE_TEXT,
|
||||
SHORT_TEXT,
|
||||
assert_counted_with_deepseek,
|
||||
chinese_minus_short,
|
||||
cl100k_tokens,
|
||||
)
|
||||
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 pydantic import JsonValue
|
||||
|
||||
from litellm.litellm_core_utils.default_encoding import encoding as cl100k_base
|
||||
|
||||
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"
|
||||
_DEEPSEEK_R1: Final = "azure_ai/DeepSeek-R1-0528"
|
||||
_QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b"
|
||||
_GROK: Final = "azure_ai/grok-4-1"
|
||||
_TOOL_CHOICE_CONTROL: Final = "azure_ai/grok-4"
|
||||
_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 = SHORT_TEXT
|
||||
_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], tokens: Callable[[str], int]) -> 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], text: str = _PROMPT) -> dict[str, JsonValue]:
|
||||
prompt: Final[dict[str, JsonValue]] = (
|
||||
{"input": text}
|
||||
if self.path == "/v1/responses/input_tokens"
|
||||
else {"messages": [{"role": "user", "content": text}]}
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
_TOOL_COUNTS: Final[Mapping[str, tuple[str, str]]] = {
|
||||
**{f"{route}-azure_ai-v4": (route, _DEEPSEEK_V4) for route in _COUNT_ROUTES},
|
||||
"messages-count-tokens-deepseek-chat": ("messages-count-tokens", "deepseek/deepseek-chat"),
|
||||
"messages-count-tokens-deepseek-reasoner": ("messages-count-tokens", "deepseek/deepseek-reasoner"),
|
||||
"messages-count-tokens-openrouter-deepseek-chat": ("messages-count-tokens", "openrouter/deepseek/deepseek-chat"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", sorted(_TOOL_COUNTS))
|
||||
def test_count_routes_count_deepseek_v4_tools_the_way_its_chat_template_renders_them(
|
||||
gateway: Gateway, case: str
|
||||
) -> None:
|
||||
route, backend = _TOOL_COUNTS[case]
|
||||
selected: Final = _COUNT_ROUTES[route]
|
||||
with wire_server(_completion) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, "count", backend, 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 == ()
|
||||
tool_tokens: Final = with_tool - without_tools
|
||||
assert tool_tokens > _DEEPSEEK_V4_TOOL_PREAMBLE, f"counted {tool_tokens} tool tokens, short of V4's tool preamble"
|
||||
assert_counted_with_deepseek(tool_tokens, lambda tokens: _deepseek_v4_tool_tokens(_FUNCTION, tokens))
|
||||
|
||||
|
||||
_TEXT_COUNTS: Final[Mapping[str, tuple[str, str]]] = {
|
||||
"messages-count-tokens-azure_ai-v4": ("messages-count-tokens", _DEEPSEEK_V4),
|
||||
"responses-input-tokens-azure_ai-v4": ("responses-input-tokens", _DEEPSEEK_V4),
|
||||
"messages-count-tokens-deepseek-chat": ("messages-count-tokens", "deepseek/deepseek-chat"),
|
||||
}
|
||||
|
||||
|
||||
def _text_count_delta(gateway: Gateway, route: _CountRoute, model: str) -> int:
|
||||
chinese: Final = _count(gateway, route, route.body(model, (), CHINESE_TEXT))
|
||||
short: Final = _count(gateway, route, route.body(model, (), SHORT_TEXT))
|
||||
return chinese - short
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", sorted(_TEXT_COUNTS))
|
||||
def test_count_routes_count_deepseek_v3_and_v4_text_with_deepseeks_tokenizer(gateway: Gateway, case: str) -> None:
|
||||
route, backend = _TEXT_COUNTS[case]
|
||||
with wire_server(_completion) as wire, gateway.scenario() as scenario:
|
||||
delta: Final = _text_count_delta(
|
||||
gateway, _COUNT_ROUTES[route], _deployment(scenario, "count", backend, wire.url)
|
||||
)
|
||||
upstream: Final = wire.drain()
|
||||
assert upstream == ()
|
||||
assert_counted_with_deepseek(delta, chinese_minus_short)
|
||||
|
||||
|
||||
def _counted_with_deepseek(delta: int) -> None:
|
||||
assert_counted_with_deepseek(delta, chinese_minus_short)
|
||||
|
||||
|
||||
def _counted_with_cl100k(delta: int) -> None:
|
||||
assert delta == chinese_minus_short(cl100k_tokens)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("backend", "tokenizer_type", "counted_with"),
|
||||
(
|
||||
(_DEEPSEEK_V4, "huggingface_tokenizer", _counted_with_deepseek),
|
||||
(_DEEPSEEK_R1, "openai_tokenizer", _counted_with_cl100k),
|
||||
),
|
||||
ids=("deepseek-v4", "deepseek-r1"),
|
||||
)
|
||||
def test_utils_token_counter_names_the_tokenizer_that_counted_deepseek_text(
|
||||
gateway: Gateway, backend: str, tokenizer_type: str, counted_with: Callable[[int], None]
|
||||
) -> None:
|
||||
route: Final = _COUNT_ROUTES["utils-token-counter"]
|
||||
with wire_server(_completion) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, "count", backend, wire.url)
|
||||
delta: Final = _text_count_delta(gateway, route, model)
|
||||
response: Final = gateway.request("POST", route.path, route.body(model, (), CHINESE_TEXT))
|
||||
upstream: Final = wire.drain()
|
||||
assert upstream == ()
|
||||
assert response.status_code == 200, response.text
|
||||
assert JSON_OBJECT.validate_json(response.content)["tokenizer_type"] == tokenizer_type, response.text
|
||||
counted_with(delta)
|
||||
|
||||
|
||||
@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_counted_with_deepseek(
|
||||
deferred - plain,
|
||||
lambda tokens: (
|
||||
_deepseek_v4_tool_tokens(deferred_function, tokens) - _deepseek_v4_tool_tokens(_FUNCTION, tokens)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
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 _owned_config(
|
||||
directory: Path,
|
||||
model_list: Sequence[JsonValue],
|
||||
general_settings: Mapping[str, JsonValue],
|
||||
router_settings: Mapping[str, JsonValue],
|
||||
) -> Path:
|
||||
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
config: Final = {
|
||||
**base,
|
||||
"model_list": list(model_list),
|
||||
"general_settings": {**object_value(base.get("general_settings", {})), **general_settings},
|
||||
"router_settings": {**object_value(base.get("router_settings", {})), **router_settings},
|
||||
}
|
||||
path: Final = directory / f"chat-template-tool-count-{uuid.uuid4().hex[:8]}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
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,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@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-{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 = _owned_config(
|
||||
directory,
|
||||
[
|
||||
_priced_deployment(group, _DEEPSEEK_V4, wire.url),
|
||||
_priced_deployment(grok_group, _TOOL_CHOICE_CONTROL, wire.url),
|
||||
],
|
||||
{"fail_closed_budget_enforcement": True},
|
||||
{},
|
||||
)
|
||||
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])
|
||||
|
||||
|
||||
def test_budget_reservation_counts_deepseek_v4_tools_like_models_without_a_chat_template(
|
||||
strict_budget: _BudgetRig,
|
||||
) -> None:
|
||||
selected: Final = _ENDPOINTS["chat-completions"]
|
||||
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
|
||||
|
||||
|
||||
_CODE_TEXT: Final = " for item in items:\n total += item.price * item.quantity\n" * 20
|
||||
|
||||
|
||||
def test_budget_reservation_counts_deepseek_v4_text_like_models_without_deepseeks_tokenizer(
|
||||
strict_budget: _BudgetRig,
|
||||
) -> None:
|
||||
chat: Final = _ENDPOINTS["chat-completions"]
|
||||
with strict_budget.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(models=[strict_budget.group, strict_budget.grok_group], max_budget=_TINY_BUDGET)
|
||||
estimates: Final = tuple(
|
||||
_estimated_cost(
|
||||
strict_budget.gateway.request(
|
||||
"POST",
|
||||
chat.path,
|
||||
{"model": group, "messages": [{"role": "user", "content": _CODE_TEXT}], "max_tokens": 1},
|
||||
key=key,
|
||||
)
|
||||
)
|
||||
for group in (strict_budget.group, strict_budget.grok_group)
|
||||
)
|
||||
assert strict_budget.wire.drain() == ()
|
||||
assert estimates[0] == estimates[1]
|
||||
|
||||
|
||||
def _project_key(scenario: Scenario, group: str, itpm_limit: int) -> str:
|
||||
team: Final = scenario.team()
|
||||
project: Final = scenario.project(team, model_itpm_limit={group: itpm_limit})
|
||||
return scenario.key(team_id=team, project_id=project)
|
||||
|
||||
|
||||
_TOOL_FRAMING_ITPM_LIMIT: Final = _deepseek_v4_tool_tokens(_FUNCTION, _tokens)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("alias", "model"), (("deepseek-v4-flash", _DEEPSEEK_V4), ("grok-4", _TOOL_CHOICE_CONTROL)))
|
||||
@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)
|
||||
key: Final = _project_key(scenario, group, _TOOL_FRAMING_ITPM_LIMIT)
|
||||
response: Final = _send(gateway, _ENDPOINTS[endpoint], group, key, 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, file_body: str) -> 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}>{file_body}",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _agent_body(group: str, marker: str, file_body: str, trigger: int) -> dict[str, JsonValue]:
|
||||
rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index, file_body) 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": trigger},
|
||||
"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)
|
||||
trigger: Final = _deepseek_v4_tool_tokens(_READ_FILE_FUNCTION, _tokens)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/messages", _agent_body(group, marker, "short file body.", trigger)
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
_BURST: Final = (CHINESE_TEXT, SHORT_TEXT) * 8
|
||||
|
||||
|
||||
def test_concurrent_first_deepseek_v4_counts_in_a_fresh_proxy_agree(tmp_path: Path) -> None:
|
||||
group: Final = f"count-{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 = _owned_config(
|
||||
tmp_path,
|
||||
[
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {"model": _DEEPSEEK_V4, "api_base": wire.url, "api_key": _API_KEY},
|
||||
}
|
||||
],
|
||||
{},
|
||||
{},
|
||||
)
|
||||
owned: Final = stack.enter_context(owned_proxy_process(rig_gateway, tmp_path, {}, config=config))
|
||||
with ThreadPoolExecutor(max_workers=len(_BURST)) as pool:
|
||||
counts: Final = tuple(
|
||||
pool.map(
|
||||
lambda text: _count(owned.gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(group, (), text)), _BURST
|
||||
)
|
||||
)
|
||||
upstream: Final = wire.drain()
|
||||
assert upstream == ()
|
||||
chinese: Final = frozenset(counts[0::2])
|
||||
short: Final = frozenset(counts[1::2])
|
||||
assert len(chinese) == 1 and len(short) == 1, counts
|
||||
assert_counted_with_deepseek(min(chinese) - min(short), chinese_minus_short)
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
_SCRIPT_SECONDS: Final = 60
|
||||
_CONCURRENT_LLAMA_COUNTS: Final = """
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
from litellm.utils import _select_tokenizer_helper
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
tuple(pool.map(lambda _: _select_tokenizer_helper("meta-llama/llama-3-8b"), range(2)))
|
||||
"""
|
||||
|
||||
|
||||
def test_concurrent_first_llama_counts_download_without_waiting_on_each_other(tmp_path: Path) -> None:
|
||||
both_downloading: Final = threading.Barrier(2, timeout=5)
|
||||
downloads: Final = SimpleQueue[str]()
|
||||
|
||||
def stalled_hub(request: Request) -> Reply:
|
||||
if request.target.endswith("/tokenizer.json"):
|
||||
try:
|
||||
both_downloading.wait()
|
||||
downloads.put("concurrent")
|
||||
except threading.BrokenBarrierError:
|
||||
downloads.put("serialized")
|
||||
return Reply(status=404)
|
||||
|
||||
with wire_server(stalled_hub) as hub:
|
||||
result: Final = subprocess.run(
|
||||
[sys.executable, "-I", "-c", _CONCURRENT_LLAMA_COUNTS, str(Path(litellm.__file__).parent.parent)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=_SCRIPT_SECONDS,
|
||||
check=False,
|
||||
env={
|
||||
**os.environ,
|
||||
"HF_ENDPOINT": hub.url,
|
||||
"HF_HOME": str(tmp_path / "home"),
|
||||
"HF_HUB_CACHE": str(tmp_path / "cache"),
|
||||
"HF_HUB_OFFLINE": "0",
|
||||
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
|
||||
},
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert downloads.get_nowait() == "concurrent"
|
||||
|
|
@ -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,11 +9,13 @@ import subprocess
|
|||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from concurrent.futures import Future, wait
|
||||
from collections.abc import Callable, Mapping
|
||||
from concurrent.futures import Future, ThreadPoolExecutor, wait
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from itertools import accumulate, chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from typing import Final, NoReturn
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import anyio.to_thread
|
||||
|
|
@ -39,6 +41,8 @@ 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 litellm.types.utils import CustomHuggingfaceTokenizer
|
||||
from tests.large_text import text
|
||||
from tests.unit.litellm_core_utils.event_loop_lag import (
|
||||
assert_loop_stayed_free,
|
||||
|
|
@ -443,6 +447,199 @@ 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/deepseek-chat", "DeepSeek-V4-Flash-0731"),
|
||||
("deepseek/deepseek-chat:free", "DeepSeek-V4-Flash-0731"),
|
||||
("deepseek-reasoner", "DeepSeek-V4-Flash-0731"),
|
||||
("deepseek/deepseek-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", "mistral-large-latest"),
|
||||
],
|
||||
)
|
||||
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"""
|
||||
|
||||
|
|
@ -686,7 +883,14 @@ def test_token_counter():
|
|||
|
||||
import unittest
|
||||
|
||||
from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding
|
||||
from litellm.utils import (
|
||||
_load_huggingface_tokenizer,
|
||||
_select_custom_tokenizer_helper,
|
||||
_select_tokenizer_helper,
|
||||
claude_json_str,
|
||||
encoding,
|
||||
select_count_route_tokenizer,
|
||||
)
|
||||
|
||||
# Clear the cache at module load to ensure clean state
|
||||
_load_huggingface_tokenizer.cache_clear()
|
||||
|
|
@ -785,6 +989,129 @@ class TestTokenizerSelection(unittest.TestCase):
|
|||
monkeypatch.undo()
|
||||
|
||||
|
||||
_CODE_SAMPLE: Final = " for item in items:\n total += item.price * item.quantity\n" * 20
|
||||
|
||||
|
||||
def _deepseek_token_ids(text: str) -> list[int]:
|
||||
deepseek_tokenizer: Final = Tokenizer.from_str(
|
||||
importlib.resources.files("litellm.litellm_core_utils.tokenizers")
|
||||
.joinpath("deepseek_tokenizer.json")
|
||||
.read_text(encoding="utf-8")
|
||||
)
|
||||
return deepseek_tokenizer.encode(text).ids
|
||||
|
||||
|
||||
def _deepseek_token_count(text: str) -> int:
|
||||
return len(_deepseek_token_ids(text))
|
||||
|
||||
|
||||
def _count_route_count(model: str, text: str) -> int:
|
||||
return litellm.token_counter(model=model, text=text, custom_tokenizer=select_count_route_tokenizer(model))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"azure_ai/DeepSeek-V4-Flash-0731",
|
||||
"deepseek/deepseek-v4-pro",
|
||||
"fireworks_ai/deepseek-v4p1-flash",
|
||||
"bedrock/us-east-1/deepseek.v3.2",
|
||||
"together_ai/deepseek-ai/DeepSeek-V3.1",
|
||||
"databricks/databricks-deepseek-v4-flash-0731",
|
||||
"replicate/deepseek-ai/deepseek-v3",
|
||||
"deepseek-chat",
|
||||
"deepseek/deepseek-reasoner",
|
||||
"openrouter/deepseek/deepseek-chat-v3-0324",
|
||||
"openrouter/deepseek/deepseek-chat:free",
|
||||
],
|
||||
)
|
||||
def test_count_routes_count_deepseek_v3_and_v4_with_the_packaged_deepseek_tokenizer(model: str):
|
||||
expected: Final = _deepseek_token_count(_CODE_SAMPLE)
|
||||
|
||||
assert expected != len(tiktoken.get_encoding("cl100k_base").encode(_CODE_SAMPLE))
|
||||
assert _count_route_count(model, _CODE_SAMPLE) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"azure_ai/DeepSeek-V4-Flash-0731",
|
||||
"deepseek/deepseek-v4-pro",
|
||||
"fireworks_ai/deepseek-v4p1-flash",
|
||||
"bedrock/us-east-1/deepseek.v3.2",
|
||||
"together_ai/deepseek-ai/DeepSeek-V3.1",
|
||||
"databricks/databricks-deepseek-v4-flash-0731",
|
||||
"deepseek-chat",
|
||||
"deepseek/deepseek-reasoner",
|
||||
"openrouter/deepseek/deepseek-chat-v3-0324",
|
||||
"openrouter/deepseek/deepseek-chat:free",
|
||||
],
|
||||
)
|
||||
def test_token_counter_keeps_its_tokenizer_for_deepseek_v3_and_v4(model: str):
|
||||
assert litellm.token_counter(model=model, text=_CODE_SAMPLE) == litellm.token_counter(model="", text=_CODE_SAMPLE)
|
||||
|
||||
|
||||
def test_count_routes_use_a_deployments_custom_tokenizer_over_the_packaged_deepseek_one():
|
||||
whitespace_tokenizer: Final = Tokenizer.from_str(_word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()))
|
||||
configured: Final = CustomHuggingfaceTokenizer(
|
||||
identifier="my-org/deepseek-v4-tokenizer", revision="v2", auth_token="private-repo-read-token"
|
||||
)
|
||||
_select_custom_tokenizer_helper.cache_clear()
|
||||
|
||||
with patch.object(litellm.utils.tokenizer_dispatch, "from_pretrained", return_value=whitespace_tokenizer) as load:
|
||||
selected: Final = select_count_route_tokenizer("deepseek/deepseek-v4-flash", configured)
|
||||
|
||||
load.assert_called_once_with("my-org/deepseek-v4-tokenizer", revision="v2", token="private-repo-read-token")
|
||||
counted: Final = litellm.token_counter(model="deepseek-chat", text=_CODE_SAMPLE, custom_tokenizer=selected)
|
||||
assert counted == len(_CODE_SAMPLE.split())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", "fireworks_ai/accounts/fireworks/models/deepseek-v2p5"],
|
||||
)
|
||||
def test_count_routes_skip_the_deepseek_tokenizer_for_deepseek_models_built_on_other_vocabularies(model: str):
|
||||
assert _count_route_count(model, _CODE_SAMPLE) != _deepseek_token_count(_CODE_SAMPLE)
|
||||
|
||||
|
||||
def test_encode_and_decode_keep_their_tokenizer_for_deepseek_v4():
|
||||
token_ids: Final = encode(model="deepseek/deepseek-v4-flash", text=_CODE_SAMPLE)
|
||||
|
||||
assert token_ids == encode(model="", text=_CODE_SAMPLE)
|
||||
assert decode(model="deepseek/deepseek-v4-flash", tokens=token_ids) == _CODE_SAMPLE
|
||||
|
||||
|
||||
def test_disable_hf_tokenizer_download_keeps_deepseek_v4_count_routes_on_cl100k(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
|
||||
|
||||
assert _count_route_count("deepseek/deepseek-v4-flash", _CODE_SAMPLE) == litellm.token_counter(
|
||||
model="", text=_CODE_SAMPLE
|
||||
)
|
||||
|
||||
|
||||
def test_count_routes_fall_back_to_cl100k_when_the_deepseek_tokenizer_fails_to_load(monkeypatch: pytest.MonkeyPatch):
|
||||
def fail_to_load(kind: str, backend: object) -> NoReturn:
|
||||
raise RuntimeError(f"{kind} tokenizer failed to load on {backend}")
|
||||
|
||||
monkeypatch.setattr(litellm.utils, "_load_huggingface_tokenizer", fail_to_load)
|
||||
|
||||
assert _count_route_count("deepseek/deepseek-v4-flash", _CODE_SAMPLE) == _count_route_count("", _CODE_SAMPLE)
|
||||
|
||||
|
||||
def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer():
|
||||
_load_huggingface_tokenizer.cache_clear()
|
||||
start: Final = threading.Barrier(8)
|
||||
|
||||
def first_tokenizer(_: int) -> object:
|
||||
start.wait()
|
||||
return select_count_route_tokenizer("deepseek/deepseek-v4-flash")["tokenizer"]
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
tokenizers: Final = tuple(pool.map(first_tokenizer, range(8)))
|
||||
|
||||
assert len({id(tokenizer) for tokenizer in tokenizers}) == 1
|
||||
|
||||
|
||||
def test_token_counter_with_anthropic_tool_use():
|
||||
"""
|
||||
Test that _count_anthropic_content() correctly handles tool_use blocks.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue