diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 1ce7b0fca..242e0aaa2 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -60,6 +60,7 @@ from ..config import ( ) from ..parser import CCR_RETRIEVAL_MARKER_RE from ..tokenizer import Tokenizer +from ..tokenizers.base import count_content_blocks from ..tokenizers.estimator import EstimatingTokenCounter from . import mixed_content as _mixed_content from .base import Transform @@ -1176,42 +1177,29 @@ def _gain_bucket(gain: float) -> str: def _netcost_message_tokens(message: dict[str, Any], tokenizer: Tokenizer) -> int: """Token count of a message for net-cost suffix (S) estimation. - String content is counted directly. Anthropic block-list content is - counted by summing the text-bearing fields (``text`` blocks and - ``tool_result`` content) rather than stringifying the whole list, which - would count Python ``repr`` punctuation and type names and badly - miscount S — the value that drives the break-even gate decision. + String content is counted directly. Block-list content is delegated to the + canonical block counter, which knows how to price non-text blocks. + + This function used to walk the list itself and fall back to + ``str(block)`` for anything that was not ``text`` or ``tool_result``, on the + stated assumption that such blocks "rarely dominate a suffix". An ``image`` + block is the exception that breaks it: ``str()`` embeds the whole base64 + payload, so one screenshot counted ~100,000 tokens instead of ~1,600 + (57x-146x over, growing with image size). + + That mattered because S is the cache-bust cost — the tokens re-written if + message *j* is mutated — so an image inflated S for **every message before + it**, and the break-even gate then refused to compress any of them. + ``BaseTokenizer._count_content_parts`` already solves this (see its "1MB + image = ~330K fake tokens without this" guard); this walk simply predated + it. Delegating also means new block types are priced in one place. """ content = message.get("content", "") if isinstance(content, str): return tokenizer.count_text(content) if not isinstance(content, list): return tokenizer.count_text(str(content)) - total = 0 - for block in content: - if not isinstance(block, dict): - total += tokenizer.count_text(str(block)) - continue - block_type = block.get("type") - if block_type == "text": - total += tokenizer.count_text(str(block.get("text", ""))) - elif block_type == "tool_result": - tc = block.get("content", "") - if isinstance(tc, str): - total += tokenizer.count_text(tc) - elif isinstance(tc, list): - for sub in tc: - if isinstance(sub, dict) and sub.get("type") == "text": - total += tokenizer.count_text(str(sub.get("text", ""))) - else: - total += tokenizer.count_text(str(sub)) - else: - total += tokenizer.count_text(str(tc)) - else: - # Other blocks (image, tool_use input, …) — repr is a rough proxy - # but bounded; these rarely dominate a suffix. - total += tokenizer.count_text(str(block)) - return total + return count_content_blocks(content, tokenizer.count_text) class CompressionCache: diff --git a/tests/test_netcost_gate.py b/tests/test_netcost_gate.py index 0bb71e8df..8a070eafe 100644 --- a/tests/test_netcost_gate.py +++ b/tests/test_netcost_gate.py @@ -148,8 +148,13 @@ class TestNetCostHelpers: assert _gain_bucket(float("inf")) == "nan" def test_message_tokens_block_list_beats_repr(self, tokenizer): - # str(content) over a block list counts repr punctuation/type names; - # the block-aware helper counts only the text-bearing payload. + # str(content) over a block list embeds the whole base64 payload; the + # block-aware helper prices the image at its pixel cost instead. + # + # This used to use a 500-char stub image, which is *smaller* than a + # single image's real token cost -- so repr looked cheap and the + # payload-scaling bug stayed invisible. Use a realistically sized + # payload, which is what actually occurs (screenshots). from headroom.transforms.content_router import _netcost_message_tokens text = "word " * 200 @@ -157,15 +162,17 @@ class TestNetCostHelpers: "role": "user", "content": [ {"type": "text", "text": text}, - {"type": "image", "source": {"data": "x" * 500}}, + {"type": "image", "source": {"data": "x" * 200_000}}, ], } helper = _netcost_message_tokens(block_msg, tokenizer) text_only = tokenizer.count_text(text) - # Helper tracks the text payload closely; the image block adds only a - # small repr proxy, far less than stringifying the whole list. - assert abs(helper - text_only) < text_only * 0.5 - assert helper < tokenizer.count_text(str(block_msg["content"])) + # The text payload is still counted in full, and the image adds a + # bounded pixel-based cost rather than a payload-scaled one. + assert helper >= text_only + assert helper - text_only <= 2000 + # ...which is dramatically less than stringifying the whole list. + assert helper < tokenizer.count_text(str(block_msg["content"])) / 10 def test_message_tokens_tool_result_blocks(self, tokenizer): from headroom.transforms.content_router import _netcost_message_tokens diff --git a/tests/test_netcost_suffix_image_tokens.py b/tests/test_netcost_suffix_image_tokens.py new file mode 100644 index 000000000..1355c309d --- /dev/null +++ b/tests/test_netcost_suffix_image_tokens.py @@ -0,0 +1,93 @@ +"""An image must not inflate the net-cost suffix (S). + +``_netcost_message_tokens`` used to walk block-list content itself and fall back +to ``str(block)`` for anything that was not ``text`` or ``tool_result``, on the +stated assumption that such blocks "rarely dominate a suffix". An ``image`` block +breaks that assumption completely: ``str()`` embeds the whole base64 payload. + + 512x512 PNG 20,034 counted vs ~349 real 57x + 1092x1092 shot 100,034 counted vs ~1,589 real 63x + 1568x1568 233,367 counted vs ~1,600 real 146x + +S is the cache-bust cost -- the tokens re-written if message *j* is mutated -- so +an image inflated S for **every message before it**, and the break-even gate then +declined to compress any of them. A single screenshot could switch off net-cost +compression for the whole earlier conversation. +""" + +from __future__ import annotations + +import base64 + +import pytest + +from headroom.tokenizer import Tokenizer +from headroom.tokenizers import get_tokenizer +from headroom.transforms.content_router import _netcost_message_tokens + + +@pytest.fixture +def tok() -> Tokenizer: + return Tokenizer(get_tokenizer("claude-sonnet-4-6"), "claude-sonnet-4-6") + + +def _image_block(payload_bytes: int) -> dict: + data = base64.b64encode(b"\x89PNG" + b"\x00" * payload_bytes).decode() + return { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": data}, + } + + +@pytest.mark.parametrize("payload_bytes", [120_000, 600_000, 1_400_000]) +def test_image_is_not_counted_as_its_base64_payload(tok: Tokenizer, payload_bytes: int) -> None: + """Cost must not scale with the base64 length.""" + message = {"role": "user", "content": [_image_block(payload_bytes)]} + + counted = _netcost_message_tokens(message, tok) + + # Anthropic caps image cost around 1600 tokens; anything in the tens of + # thousands means the payload is being counted as text. + assert counted <= 2000, f"image counted as {counted:,} tokens" + + +def test_image_cost_does_not_grow_with_payload_size(tok: Tokenizer) -> None: + """A 12x larger payload must not cost ~12x more.""" + small = _netcost_message_tokens({"role": "user", "content": [_image_block(120_000)]}, tok) + large = _netcost_message_tokens({"role": "user", "content": [_image_block(1_400_000)]}, tok) + + assert large == small + + +@pytest.mark.parametrize( + "content", + [ + [{"type": "text", "text": "hello world " * 50}], + [{"type": "tool_result", "content": "result text " * 40}], + [{"type": "tool_result", "content": [{"type": "text", "text": "x " * 60}]}], + ], + ids=["text", "tool_result_str", "tool_result_list"], +) +def test_text_bearing_blocks_are_unchanged(tok: Tokenizer, content: list) -> None: + """Delegation must be behaviour-preserving for what already worked. + + These are the shapes the old local walk handled correctly; pinning them + keeps the delegation from quietly changing suffix sizes on normal traffic. + """ + counted = _netcost_message_tokens({"role": "user", "content": content}, tok) + text = "".join( + block.get("text", "") + or (block.get("content") if isinstance(block.get("content"), str) else "") + or "".join( + sub.get("text", "") for sub in (block.get("content") or []) if isinstance(sub, dict) + ) + for block in content + ) + + assert counted == pytest.approx(tok.count_text(text), abs=2) + + +def test_plain_string_content_still_counted(tok: Tokenizer) -> None: + message = {"role": "user", "content": "plain " * 20} + + assert _netcost_message_tokens(message, tok) == tok.count_text(message["content"])