diff --git a/haystack/token_counters/tiktoken_counter.py b/haystack/token_counters/tiktoken_counter.py index 57298234074..80f9c93fe52 100644 --- a/haystack/token_counters/tiktoken_counter.py +++ b/haystack/token_counters/tiktoken_counter.py @@ -44,7 +44,8 @@ def __init__(self, encoding: str = "o200k_base", tokens_per_image: int = 85, tok Initialize the counter. :param encoding: The `tiktoken` encoding to count with. The default, `o200k_base`, is what current OpenAI - models use. + models use. Special-token strings in message or tool text, such as `<|endoftext|>`, are counted as + ordinary text rather than raising. :param tokens_per_image: Tokens to charge per image, which the tokenizer cannot measure. The default is what OpenAI charges for a small image; raise it if you send large ones. :param tokens_per_file: Tokens to charge per file. A rough stand-in for a short document, since the real @@ -74,7 +75,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-counter-special-tokens-3a8f27d1e6b4c05a.yaml b/releasenotes/notes/fix-tiktoken-counter-special-tokens-3a8f27d1e6b4c05a.yaml new file mode 100644 index 00000000000..b6fefe0dc22 --- /dev/null +++ b/releasenotes/notes/fix-tiktoken-counter-special-tokens-3a8f27d1e6b4c05a.yaml @@ -0,0 +1,7 @@ +--- +fixes: + - | + Fix ``TiktokenCounter.count`` raising ``ValueError`` when message or tool text + contained literal special-token strings such as ``<|endoftext|>``. + These strings are now counted as ordinary text, consistent with the behavior + already fixed for the token-based document splitters. diff --git a/test/token_counters/test_tiktoken_counter.py b/test/token_counters/test_tiktoken_counter.py index bd19f1e5871..2785f039422 100644 --- a/test/token_counters/test_tiktoken_counter.py +++ b/test/token_counters/test_tiktoken_counter.py @@ -27,10 +27,13 @@ class _FakeEncoder: def __init__(self) -> None: self.encoded: list[str] = [] - def encode(self, text: str) -> list[int]: + def encode_ordinary(self, text: str) -> list[int]: self.encoded.append(text) return list(range(len(text.split()))) + def encode(self, text: str, **kwargs) -> list[int]: # pragma: no cover + raise AssertionError("encode() with default disallowed_special='all' would raise on special-token text") + @pytest.fixture def fake_encoder(monkeypatch: pytest.MonkeyPatch) -> _FakeEncoder: @@ -97,6 +100,16 @@ def test_encodes_the_rendered_conversation(self, fake_encoder): assert "" in rendered assert "" in rendered + def test_counts_literal_special_token_text_as_ordinary(self, fake_encoder): + # Regression for #12869: count() must route text through encode_ordinary so a literal + # special-token marker cannot raise the way encode()'s disallowed_special="all" default would. + marker = "<" + "|endoftext" + "|>" + counter = TiktokenCounter() + + counter.count([ChatMessage.from_user(f"text with {marker} inside")]) + + assert any(marker in text for text in fake_encoder.encoded) + def test_serde_round_trip(self): data = TiktokenCounter(encoding="cl100k_base", tokens_per_image=200, tokens_per_file=3000).to_dict() @@ -160,3 +173,15 @@ 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 + + @pytest.mark.parametrize("encoding", ["o200k_base", "cl100k_base"]) + def test_literal_special_token_text_is_counted_as_ordinary(self, encoding): + # Regression for #12869: with the real encoder, count() must not raise on text that merely + # contains a special-token marker; the marker contributes ordinary tokens. + marker = "<" + "|endoftext" + "|>" + counter = TiktokenCounter(encoding=encoding) + + with_marker = counter.count([ChatMessage.from_user(f"The manual documents {marker} as a literal.")]) + without = counter.count([ChatMessage.from_user("The manual documents x as a literal.")]) + + assert without < with_marker