feat: add token counters (#12195)

Author: sjrl

haystack/token_counters/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from typing import TYPE_CHECKING
+
+from lazy_imports import LazyImporter
+
+_import_structure = {
+    "approximate_counter": ["ApproximateTokenCounter"],
+    "types": ["TokenCounter"],
+    "tiktoken_counter": ["TiktokenCounter"],
+}
+
+if TYPE_CHECKING:
+    from .approximate_counter import ApproximateTokenCounter as ApproximateTokenCounter
+    from .tiktoken_counter import TiktokenCounter as TiktokenCounter
+    from .types import TokenCounter as TokenCounter
+else:
+    sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)

haystack/token_counters/approximate_counter.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from haystack.core.serialization import default_to_dict
+from haystack.dataclasses import ChatMessage
+from haystack.token_counters.types import TokenCounter
+from haystack.token_counters.utils import _non_text_tokens, _rendered_conversation, _rendered_tools
+from haystack.tools import ToolsType
+
+
+class ApproximateTokenCounter(TokenCounter):
+    """
+    Estimates tokens from text length using a flat ratio of characters to tokens.
+
+    ## Usage Example:
+    ```python
+    from haystack.dataclasses import ChatMessage
+    from haystack.token_counters import ApproximateTokenCounter
+
+    counter = ApproximateTokenCounter(chars_per_token=4.0)
+    messages = [
+        ChatMessage.from_user("Hello, how are you?"),
+        ChatMessage.from_assistant("I'm good, thank you! How can I assist you today?")
+    ]
+    token_count = counter.count(messages)
+    print(f"Estimated token count: {token_count}")
+    ```
+    """
+
+    def __init__(self, chars_per_token: float = 4.0, tokens_per_image: int = 85, tokens_per_file: int = 1000) -> None:
+        """
+        Initialize the counter.
+
+        :param chars_per_token: How many characters to treat as one token.
+        :param tokens_per_image: Tokens to charge per image, which has no text to 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
+            cost depends on the page count; raise it if you send long ones.
+        :raises ValueError: If `chars_per_token` is not positive.
+        """
+        if chars_per_token <= 0:
+            raise ValueError(f"`chars_per_token` must be greater than 0, got {chars_per_token}.")
+        self.chars_per_token = chars_per_token
+        self.tokens_per_image = tokens_per_image
+        self.tokens_per_file = tokens_per_file
+
+    def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int:
+        """
+        Return the estimated number of tokens the given messages occupy.
+
+        :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.
+        """
+        if not messages and not tools:
+            return 0
+        text = _rendered_conversation(messages) + _rendered_tools(tools)
+        text_tokens = int(len(text) / self.chars_per_token)
+        return text_tokens + _non_text_tokens(
+            messages=messages, tokens_per_image=self.tokens_per_image, tokens_per_file=self.tokens_per_file
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize the counter.
+
+        :returns: A dictionary representation of the counter.
+        """
+        return default_to_dict(
+            self,
+            chars_per_token=self.chars_per_token,
+            tokens_per_image=self.tokens_per_image,
+            tokens_per_file=self.tokens_per_file,
+        )

haystack/token_counters/tiktoken_counter.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from haystack.core.serialization import default_to_dict
+from haystack.dataclasses import ChatMessage
+from haystack.lazy_imports import LazyImport
+from haystack.token_counters.types import TokenCounter
+from haystack.token_counters.utils import _non_text_tokens, _rendered_conversation, _rendered_tools
+from haystack.tools import ToolsType
+
+with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports:
+    import tiktoken
+
+
+class TiktokenCounter(TokenCounter):
+    """
+    Counts tokens locally with `tiktoken`, OpenAI's byte-pair encoder.
+
+    Counting is an estimate, and two limits are worth knowing before relying on it:
+    - **It is text-only**, so images and files get the flat `tokens_per_image` / `tokens_per_file` estimate rather
+      than a real count.
+    - **It is OpenAI's encoder.** Other providers tokenize differently, so expect the count to drift on them.
+
+    ## Usage Example:
+    ```python
+    from haystack.dataclasses import ChatMessage
+    from haystack.token_counters import TiktokenCounter
+
+    counter = TiktokenCounter(encoding="o200k_base")
+    messages = [
+        ChatMessage.from_user("Hello, how are you?"),
+        ChatMessage.from_assistant("I'm good, thank you! How can I assist you today?")
+    ]
+    token_count = counter.count(messages)
+    print(f"Token count: {token_count}")
+    ```
+    """
+
+    def __init__(self, encoding: str = "o200k_base", tokens_per_image: int = 85, tokens_per_file: int = 1000) -> None:
+        """
+        Initialize the counter.
+
+        :param encoding: The `tiktoken` encoding to count with. The default, `o200k_base`, is what current OpenAI
+            models use.
+        :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
+            cost depends on the page count; raise it if you send long ones.
+        :raises ImportError: If `tiktoken` is not installed.
+        """
+        tiktoken_imports.check()
+        self.encoding = encoding
+        self.tokens_per_image = tokens_per_image
+        self.tokens_per_file = tokens_per_file
+        self._encoder: Any = None
+
+    def warm_up(self) -> None:
+        """Load the encoder, downloading its vocabulary if it is not already cached."""
+        if self._encoder is not None:
+            return
+        self._encoder = tiktoken.get_encoding(self.encoding)
+
+    def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int:
+        """
+        Return the estimated number of tokens used by the given messages.
+
+        :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.
+        """
+        if not messages and not tools:
+            return 0
+        self.warm_up()
+        text_tokens = len(self._encoder.encode(_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
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize the counter.
+
+        :returns: A dictionary representation of the counter.
+        """
+        return default_to_dict(
+            self, encoding=self.encoding, tokens_per_image=self.tokens_per_image, tokens_per_file=self.tokens_per_file
+        )

haystack/token_counters/types/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .protocol import TokenCounter
+
+__all__ = ["TokenCounter"]

haystack/token_counters/types/protocol.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Protocol
+
+from haystack.core.serialization import default_from_dict
+from haystack.dataclasses import ChatMessage
+from haystack.tools import ToolsType
+
+
+class TokenCounter(Protocol):
+    """
+    Estimates the number tokens used by a list of messages.
+
+    Implement `to_dict` so the counter's settings survive serialization. The default `from_dict` passes them straight
+    back to the constructor, which is enough for plain values; override it when `to_dict` emitted something that has to
+    be rebuilt first, such as a `Secret` or a nested component.
+    """
+
+    def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int:
+        """
+        Return the estimated number of tokens in the given messages.
+
+        :param messages: The messages to measure.
+        :param tools: Tools whose schemas are sent alongside the messages, and so consume tokens too. Pass them to have
+            them counted; leave as None to measure the messages alone.
+        :returns: The estimated token count.
+        """
+        ...
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize the counter to a dictionary."""
+        ...
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "TokenCounter":
+        """Deserialize the counter from a dictionary."""
+        return default_from_dict(cls, data)

haystack/token_counters/utils.py

@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent
+from haystack.dataclasses.chat_message import ChatMessageContentT, ToolCallResultContentT
+from haystack.tools import ToolsType, flatten_tools_or_toolsets
+
+
+def _non_text_placeholder(content: ChatMessageContentT) -> str:
+    """A short stand-in, such as `<image>`, for message content that has no text form."""
+    if isinstance(content, ImageContent):
+        return "<image>"
+    if isinstance(content, FileContent):
+        return f"<file: {content.filename or 'unnamed'}>"
+    return f"<{type(content).__name__}>"
+
+
+def _tool_result_text(result: ToolCallResultContentT) -> str:
+    """A tool result as a single string, with placeholders standing in for any non-text parts."""
+    if isinstance(result, str):
+        return result
+    return "".join(block.text if isinstance(block, TextContent) else _non_text_placeholder(block) for block in result)
+
+
+def _render_message(message: ChatMessage) -> str:
+    """
+    One message as one or more lines of plain text.
+
+    Reasoning content is deliberately left out: providers discard it between turns, so it is not part of the context
+    being measured.
+    """
+    role = message.role.value
+    # A tool-result msg only carries tool_call_results, so it is rendered on its own and labelled with the tool that
+    # produced it.
+    if results := message.tool_call_results:
+        return "\n".join(
+            f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] {_tool_result_text(result.result)}"
+            for result in results
+        )
+
+    lines: list[str] = []
+    if texts := message.texts:
+        lines.append(f"[{role}] " + "\n".join(texts))
+    for call in message.tool_calls:
+        arguments = json.dumps(call.arguments, default=str, sort_keys=True)
+        lines.append(f"[{role} -> tool_call] {call.tool_name}({arguments})")
+    # Images and files cost tokens too, so they need a stand-in rather than being skipped.
+    non_text: list[ChatMessageContentT] = [*message.images, *message.files]
+    for content in non_text:
+        lines.append(f"[{role}] {_non_text_placeholder(content)}")
+    return "\n".join(lines) if lines else f"[{role}] <no content>"
+
+
+def _rendered_conversation(messages: list[ChatMessage]) -> str:
+    """The whole conversation as one plain-text block, which is what a counter measures."""
+    return "\n".join(_render_message(message) for message in messages)
+
+
+def _rendered_tools(tools: ToolsType | None) -> str:
+    """
+    Tool schemas as one JSON block, standing in for what a provider sends alongside the messages.
+
+    :param tools: The tools to render, or None.
+    :returns: The rendered schemas, or an empty string when there are none.
+    """
+    if not tools:
+        return ""
+    return json.dumps([tool.tool_spec for tool in flatten_tools_or_toolsets(tools)], default=str, sort_keys=True)
+
+
+def _non_text_tokens(messages: list[ChatMessage], *, tokens_per_image: int, tokens_per_file: int) -> int:
+    """
+    A flat estimate for the images and files in a conversation, which have no text to measure.
+
+    :param messages: The conversation to measure.
+    :param tokens_per_image: Tokens to charge per image.
+    :param tokens_per_file: Tokens to charge per file.
+    :returns: The estimated token count for the non-text content.
+    """
+    images = 0
+    files = 0
+    for message in messages:
+        images += len(message.images)
+        files += len(message.files)
+        # Images and files returned by a tool are nested inside the tool result
+        for result in message.tool_call_results:
+            if isinstance(result.result, str):
+                continue
+            images += sum(isinstance(block, ImageContent) for block in result.result)
+            files += sum(isinstance(block, FileContent) for block in result.result)
+    return images * tokens_per_image + files * tokens_per_file

pydoc/token_counters_api.yml

@@ -0,0 +1,13 @@
+loaders:
+  - search_path: [../haystack/token_counters]
+    modules: ["types/protocol", "approximate_counter", "tiktoken_counter"]
+processors:
+  - type: filter
+    documented_only: true
+    skip_empty_modules: true
+renderer:
+  title: Token Counters
+  id: token-counters-api
+  description: Estimate how many tokens a conversation occupies, for features that need a size before sending it to a
+    model.
+  filename: token_counters_api.md

pyproject.toml

@@ -113,7 +113,7 @@ dependencies = [
   "python-oxmsg",                     # MSGToDocument
 
   "nltk>=3.9.1", # NLTKDocumentSplitter, RecursiveDocumentSplitter
-  "tiktoken", # RecursiveDocumentSplitter
+  "tiktoken", # RecursiveDocumentSplitter, TiktokenCounter
 
   # Human in the Loop
   "rich",                 # Console rendering for HITL

releasenotes/notes/add-token-counters-f2685e989bff106e.yaml

@@ -0,0 +1,42 @@
+---
+features:
+  - |
+    Added ``haystack.token_counters``: a ``TokenCounter`` protocol for estimating how many tokens a list of
+    ``ChatMessage`` objects occupies, with two implementations.
+
+    Providers report token usage only after a call, and only for the call as a whole, so anything that needs a size
+    beforehand - deciding whether a conversation still fits a model's context window, or how much of it to drop - has to
+    estimate one.
+
+    .. code-block:: python
+
+        from haystack.dataclasses import ChatMessage
+        from haystack.token_counters import ApproximateTokenCounter, TiktokenCounter
+
+        messages = [ChatMessage.from_user("Hello, how are you?")]
+
+        # No dependencies: estimates from text length.
+        ApproximateTokenCounter(chars_per_token=4.0).count(messages)
+
+        # Closer for OpenAI models; needs: pip install tiktoken
+        TiktokenCounter(encoding="o200k_base").count(messages)
+
+    ``ApproximateTokenCounter`` needs nothing installed and estimates from text length at a configurable
+    ``chars_per_token``. ``TiktokenCounter`` counts with OpenAI's byte-pair encoder, which is closer for OpenAI models
+    but requires ``tiktoken`` and drifts on other providers; it raises at construction when the dependency is missing,
+    and loads its encoding on first use.
+
+    Neither can measure an image or a file, since a tokenizer only sees text and providers derive an image's cost from
+    its dimensions. Both charge a flat ``tokens_per_image`` and ``tokens_per_file`` instead, counting images a tool
+    returned inside its result as well as those a message carries directly. Raise those values if you send large images
+    or long documents.
+
+    Tool schemas are sent alongside the messages and consume tokens too, so ``count`` takes an optional ``tools``
+    argument to have them included:
+
+    .. code-block:: python
+
+        counter.count(messages, tools=[my_tool])
+
+    Implement ``TokenCounter`` to count differently - for instance against a provider's own token-counting endpoint,
+    which is the only way to have images counted exactly.