mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
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:
parent
d56ff281de
commit
2d181e6633
9 changed files with 491 additions and 51 deletions
|
|
@ -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
|
|
@ -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.
|
||||
|
|
@ -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":
|
||||
|
|
|
|||
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}"
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue