fix: TiktokenCounter counts literal special-token strings as ordinary text

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
         )

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.

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 "<image>" in rendered
         assert "<file: report.pdf>" 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