diff --git a/headroom/providers/anthropic.py b/headroom/providers/anthropic.py index 107fc53c5..0bca5be84 100644 --- a/headroom/providers/anthropic.py +++ b/headroom/providers/anthropic.py @@ -24,6 +24,7 @@ import warnings from typing import Any, cast from headroom import paths as _paths +from headroom.tokenizers.base import count_content_blocks from .base import Provider, TokenCounter @@ -395,15 +396,16 @@ class AnthropicTokenCounter(TokenCounter): if isinstance(content, str): tokens += self.count_text(content) elif isinstance(content, list): - for block in content: - if isinstance(block, dict): - if block.get("type") == "text": - tokens += self.count_text(block.get("text", "")) - elif block.get("type") == "tool_use": - tokens += self.count_text(block.get("name", "")) - tokens += self.count_text(str(block.get("input", {}))) - elif block.get("type") == "tool_result": - tokens += self.count_text(str(block.get("content", ""))) + # Delegate to the audited shared walker instead of a partial + # per-provider one. Each provider counter had grown its own + # shortened branch list, so every modern block priced at ~0: + # measured on a 6,800-char block this returned 8 tokens for + # tool_result, thinking, document, mcp_tool_result — and for + # output_text / refusal, which are OpenAI's OWN Responses shapes. + # The shared walker is also image-safe: a 200KB base64 image gets + # a pixel-based 1600, not the ~50K phantom text tokens a naive + # str(block) catch-all would produce. + tokens += count_content_blocks(content, self.count_text) # OpenAI format tool calls if "tool_calls" in message: diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index cc91d2b71..73a953c21 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -16,6 +16,7 @@ from functools import lru_cache from typing import Any, cast from headroom import paths as _paths +from headroom.tokenizers.base import count_content_blocks from .base import Provider, TokenCounter @@ -331,14 +332,16 @@ class OpenAITokenCounter: if isinstance(content, str): tokens += self.count_text(content) elif isinstance(content, list): - for part in content: - if isinstance(part, dict): - if part.get("type") == "text": - tokens += self.count_text(part.get("text", "")) - elif part.get("type") == "image_url": - tokens += 85 # Low detail image estimate - elif isinstance(part, str): - tokens += self.count_text(part) + # Delegate to the audited shared walker instead of a partial + # per-provider one. Each provider counter had grown its own + # shortened branch list, so every modern block priced at ~0: + # measured on a 6,800-char block this returned 8 tokens for + # tool_result, thinking, document, mcp_tool_result — and for + # output_text / refusal, which are OpenAI's OWN Responses shapes. + # The shared walker is also image-safe: a 200KB base64 image gets + # a pixel-based 1600, not the ~50K phantom text tokens a naive + # str(block) catch-all would produce. + tokens += count_content_blocks(content, self.count_text) # Name field name = message.get("name") diff --git a/headroom/providers/openai_compatible.py b/headroom/providers/openai_compatible.py index 32c88db70..da0ebd3f4 100644 --- a/headroom/providers/openai_compatible.py +++ b/headroom/providers/openai_compatible.py @@ -24,6 +24,7 @@ from dataclasses import dataclass from typing import Any from headroom.tokenizers import get_tokenizer +from headroom.tokenizers.base import count_content_blocks from .base import Provider @@ -159,12 +160,16 @@ class OpenAICompatibleTokenCounter: if isinstance(content, str): tokens += self.count_text(content) elif isinstance(content, list): - for part in content: - if isinstance(part, dict): - if part.get("type") == "text": - tokens += self.count_text(part.get("text", "")) - elif isinstance(part, str): - tokens += self.count_text(part) + # Delegate to the audited shared walker instead of a partial + # per-provider one. Each provider counter had grown its own + # shortened branch list, so every modern block priced at ~0: + # measured on a 6,800-char block this returned 8 tokens for + # tool_result, thinking, document, mcp_tool_result — and for + # output_text / refusal, which are OpenAI's OWN Responses shapes. + # The shared walker is also image-safe: a 200KB base64 image gets + # a pixel-based 1600, not the ~50K phantom text tokens a naive + # str(block) catch-all would produce. + tokens += count_content_blocks(content, self.count_text) name = message.get("name") if name: diff --git a/headroom/tokenizers/base.py b/headroom/tokenizers/base.py index e90c5534c..fcd38ad13 100644 --- a/headroom/tokenizers/base.py +++ b/headroom/tokenizers/base.py @@ -8,6 +8,7 @@ from __future__ import annotations import json from abc import ABC, abstractmethod +from collections.abc import Callable from typing import Any, Protocol, runtime_checkable @@ -359,3 +360,39 @@ class BaseTokenizer(ABC): NotImplementedError: If decoding is not supported. """ raise NotImplementedError(f"{self.__class__.__name__} does not support decoding") + + +class _DelegatingBlockCounter(BaseTokenizer): + """Adapter exposing :meth:`BaseTokenizer._count_content_parts` to non-subclasses. + + The provider token counters in ``headroom/providers/`` are not + ``BaseTokenizer`` subclasses, and each grew its own shortened content-block + walker that handled only the shapes its provider was expected to send. The + result was that every one of them priced most modern blocks at ~0: measured + on a 6,800-char block, ``OpenAITokenCounter`` returned 8 tokens for + ``tool_result``/``thinking``/``document``/``mcp_tool_result`` and — its own + Responses shapes — ``output_text``/``refusal``; ``AnthropicTokenCounter`` + returned 7 for ``thinking``/``document``, which are Anthropic's own. + + Rather than add a fifth partial walker, this lets them borrow the audited one. + It is image-safe (base64 blobs get a pixel-based estimate instead of being + serialized and priced as text) and bounds oversized blobs, which a naive + ``count_text(str(block))`` catch-all does not. + """ + + __slots__ = ("_count_text_fn",) + + def __init__(self, count_text_fn: Callable[[str], int]) -> None: + self._count_text_fn = count_text_fn + + def count_text(self, text: str) -> int: + return self._count_text_fn(text) + + +def count_content_blocks(parts: list[Any], count_text_fn: Callable[[str], int]) -> int: + """Count a multi-part content list using *count_text_fn* for text. + + Shared entry point for provider counters. See + :class:`_DelegatingBlockCounter` for why this exists. + """ + return _DelegatingBlockCounter(count_text_fn)._count_content_parts(parts) diff --git a/tests/test_provider_counter_content_blocks.py b/tests/test_provider_counter_content_blocks.py new file mode 100644 index 000000000..c4ad57088 --- /dev/null +++ b/tests/test_provider_counter_content_blocks.py @@ -0,0 +1,106 @@ +"""Provider token counters must price every content block, not just ``text``. + +Each counter in ``headroom/providers/`` had grown its own shortened content-block +walker handling only the shapes its provider was expected to send. Everything else +fell through and contributed nothing. Measured on one 6,800-char block +(``count_messages`` of a single-block message, so 7-8 is message overhead alone): + + block type OpenAI ctr Anthropic ctr + text (control) 3409 3748 + tool_result 8 3748 + thinking 8 7 + document 8 7 + mcp_tool_result 8 7 + output_text 8 7 + refusal 8 7 + +Note each counter zeroed blocks from its OWN provider — ``output_text`` and +``refusal`` are OpenAI Responses shapes, ``thinking`` and ``document`` are +Anthropic's. And these counters are what the LIVE proxy pipelines use +(``proxy/server.py`` builds them with ``AnthropicProvider`` / ``OpenAIProvider``), +so this was the main request path, not an edge case. + +The counters now delegate to ``count_content_blocks``, which reuses +``BaseTokenizer._count_content_parts``. That matters beyond coverage: a naive +``count_text(str(block))`` catch-all would serialize a base64 image and price it +as text — a 1MB screenshot reads as ~330K phantom tokens. The shared walker gives +media a pixel/byte-based estimate instead. +""" + +from __future__ import annotations + +import pytest + +from headroom.providers.anthropic import AnthropicProvider +from headroom.providers.openai import OpenAIProvider +from headroom.tokenizers.base import count_content_blocks + +_BIG = "x " * 3400 # ~6,800 chars + + +def _counters(): + return { + "openai": OpenAIProvider().get_token_counter("gpt-4o"), + "anthropic": AnthropicProvider(warn=False).get_token_counter("claude-sonnet-4-6"), + } + + +@pytest.mark.parametrize( + "block", + [ + {"type": "tool_result", "tool_use_id": "t", "content": _BIG}, + {"type": "tool_use", "id": "t", "name": "grep", "input": {"pattern": _BIG}}, + {"type": "thinking", "thinking": _BIG}, + {"type": "document", "source": {"data": _BIG}}, + {"type": "mcp_tool_result", "content": _BIG}, + {"type": "output_text", "text": _BIG}, + {"type": "refusal", "refusal": _BIG}, + {"type": "search_result", "content": _BIG}, + ], + ids=lambda b: str(b.get("type")), +) +def test_no_provider_counter_prices_a_large_block_at_zero(block: dict) -> None: + """Every one of these returned 7-8 tokens — message overhead only.""" + message = {"role": "user", "content": [block]} + for name, counter in _counters().items(): + got = counter.count_messages([message]) + assert got > 1_000, f"{name} priced a ~6,800-char {block['type']} block at {got}" + + +def test_base64_media_is_not_priced_as_text() -> None: + """The reason a str(block) catch-all would have been the wrong fix.""" + image = {"type": "image", "source": {"type": "base64", "data": "A" * 200_000}} + message = {"role": "user", "content": [image]} + + for name, counter in _counters().items(): + got = counter.count_messages([message]) + # 200KB of base64 as text would be ~50,000 tokens; the pixel estimate is 1600. + assert got < 5_000, f"{name} priced a 200KB base64 image as text: {got}" + assert got > 1_000, f"{name} ignored a declared image entirely: {got}" + + +def test_plain_text_blocks_are_unchanged() -> None: + """The control: the shape both counters already handled must not move.""" + message = {"role": "user", "content": [{"type": "text", "text": _BIG}]} + for name, counter in _counters().items(): + text_form = {"role": "user", "content": _BIG} + block_form = counter.count_messages([message]) + # A text block and the equivalent string should agree closely. + assert abs(block_form - counter.count_messages([text_form])) <= 5, name + + +def test_shared_walker_ignores_non_block_parts() -> None: + """A bare int is not a block and must contribute nothing.""" + assert count_content_blocks([123], len) == 0 + assert count_content_blocks([], len) == 0 + + +def test_shared_walker_counts_nested_tool_result_blocks() -> None: + """A tool that returns blocks nests them; they must be walked, not serialized.""" + nested = { + "type": "tool_result", + "tool_use_id": "t", + "content": [{"type": "text", "text": _BIG}], + } + flat = {"type": "tool_result", "tool_use_id": "t", "content": _BIG} + assert abs(count_content_blocks([nested], len) - count_content_blocks([flat], len)) < 100 diff --git a/tests/test_providers/test_universal.py b/tests/test_providers/test_universal.py index 26190ed46..b840837fc 100644 --- a/tests/test_providers/test_universal.py +++ b/tests/test_providers/test_universal.py @@ -200,7 +200,18 @@ class TestOpenAICompatibleProvider: assert tokens == 55 assert total == 34 - def test_openai_compatible_token_counter_ignores_unhandled_content_shapes(self, monkeypatch): + def test_openai_compatible_token_counter_prices_declared_media(self, monkeypatch): + """An image block costs tokens; a non dict/str part still contributes none. + + This previously asserted that BOTH contribute 0 — i.e. it pinned the + defect. The counter handled only ``type == "text"``, so every other block + priced at ~0: measured on a 6,800-char block, tool_result / thinking / + document / mcp_tool_result all returned 8 tokens, overhead only. Counters + now delegate to the shared walker, which prices a declared image with the + pixel-based estimate (1600, the max after provider auto-resize) rather + than either ignoring it or serializing its base64 as text. + """ + class DummyTokenizer: def count_text(self, text: str) -> int: return len(text) @@ -211,8 +222,12 @@ class TestOpenAICompatibleProvider: ) counter = OpenAICompatibleProvider().get_token_counter("demo-model") + # Non-list, non-str content is still ignored. assert counter.count_message({"role": "user", "content": {}}) == 8 - assert counter.count_message({"role": "user", "content": [{"type": "image"}, 123]}) == 8 + # A bare int is not a block and still contributes nothing. + assert counter.count_message({"role": "user", "content": [123]}) == 8 + # A declared image is now priced instead of silently free. + assert counter.count_message({"role": "user", "content": [{"type": "image"}, 123]}) == 1608 def test_get_context_limit_prefix_output_buffer_and_partial_pricing(self): provider = OpenAICompatibleProvider(