diff --git a/haystack/token_counters/tiktoken_counter.py b/haystack/token_counters/tiktoken_counter.py index 57298234074..e76d00df892 100644 --- a/haystack/token_counters/tiktoken_counter.py +++ b/haystack/token_counters/tiktoken_counter.py @@ -69,12 +69,16 @@ def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> :param messages: The messages to measure. :param tools: Tools whose schemas are sent alongside the messages, and so consume tokens too. - :returns: The estimated token count, or `0` when there is nothing to measure. + :returns: The estimated token count, or `0` when there is nothing to measure. Special-token strings in the + messages are counted as ordinary text. """ if not messages and not tools: return 0 self.warm_up() - text_tokens = len(self._encoder.encode(_rendered_conversation(messages) + _rendered_tools(tools))) + # `encode_ordinary` treats special-token strings such as `<|endoftext|>` as ordinary text. Plain `encode` + # raises on them, which would break counting for conversations that merely mention them. + rendered = _rendered_conversation(messages) + _rendered_tools(tools) + text_tokens = len(self._encoder.encode_ordinary(rendered)) 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-special-tokens-in-tiktoken-counter-7b1c2d4e5f6a8b90.yaml b/releasenotes/notes/fix-special-tokens-in-tiktoken-counter-7b1c2d4e5f6a8b90.yaml new file mode 100644 index 00000000000..74d6cb8ff60 --- /dev/null +++ b/releasenotes/notes/fix-special-tokens-in-tiktoken-counter-7b1c2d4e5f6a8b90.yaml @@ -0,0 +1,5 @@ +--- +fixes: + - | + `TiktokenCounter` no longer raises a `ValueError` when the messages or tool schemas being counted contain + literal special-token strings such as `<|endoftext|>`. They are now counted as ordinary text. diff --git a/test/token_counters/test_tiktoken_counter.py b/test/token_counters/test_tiktoken_counter.py index bd19f1e5871..2701ca9886c 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()))) @@ -160,3 +160,13 @@ 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.integration +def test_special_token_strings_are_counted_as_literal_text(): + # `tiktoken.Encoding.encode` rejects these by default, so a conversation that merely mentions one + # must not make counting fail. + counter = TiktokenCounter() + messages = [ChatMessage.from_user("What does <|endoftext|> mean?")] + + assert counter.count(messages) > 0