PR comments

Author: sjrl

haystack/token_counters/approximate_counter.py

@@ -7,7 +7,8 @@
 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
+from haystack.token_counters.utils import _non_text_tokens, _rendered_conversation, _rendered_tools
+from haystack.tools import ToolsType
 
 
 class ApproximateTokenCounter(TokenCounter):
@@ -46,17 +47,21 @@ def __init__(self, chars_per_token: float = 4.0, tokens_per_image: int = 85, tok
         self.tokens_per_image = tokens_per_image
         self.tokens_per_file = tokens_per_file
 
-    def count(self, messages: list[ChatMessage]) -> int:
+    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.
-        :returns: The estimated token count, or `0` for an empty list.
+        :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:
+        if not messages and not tools:
             return 0
-        text_tokens = int(len(_rendered_conversation(messages)) / self.chars_per_token)
-        return text_tokens + _non_text_tokens(messages, self.tokens_per_image, self.tokens_per_file)
+        text = _rendered_conversation(messages) + _rendered_tools(tools)
+        text_tokens = int(len(text) / self.chars_per_token)
+        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]:
         """

haystack/token_counters/tiktoken_counter.py

@@ -8,7 +8,8 @@
 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
+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
@@ -62,18 +63,21 @@ def warm_up(self) -> None:
             return
         self._encoder = tiktoken.get_encoding(self.encoding)
 
-    def count(self, messages: list[ChatMessage]) -> int:
+    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.
-        :returns: The estimated token count, or `0` for an empty list.
+        :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:
+        if not messages and not tools:
             return 0
         self.warm_up()
-        text_tokens = len(self._encoder.encode(_rendered_conversation(messages)))
-        return text_tokens + _non_text_tokens(messages, self.tokens_per_image, self.tokens_per_file)
+        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]:
         """

haystack/token_counters/types/protocol.py

@@ -4,30 +4,34 @@
 
 from typing import Any, Protocol
 
-from haystack.core.serialization import default_from_dict, default_to_dict
+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.
 
