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/tiktoken-literal-special-tokens-3a5e7d2b9c4f1a06.yaml b/releasenotes/notes/tiktoken-literal-special-tokens-3a5e7d2b9c4f1a06.yaml new file mode 100644 index 00000000000..a31ea4de375 --- /dev/null +++ b/releasenotes/notes/tiktoken-literal-special-tokens-3a5e7d2b9c4f1a06.yaml @@ -0,0 +1,5 @@ +--- +fixes: + - | + ``TiktokenCounter`` now counts literal strings such as ``<|endoftext|>`` in messages and tool descriptions + as ordinary text instead of raising a ``ValueError``. diff --git a/test/token_counters/test_tiktoken_counter.py b/test/token_counters/test_tiktoken_counter.py index bd19f1e5871..e4b279614e1 100644 --- a/test/token_counters/test_tiktoken_counter.py +++ b/test/token_counters/test_tiktoken_counter.py @@ -27,7 +27,7 @@ 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()))) @@ -97,6 +97,14 @@ def test_encodes_the_rendered_conversation(self, fake_encoder): assert "" in rendered assert "" in rendered + def test_counts_literal_special_token_in_message(self, fake_encoder): + counter = TiktokenCounter() + + count = counter.count([ChatMessage.from_user("Literal <|endoftext|> marker")]) + + assert count > 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() @@ -131,11 +139,34 @@ def test_nothing_to_measure_is_zero(self): assert TiktokenCounter().count([]) == 0 assert TiktokenCounter().count([], tools=None) == 0 + def test_counts_literal_special_token_in_tool_description(self, fake_encoder): + @tool + def describe_marker() -> str: + """Describe the literal <|endoftext|> marker.""" + return "marker" + + count = TiktokenCounter().count([], tools=[describe_marker]) + + assert count > 0 + assert "<|endoftext|>" in fake_encoder.encoded[0] + @pytest.mark.integration class TestTiktokenCounterIntegration: """Exercises the real encoder, which downloads its vocabulary on first use.""" + def test_literal_special_token_in_messages_and_tools(self): + @tool + def describe_marker() -> str: + """Explain the literal <|endoftext|> marker.""" + return "marker" + + counter = TiktokenCounter() + messages = [ChatMessage.from_user("Literal <|endoftext|> marker")] + + assert counter.count(messages) > 0 + assert counter.count(messages, tools=[describe_marker]) > counter.count(messages) + def test_counts_grow_with_content(self): counter = TiktokenCounter()