diff --git a/haystack/token_counters/tiktoken_counter.py b/haystack/token_counters/tiktoken_counter.py index 57298234074..cb54062be94 100644 --- a/haystack/token_counters/tiktoken_counter.py +++ b/haystack/token_counters/tiktoken_counter.py @@ -74,7 +74,9 @@ 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))) + # Encode as ordinary text so literal special-token strings (e.g. `<|endoftext|>`) in + # messages or tool schemas are counted instead of raising `ValueError`. + 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-counter-special-tokens-12869.yaml b/releasenotes/notes/fix-tiktoken-counter-special-tokens-12869.yaml new file mode 100644 index 00000000000..cc3f2918314 --- /dev/null +++ b/releasenotes/notes/fix-tiktoken-counter-special-tokens-12869.yaml @@ -0,0 +1,6 @@ +fixes: + - | + Fix ``TiktokenCounter`` failing on messages or tool schemas containing + strings such as ``<|endoftext|>``. + These strings are now counted as ordinary text instead of raising + ``ValueError``. diff --git a/test/token_counters/test_tiktoken_counter.py b/test/token_counters/test_tiktoken_counter.py index bd19f1e5871..f6c407e19cc 100644 --- a/test/token_counters/test_tiktoken_counter.py +++ b/test/token_counters/test_tiktoken_counter.py @@ -31,6 +31,11 @@ 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]: + # Mirrors `tiktoken.Encoding.encode_ordinary`: counts special-token strings as ordinary text. + self.encoded.append(text) + return list(range(len(text.split()))) + @pytest.fixture def fake_encoder(monkeypatch: pytest.MonkeyPatch) -> _FakeEncoder: @@ -97,6 +102,13 @@ def test_encodes_the_rendered_conversation(self, fake_encoder): assert "" in rendered assert "" in rendered + def test_special_token_strings_are_counted_as_ordinary_text(self, fake_encoder): + # Literal markers such as `<|endoftext|>` must not raise; they are counted as ordinary text. + messages = [ChatMessage.from_user("The manual documents <|endoftext|> as a literal marker.")] + + assert TiktokenCounter().count(messages) > 0 + assert "<|endoftext|>" in fake_encoder.encoded[0] + def test_serde_round_trip(self): data = TiktokenCounter(encoding="cl100k_base", tokens_per_image=200, tokens_per_file=3000).to_dict() @@ -160,3 +172,24 @@ def test_an_image_is_charged_at_the_flat_rate(self): counter = TiktokenCounter(tokens_per_image=85) assert counter.count([ChatMessage.from_user(content_parts=[IMAGE])]) > 85 + + def test_special_token_string_in_message_is_counted(self): + # Regression test for https://github.com/deepset-ai/haystack/issues/12869 + counter = TiktokenCounter() + + count = counter.count([ChatMessage.from_user("The manual documents <|endoftext|> as a literal marker.")]) + + assert count > 0 + + def test_special_token_string_in_tool_description_is_counted(self): + # The tools path concatenates schemas into the same encode call, so it must not raise either. + @tool + def marked_search(query: Annotated[str, "the search query"]) -> str: + """Search the web. Documents <|endoftext|> as a literal marker.""" + return "result" + + counter = TiktokenCounter() + + count = counter.count([ChatMessage.from_user("hi")], tools=[marked_search]) + + assert count > 0