-    Override `to_dict` when the counter takes constructor arguments, so that they are serialized; the default emits
-    none. `from_dict` rebuilds the counter from whatever `to_dict` emitted and rarely needs overriding.
+    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]) -> int:
+    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."""
-        return default_to_dict(self)
+        ...
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "TokenCounter":

haystack/token_counters/utils.py

@@ -6,6 +6,7 @@
 
 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:
@@ -58,7 +59,19 @@ def _rendered_conversation(messages: list[ChatMessage]) -> str:
     return "\n".join(_render_message(message) for message in messages)
 
 
-def _non_text_tokens(messages: list[ChatMessage], tokens_per_image: int, tokens_per_file: int) -> int:
+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.
 

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

@@ -31,5 +31,12 @@ features:
     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.

test/token_counters/test_approximate_counter.py

@@ -2,11 +2,14 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Annotated
+
 import pytest
 
 from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCall
 from haystack.token_counters import ApproximateTokenCounter
 from haystack.token_counters.utils import _rendered_conversation
+from haystack.tools import tool
 
 IMAGE = ImageContent(base64_image="Zm9v", mime_type="image/png")
 FILE = FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="report.pdf")
@@ -79,3 +82,25 @@ def test_serde_round_trip(self):
         restored = ApproximateTokenCounter.from_dict(data)
         assert restored.chars_per_token == 3.5
         assert restored.tokens_per_file == 3000
+
+
+@tool
+def search(query: Annotated[str, "the search query"]) -> str:
+    """Search the web for a query and return the top results."""
+    return "result"
+
+
+class TestApproximateTokenCounterTools:
+    def test_tool_schemas_add_to_the_count(self):
+        # A provider is sent the schemas alongside the messages, so they consume tokens too.
+        counter = ApproximateTokenCounter()
+        messages = [ChatMessage.from_user("hi")]
+
+        assert counter.count(messages, tools=[search]) > counter.count(messages)
+
+    def test_tools_can_be_counted_without_messages(self):
+        assert ApproximateTokenCounter().count([], tools=[search]) > 0
+
+    def test_nothing_to_measure_is_zero(self):
+        assert ApproximateTokenCounter().count([]) == 0
+        assert ApproximateTokenCounter().count([], tools=None) == 0

test/token_counters/test_tiktoken_counter.py

@@ -2,11 +2,14 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Annotated
+
 import pytest
 
 from haystack.dataclasses import ChatMessage, FileContent, ImageContent, ToolCall
 from haystack.token_counters import TiktokenCounter
 from haystack.token_counters import tiktoken_counter as tiktoken_counter_module
+from haystack.tools import tool
 
 IMAGE = ImageContent(base64_image="Zm9v", mime_type="image/png")
 FILE = FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="report.pdf")
@@ -134,3 +137,25 @@ 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
+
+
+@tool
+def search(query: Annotated[str, "the search query"]) -> str:
+    """Search the web for a query and return the top results."""
+    return "result"
+
+
+class TestTiktokenCounterTools:
+    def test_tool_schemas_add_to_the_count(self):
+        # A provider is sent the schemas alongside the messages, so they consume tokens too.
+        counter = TiktokenCounter()
+        messages = [ChatMessage.from_user("hi")]
+
+        assert counter.count(messages, tools=[search]) > counter.count(messages)
+
+    def test_tools_can_be_counted_without_messages(self):
+        assert TiktokenCounter().count([], tools=[search]) > 0
+
+    def test_nothing_to_measure_is_zero(self):
+        assert TiktokenCounter().count([]) == 0
+        assert TiktokenCounter().count([], tools=None) == 0

test/token_counters/test_utils.py

@@ -2,10 +2,19 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Annotated
+
 import pytest
 
-from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCall
-from haystack.token_counters.utils import _render_message, _rendered_conversation, _tool_result_text
+from haystack.dataclasses import ChatMessage, FileContent, ImageContent, ReasoningContent, TextContent, ToolCall
+from haystack.token_counters.utils import (
+    _non_text_placeholder,
+    _render_message,
+    _rendered_conversation,
+    _rendered_tools,
+    _tool_result_text,
+)
+from haystack.tools import tool
 
 IMAGE = ImageContent(base64_image="Zm9v", mime_type="image/png")
 FILE = FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="report.pdf")
@@ -65,3 +74,39 @@ def test_joins_messages_with_newlines(self):
 
     def test_empty_conversation(self):
         assert _rendered_conversation([]) == ""
+
+
+@tool
+def search(query: Annotated[str, "the search query"]) -> str:
+    """Search the web for a query."""
+    return "result"
+
+
+class TestNonTextPlaceholder:
+    @pytest.mark.parametrize(
+        ("content", "expected"),
+        [
+            pytest.param(IMAGE, "<image>", id="image"),
+            pytest.param(FILE, "<file: report.pdf>", id="file-with-name"),
+            pytest.param(
+                FileContent(base64_data="Zm9v", mime_type="application/pdf"), "<file: unnamed>", id="file-unnamed"
+            ),
+            # Anything else falls back to naming its type, so an unexpected block still shows up in the text.
+            pytest.param(ReasoningContent(reasoning_text="thinking"), "<ReasoningContent>", id="unknown-type"),
+        ],
+    )
+    def test_placeholder(self, content, expected):
+        assert _non_text_placeholder(content) == expected
+
+
+class TestRenderedTools:
+    def test_no_tools_renders_nothing(self):
+        assert _rendered_tools(None) == ""
+        assert _rendered_tools([]) == ""
+
+    def test_renders_the_schema_a_provider_would_be_sent(self):
+        rendered = _rendered_tools([search])
+
+        assert "search" in rendered
+        assert "Search the web for a query." in rendered
+        assert "the search query" in rendered