mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(providers): stop pricing modern content blocks at zero (#2760)
## Description Each token 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, via `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** | Two things make this worse than a coverage gap: 1. **Each counter zeroed blocks from its own provider.** `output_text` and `refusal` are OpenAI Responses shapes; `thinking` and `document` are Anthropic's. 2. **These are the counters the live pipelines use.** `proxy/server.py` builds them with `AnthropicProvider` / `OpenAIProvider`, so this is the main request path — not an edge case. #2743 fixed this for `/v1/compress` only, by routing that route to the registry tokenizers, whose `BaseTokenizer` walker is complete. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made Rather than add a **fifth** partial walker, the counters now delegate to the audited one: - `tokenizers/base.py` — new `count_content_blocks(parts, count_text_fn)` plus a thin `_DelegatingBlockCounter` adapter, since the provider counters are not `BaseTokenizer` subclasses. `BaseTokenizer` itself is untouched. - `providers/openai.py`, `providers/anthropic.py`, `providers/openai_compatible.py` — list-content branches delegate. **Why delegate instead of adding a `count_text(str(block))` catch-all:** that would serialize a base64 blob and price it as text. `tiktoken_counter.py` already documents the failure — a 1MB image becomes ~330K phantom tokens. The shared walker gives media a pixel/byte-based estimate. **Scope:** the three counters that accumulate token counts. `google.py` and `cohere.py` extract a *text string* first and count that, so the same defect there needs a differently-shaped fix — left as a follow-up. ## Testing - [x] Unit tests pass - [x] Linting passes (`ruff check` + `format --check`, pinned 0.15.17) - [x] New tests added - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_provider_counter_content_blocks.py -q 12 passed in 0.52s $ uvx ruff@0.15.17 check headroom/ tests/... --exclude headroom/dashboard/templates All checks passed! ``` **After the fix**, every shape lands within ~1% of the equivalent plain text, and media stays bounded: ```text block OpenAI Anthropic text (control) 3409 3748 tool_result 3409 3748 thinking 3419 3759 document 3423 3763 mcp_tool_result 3422 3762 output_text 3420 3760 refusal 3421 3761 image b64 200KB 1608 1607 <- pixel estimate, not ~50K as text ``` ## Real Behavior Proof — including a regression I caught **This change flipped an existing test**, and I only found it because every suite was run against clean `upstream/main` in the same environment with the failure sets diffed: ```text before the test rewrite: upstream/main : 1 failed, 104 passed this branch : 2 failed, 103 passed <- regression diff : + test_openai_compatible_token_counter_ignores_unhandled_content_shapes ``` That test asserted `content: [{"type": "image"}, 123] == 8` — i.e. it **pinned the defect**, that unhandled shapes contribute nothing. Rewritten as `..._prices_declared_media`: a declared image is now priced (1608) while a bare int is still correctly ignored (8), with the rationale in the docstring. ```text after the rewrite: upstream/main : 1 failed, 104 passed, 10 skipped, 25 errors this branch : 1 failed, 104 passed, 10 skipped, 25 errors failure sets : IDENTICAL ``` - **Pre-existing, not from this change:** the 1 failure and all 25 errors. The errors are all in `test_compress_route_tokenizer_by_model.py`, whose loopback `TestClient` fixture this throwaway env cannot satisfy. - **Environment note:** `content_router` and several suites need the compiled `headroom._core`, which isn't in a fresh worktree (gitignored, built in-place). I copied the built `.so` in to run these and removed it before committing. - **Not tested:** no live provider call, so the *absolute* accuracy of the 1600 image estimate against a real Anthropic/OpenAI bill is unverified — it is the value `BaseTokenizer` already used, and this PR only changes which blocks reach it. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] I did **not** edit `CHANGELOG.md` ## Related Fourth PR from one tokenizer-consistency audit: #2757 (litellm total prompt / `--budget`), #2758 (HuggingFace chat templates, `gpt-5`, gateway-wrapped names), #2759 (router token units). Plus #2756, which splits the local/provider token scales in `RequestOutcome`. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
Tejas Chopra
GitHub
parent
0ed306b22b
commit
06add9e9d8
6 changed files with 193 additions and 25 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
106
tests/test_provider_counter_content_blocks.py
Normal file
106
tests/test_provider_counter_content_blocks.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue