fix(token_counter): count DeepSeek V4 text with DeepSeek's own tokenizer (#47353)

Local token counts for DeepSeek V3 and V4 model names now use DeepSeek's own bundled tokenizer instead of cl100k_base, and tool definitions on V4 names count with V4's chat-template tool framing
This commit is contained in:
Mateo Wang 2026-10-20 03:43:35 +00:00 • committed by GitHub
parent d56ff281de
commit 2d181e6633
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 491 additions and 51 deletions

View file

@ -24,7 +24,9 @@ def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str:
_TOOL_SECTIONS: Final = (
(
re.compile(r"deepseek-v4(?![.\d]|p\d)", re.IGNORECASE),
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)),

File diff suppressed because one or more lines are too long

View file

@ -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.

View file

@ -28,6 +28,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
@ -2589,7 +2590,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:
@ -2598,6 +2607,8 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
return "cohere"
if uses_anthropic_tokenizer(model):
return "anthropic"
if _DEEPSEEK_TOKENIZER_MODEL.search(model):
return "deepseek"
if "llama-2" in model.lower() or "replicate" in model.lower():
return "llama2"
if "llama-3" in model.lower():
@ -2609,10 +2620,9 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
kind: Final = huggingface_tokenizer_kind(model)
if kind is None:
return None
return {
"type": "huggingface_tokenizer",
"tokenizer": _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend()),
}
with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()):
tokenizer: Final = _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend())
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
@ -2623,6 +2633,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":

View 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}"

View file

@ -1,7 +1,8 @@
import json
import re
import uuid
from collections.abc import Iterator, Mapping, Sequence
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
@ -19,6 +20,13 @@ from integration._support.client import (
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
@ -29,6 +37,7 @@ 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"
_DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222
@ -47,7 +56,7 @@ _OPENAI_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "function": _FU
_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."
_PROMPT: Final = SHORT_TEXT
_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _PROMPT}]
_INPUT_PRICE: Final = 0.001
_OUTPUT_PRICE: Final = 0.002
@ -59,8 +68,8 @@ 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 _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:
@ -93,9 +102,11 @@ class _CountRoute:
field: str
tool: dict[str, JsonValue]
def body(self, model: str, tools: Sequence[JsonValue]) -> dict[str, JsonValue]:
def body(self, model: str, tools: Sequence[JsonValue], text: str = _PROMPT) -> dict[str, JsonValue]:
prompt: Final[dict[str, JsonValue]] = (
{"input": _PROMPT} if self.path == "/v1/responses/input_tokens" else {"messages": _MESSAGES}
{"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 {})}
@ -116,18 +127,85 @@ def _count(gateway: Gateway, route: _CountRoute, body: Mapping[str, JsonValue])
return counted
@pytest.mark.parametrize("route", sorted(_COUNT_ROUTES))
_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, route: str
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", _DEEPSEEK_V4, wire.url)
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 == ()
assert with_tool - without_tools == _deepseek_v4_tool_tokens(_FUNCTION)
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"))
@ -162,7 +240,12 @@ def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_fo
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)
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:
@ -215,17 +298,36 @@ def _send(gateway: Gateway, endpoint: _Endpoint, group: str, key: str, *, tools:
@dataclass(frozen=True, slots=True)
class _BudgetRig:
class _OwnedRig:
gateway: Gateway
wire: Wire
group: str
itpm_wire: Wire
itpm_group: str
def _strict_budget_config(directory: Path, group: str, api_base: str) -> Path:
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": [
"model_list": list(model_list),
"general_settings": {**object_value(base["general_settings"]), **general_settings},
"router_settings": {**object_value(base["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 _owned_rig_config(directory: Path, group: str, api_base: str, itpm_group: str, itpm_api_base: str) -> Path:
return _owned_config(
directory,
[
{
"model_name": group,
"litellm_params": {
@ -235,26 +337,35 @@ def _strict_budget_config(directory: Path, group: str, api_base: str) -> Path:
"input_cost_per_token": _INPUT_PRICE,
"output_cost_per_token": _OUTPUT_PRICE,
},
}
},
{
"model_name": itpm_group,
"litellm_params": {
"model": _DEEPSEEK_V4,
"api_base": itpm_api_base,
"api_key": _API_KEY,
"itpm": cl100k_tokens(CHINESE_TEXT),
},
"model_info": {"id": uuid.uuid4().hex},
},
],
"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
{"fail_closed_budget_enforcement": True},
{"optional_pre_call_checks": ["enforce_model_rate_limits"]},
)
@pytest.fixture(scope="module")
def strict_budget(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_BudgetRig]:
directory: Final = tmp_path_factory.mktemp("chat-template-tool-budget")
def owned_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_OwnedRig]:
directory: Final = tmp_path_factory.mktemp("chat-template-tool-count")
group: Final = f"deepseek-v4-flash-{uuid.uuid4().hex}"
itpm_group: Final = f"itpm-{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)
itpm_wire: Final = stack.enter_context(wire_server(_completion))
config: Final = _owned_rig_config(directory, group, wire.url, itpm_group, itpm_wire.url)
owned: Final = stack.enter_context(owned_proxy_process(rig_gateway, directory, {}, config=config))
yield _OwnedRig(owned.gateway, wire, group, itpm_wire, itpm_group)
def _estimated_cost(response: httpx.Response) -> float:
@ -267,30 +378,47 @@ def _estimated_cost(response: httpx.Response) -> float:
@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
owned_rig: _OwnedRig, 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)
with owned_rig.gateway.scenario() as scenario:
key: Final = scenario.key(models=[owned_rig.group], max_budget=_TINY_BUDGET)
without_tools: Final = _estimated_cost(_send(owned_rig.gateway, selected, owned_rig.group, key, tools=False))
with_tools: Final = _estimated_cost(_send(owned_rig.gateway, selected, owned_rig.group, key, tools=True))
assert owned_rig.wire.drain() == ()
assert_counted_with_deepseek(
round((with_tools - without_tools) / _INPUT_PRICE), lambda tokens: _deepseek_v4_tool_tokens(_FUNCTION, tokens)
)
def _project_key(scenario: Scenario, group: str) -> str:
def _chinese_chat(group: str) -> dict[str, JsonValue]:
return {"model": group, "messages": [{"role": "user", "content": CHINESE_TEXT}], "max_tokens": 1}
def test_router_itpm_reserves_deepseek_v4_text_with_deepseeks_tokenizer(owned_rig: _OwnedRig) -> None:
response: Final = owned_rig.gateway.request("POST", "/v1/chat/completions", _chinese_chat(owned_rig.itpm_group))
upstream: Final = owned_rig.itpm_wire.drain()
assert response.status_code == 200, response.text
assert len(upstream) == 1, upstream
def _project_key(scenario: Scenario, group: str, itpm_limit: int) -> str:
team: Final = scenario.team()
project: Final = scenario.project(team, model_itpm_limit={group: _deepseek_v4_tool_tokens(_FUNCTION)})
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("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)
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 == 429, response.text
assert "model_per_project_itpm" in response.text, response.text
@ -305,13 +433,23 @@ def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gatewa
"POST",
chat.path,
{"model": group, **chat.body, "tools": [_OPENAI_TOOL]},
key=_project_key(scenario, group),
key=_project_key(scenario, group, _TOOL_FRAMING_ITPM_LIMIT),
)
upstream: Final = wire.drain()
assert response.status_code == 200, response.text
assert len(upstream) == 1, upstream
def test_project_itpm_reserves_deepseek_v4_text_with_deepseeks_tokenizer(gateway: Gateway) -> None:
with wire_server(_completion) as wire, gateway.scenario() as scenario:
group: Final = _deployment(scenario, "deepseek-v4-flash", _DEEPSEEK_V4, wire.url)
key: Final = _project_key(scenario, group, cl100k_tokens(CHINESE_TEXT))
response: Final = gateway.request("POST", "/v1/chat/completions", _chinese_chat(group), key=key)
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",
@ -323,7 +461,7 @@ _READ_FILE_TOOL: Final[dict[str, JsonValue]] = {
}
def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]:
def _read_file_round(marker: str, index: int, file_body: str) -> tuple[JsonValue, ...]:
return (
{"role": "user", "content": f"Read file {index}."},
{
@ -338,15 +476,15 @@ def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]:
{
"type": "tool_result",
"tool_use_id": f"toolu_{index}",
"content": f"<file {marker}-{index}>short file body.",
"content": f"<file {marker}-{index}>{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)))
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,
@ -356,7 +494,7 @@ def _agent_body(group: str, marker: str) -> dict[str, JsonValue]:
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": _deepseek_v4_tool_tokens(_READ_FILE_FUNCTION)},
"trigger": {"type": "input_tokens", "value": trigger},
"keep": {"type": "tool_uses", "value": 1},
}
]
@ -375,9 +513,60 @@ def test_messages_context_management_trigger_counts_tools_the_way_the_models_cha
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))
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)) == files_kept
def test_messages_context_management_trigger_counts_deepseek_v4_text_with_deepseeks_tokenizer(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(_completion) as wire, gateway.scenario() as scenario:
group: Final = _deployment(scenario, "context", _DEEPSEEK_V4, wire.url)
trigger: Final = _FILES * cl100k_tokens(CHINESE_TEXT)
response: Final = gateway.request("POST", "/v1/messages", _agent_body(group, marker, CHINESE_TEXT, trigger))
upstream: Final = wire.drain()
assert response.status_code == 200, response.text
assert len(upstream) == 1, upstream
sent: Final = upstream[0].body.decode()
kept: Final = tuple(f"<file {marker}-{index}>" in sent for index in range(_FILES))
assert kept == (True,) * _FILES, f"cleared tool results under a trigger only the cl100k_base count reaches: {kept}"
_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)

View file

@ -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"

View file

@ -6,6 +6,12 @@ from typing import Final
import pytest
from integration._support.client import Gateway, eventually, list_value
from integration._support.database import read_rows
from integration._support.deepseek_tokens import (
CHINESE_TEXT,
SHORT_TEXT,
assert_counted_with_deepseek,
chinese_minus_short,
)
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
@ -126,3 +132,38 @@ def test_a_provider_500_failure_row_keeps_the_tags_the_caller_sent(
_, row = _provider_500_failure_row(gateway, route, body, {"x-litellm-tags": tag})
assert tag in list_value(row["request_tags"]), row
def _provider_500(request: Request) -> Reply:
assert request.method == "POST"
return Reply(
status=500,
body=b'{"error":{"message":"synthetic provider outage","type":"server_error","code":"500"}}',
)
def _failure_row_prompt_tokens(gateway: Gateway, model: str, key: str, text: str) -> int:
failed: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, text), key=key)
assert failed.status_code == 500, failed.text
rows: Final = eventually(
lambda: read_rows(
'SELECT status, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(failed.headers["x-litellm-call-id"],),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["status"] == "failure", rows
prompt_tokens: Final = rows[0]["prompt_tokens"]
assert isinstance(prompt_tokens, int), rows
return prompt_tokens
def test_provider_500_failure_rows_count_deepseek_v4_prompt_tokens_with_deepseeks_tokenizer(gateway: Gateway) -> None:
with wire_server(_provider_500) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model="azure_ai/DeepSeek-V4-Flash-0731", api_base=wire.url, num_retries=0)
key: Final = scenario.key(models=[model])
chinese: Final = _failure_row_prompt_tokens(gateway, model, key, CHINESE_TEXT)
short: Final = _failure_row_prompt_tokens(gateway, model, key, SHORT_TEXT)
assert wire.drain()
assert_counted_with_deepseek(chinese - short, chinese_minus_short)

View file

@ -10,7 +10,7 @@ import sys
import threading
import time
from collections.abc import Callable, Mapping
from concurrent.futures import Future, wait
from concurrent.futures import Future, ThreadPoolExecutor, wait
from itertools import accumulate, chain
from pathlib import Path
from typing import Final
@ -586,10 +586,14 @@ def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates(
[
("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", "deepseek-v3.2"),
("qwen3-235b-a22b", "mistral-large-latest"),
],
)
def test_token_counter_counts_tools_by_the_models_chat_template_family(model: str, counted_like: str):
@ -1035,6 +1039,83 @@ 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))
@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_token_counter_counts_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 litellm.token_counter(model=model, text=_CODE_SAMPLE) == expected
@pytest.mark.parametrize(
"model",
["together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", "fireworks_ai/accounts/fireworks/models/deepseek-v2p5"],
)
def test_token_counter_skips_the_deepseek_tokenizer_for_deepseek_models_built_on_other_vocabularies(model: str):
assert litellm.token_counter(model=model, text=_CODE_SAMPLE) != _deepseek_token_count(_CODE_SAMPLE)
def test_encode_and_decode_use_deepseeks_token_ids_for_deepseek_v4():
token_ids: Final = encode(model="deepseek/deepseek-v4-flash", text=_CODE_SAMPLE)
assert token_ids == _deepseek_token_ids(_CODE_SAMPLE)
assert decode(model="deepseek/deepseek-v4-flash", tokens=token_ids) == _CODE_SAMPLE
def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
counted: Final = litellm.token_counter(model="deepseek/deepseek-v4-flash", text=_CODE_SAMPLE)
assert counted != _deepseek_token_count(_CODE_SAMPLE)
assert counted == litellm.token_counter(model="", text=_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_tokenizer_helper("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.