From d26c1b1a9ff5d6971a793af335cefef15c48fd50 Mon Sep 17 00:00:00 2001 From: ducanhnguyen223 Date: Sun, 27 Sep 2026 20:02:42 +0700 Subject: [PATCH] fix(token-counter): treat literal special tokens as text --- haystack/token_counters/tiktoken_counter.py | 2 +- .../fix-tiktoken-literal-special-tokens.yaml | 5 +++++ test/token_counters/test_tiktoken_counter.py | 21 ++++++++++++++++++- 3 files changed, 26 insertions(+), 2 deletions(-) create mode 100644 releasenotes/notes/fix-tiktoken-literal-special-tokens.yaml diff --git a/haystack/token_counters/tiktoken_counter.py b/haystack/token_counters/tiktoken_counter.py index 57298234074..f505b89bf99 100644 --- a/haystack/token_counters/tiktoken_counter.py +++ b/haystack/token_counters/tiktoken_counter.py @@ -74,7 +74,7 @@ def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> if not messages and not tools: return 0 self.warm_up() - text_tokens = len(self._encoder.encode(_rendered_conversation(messages) + _rendered_tools(tools))) + text_tokens = len(self._encoder.encode_ordinary(_rendered_conversation(messages) + _rendered_tools(tools))) return text_tokens + _non_text_tokens( messages, tokens_per_image=self.tokens_per_image, tokens_per_file=self.tokens_per_file ) diff --git a/releasenotes/notes/fix-tiktoken-literal-special-tokens.yaml b/releasenotes/notes/fix-tiktoken-literal-special-tokens.yaml new file mode 100644 index 00000000000..ea223ec56dd --- /dev/null +++ b/releasenotes/notes/fix-tiktoken-literal-special-tokens.yaml @@ -0,0 +1,5 @@ +fixes: + - | + Fix ``TiktokenCounter`` raising ``ValueError`` when user messages or tool schemas contain a literal tiktoken + special-token marker such as ``<|endoftext|>``. These markers are now counted as ordinary text, matching the + behavior of the document splitters. diff --git a/test/token_counters/test_tiktoken_counter.py b/test/token_counters/test_tiktoken_counter.py index bd19f1e5871..8370d5e1032 100644 --- a/test/token_counters/test_tiktoken_counter.py +++ b/test/token_counters/test_tiktoken_counter.py @@ -26,11 +26,16 @@ class _FakeEncoder: def __init__(self) -> None: self.encoded: list[str] = [] + self.ordinary_encoded: list[str] = [] def encode(self, text: str) -> list[int]: self.encoded.append(text) return list(range(len(text.split()))) + def encode_ordinary(self, text: str) -> list[int]: + self.ordinary_encoded.append(text) + return list(range(len(text.split()))) + @pytest.fixture def fake_encoder(monkeypatch: pytest.MonkeyPatch) -> _FakeEncoder: @@ -90,7 +95,7 @@ def test_encodes_the_rendered_conversation(self, fake_encoder): TiktokenCounter().count(messages) - rendered = fake_encoder.encoded[0] + rendered = fake_encoder.ordinary_encoded[0] assert "[assistant] looking" in rendered assert '[assistant -> tool_call] search({"q": "x"})' in rendered assert "[tool:search] found it" in rendered @@ -116,6 +121,12 @@ def search(query: Annotated[str, "the search query"]) -> str: return "result" +@tool +def search_with_literal_marker(query: Annotated[str, "the search query"]) -> str: + """Search documentation that may contain the literal marker <|endoftext|>.""" + return "result" + + class TestTiktokenCounterTools: def test_tool_schemas_add_to_the_count(self, fake_encoder): # A provider is sent the schemas alongside the messages, so they consume tokens too. @@ -127,6 +138,14 @@ def test_tool_schemas_add_to_the_count(self, fake_encoder): def test_tools_can_be_counted_without_messages(self, fake_encoder): assert TiktokenCounter().count([], tools=[search]) > 0 + def test_literal_special_token_in_message_is_counted_as_text(self, fake_encoder): + assert TiktokenCounter().count([ChatMessage.from_user("literal <|endoftext|> marker")]) > 0 + assert "<|endoftext|>" in fake_encoder.ordinary_encoded[0] + + def test_literal_special_token_in_tool_schema_is_counted_as_text(self, fake_encoder): + assert TiktokenCounter().count([], tools=[search_with_literal_marker]) > 0 + assert "<|endoftext|>" in fake_encoder.ordinary_encoded[0] + def test_nothing_to_measure_is_zero(self): assert TiktokenCounter().count([]) == 0 assert TiktokenCounter().count([], tools=None) == 0