scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(token-counter): skip empty reasoning text for custom tokenizers

Browse source
This commit is contained in:
Mingyang Wu 2026-09-11 18:03:18 +08:00
parent af81c16019
commit 37ed3e5c7b
2 changed files with 55 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/litellm_core_utils/token_counter.py
View file

@ -839,6 +839,10 @@ def _count_anthropic_content(
return tokens
def _count_nonempty_text_tokens(text: str, count_function: TokenCounterFunction) -> int:
return count_function(text) if text else 0
def _count_content_list(
count_function: TokenCounterFunction,
content_list: str
@ -892,9 +896,12 @@ def _count_content_list(
elif c["type"] == "thinking":
# Claude extended thinking content block
# Count the thinking text and skip signature (opaque signature blob)
num_tokens += count_function(str(c.get("thinking", "")))
num_tokens += _count_nonempty_text_tokens(str(c.get("thinking", "")), count_function)
elif c["type"] == "reasoning":
num_tokens += sum(count_function(summary.get("text", "")) for summary in c.get("summary", ()))
num_tokens += sum(
_count_nonempty_text_tokens(summary.get("text", ""), count_function)
for summary in c.get("summary", ())
)
elif c["type"] == "tool_reference":
# Anthropic tool-search reference block: a lightweight pointer to
# a deferred tool, e.g. {"type": "tool_reference", "tool_name": ...}.

47
tests/test_litellm/litellm_core_utils/test_token_counter.py
View file

@ -32,7 +32,7 @@ from litellm.litellm_core_utils.token_counter import (
offload_token_count,
)
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
from litellm.types.llms.openai import ChatCompletionReasoningItem
from litellm.types.llms.openai import ChatCompletionReasoningItem, ChatCompletionThinkingBlock
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
@ -1302,6 +1302,51 @@ def test_token_counter_with_reasoning_content(
assert reasoning == original
@pytest.mark.parametrize(
("block", "texts"),
[
({"type": "thinking"}, ()),
({"type": "thinking", "thinking": ""}, ()),
({"type": "reasoning", "summary": []}, ()),
({"type": "reasoning", "summary": [{"type": "summary_text"}]}, ()),
({"type": "reasoning", "summary": [{"type": "summary_text", "text": ""}]}, ()),
(
{
"type": "reasoning",
"summary": [
{"type": "summary_text", "text": ""},
{"type": "summary_text", "text": "Thought"},
],
},
("Thought",),
),
],
ids=["missing-thinking", "empty-thinking", "empty-summary", "missing-text", "empty-text", "mixed-texts"],
)
def test_empty_reasoning_text_does_not_count_tokenizer_special_tokens(
block: ChatCompletionThinkingBlock | ChatCompletionReasoningItem, texts: tuple[str, ...]
) -> None:
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.processors import TemplateProcessing
tokenizer: Final = Tokenizer(WordLevel({"[UNK]": 0, "[BOS]": 1, "[EOS]": 2, "Ready": 3, "Thought": 4}, unk_token="[UNK]"))
tokenizer.post_processor = TemplateProcessing(single="[BOS] $A [EOS]", special_tokens=[("[BOS]", 1), ("[EOS]", 2)])
custom_tokenizer: Final = {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
messages: Final = [{"role": "assistant", "content": [block, {"type": "text", "text": "Ready"}]}]
equivalent_messages: Final = [
{
"role": "assistant",
"content": [*({"type": "text", "text": value} for value in texts), {"type": "text", "text": "Ready"}],
}
]
assert tokenizer.encode("").ids == [1, 2]
assert token_counter(custom_tokenizer=custom_tokenizer, messages=messages) == token_counter(
custom_tokenizer=custom_tokenizer, messages=equivalent_messages
)
def test_reasoning_content_preserves_prompt_cache_eligibility() -> None:
model: Final = "gemini/gemini-3.8-flash"
messages: Final = [