fix(tokenizers): don't tokenize image blocks as text in TiktokenCounter (#2093)

## Description

`TiktokenCounter.count_messages` explodes the token count for any
content block that isn't plain text or an OpenAI `image_url`.

The multi-part content loop handles exactly two shapes:

```python
if part.get("type") == "text":
    total += self.count_text(part.get("text", ""))
elif part.get("type") == "image_url":
    ...  # 85 / 170 tokens by detail
else:
    total += self.count_text(str(part))   # <-- everything else
```

Every other block shape reaching the `else` gets `str(part)`-ified and
tokenized as text. That includes Anthropic's `{"type": "image",
"source": {"type": "base64", "data": "<...>"}}`, `tool_result`,
`tool_use`, and the Strands SDK blocks. Over the wire the image `data`
is a base64 string, so a 1MB image turns into ~1.4M characters of "text"
and is counted as **~330K tokens** for a single image (a ~218x overcount
in a standalone repro). Anything that relies on the count — budget
gating, the compress/skip decision, savings math — is thrown off for
multimodal requests that route through the tiktoken counter.

The base class already solved this: `BaseTokenizer._count_content_parts`
prices `image`/`image_url`/`input_image` at a flat bounded estimate and
has a comment stating it exists specifically to stop "a 1MB image =
~330K fake tokens". The tiktoken override just never delegated to it for
the non-text shapes.

## Fix

Delegate unknown block shapes in the `else` branch to
`self._count_content_parts([part])` instead of stringifying them. `text`
and `image_url` keep the existing tiktoken-specific handling (including
the 85/170 detail split); everything else now gets the base handler's
bounded pricing.

Closes #

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/tokenizers/tiktoken_counter.py`: the `count_messages`
multi-part `else` branch delegates to the base
`_count_content_parts([part])` rather than `count_text(str(part))`.
- `tests/test_tokenizers.py`: add
`test_count_messages_image_block_is_not_stringified` — a base64 image
block inside list content must stay bounded (well under the tens of
thousands of tokens the blob would produce as text).
- `CHANGELOG.md`: Bug Fixes entry.

## Testing

- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed

### Test Output

```text
$ uvx ruff@0.15.17 check headroom/tokenizers/tiktoken_counter.py tests/test_tokenizers.py
All checks passed!
$ python -m py_compile headroom/tokenizers/tiktoken_counter.py tests/test_tokenizers.py
OK
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17`. Importing
`headroom` pulls in the torch/transformers stack and a full `pytest`
gets OOM-killed on this box, so I verified the magnitude with a
dependency-free script that models the old `count_text(str(part))` path
against the base handler's bounded image estimate, and left the full
pytest to CI.
- Exact command / steps: built a ~1MB PNG as an Anthropic `image` block
with the payload base64-encoded (as it arrives over the wire), computed
the old path (`len(str(part)) / ~4` chars-per-token) versus the new path
(base handler prices an image block at a flat 1600).
- Observed result: base64 payload ~1,398,112 chars; old path ~349,549
tokens; new path 1,600 tokens; ~218x overcount removed. The new
regression test asserts the counted total for such a message stays under
5000.
- Not tested: a live tiktoken end-to-end count through the proxy; full
local `pytest` deferred to CI (OOM, per above).

## 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
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

The "unit tests pass locally" and "type checking" boxes are unchecked
because the full suite imports the ML stack, which I can't run in this
environment; the change is a localized delegation to an existing base
method, verified by the standalone magnitude proof and the new
regression test for CI. This mirrors the earlier base-handler
`tool_result` list-recursion fix — same class of "don't count a base64
blob as text" bug, in the tiktoken override this time.

Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
This commit is contained in:
Abhay Singh 2026-07-13 19:54:51 +05:30 committed by GitHub
parent b097ef3e25
commit ae10d6c99d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 47 additions and 1 deletions

View file

@ -102,6 +102,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Bug Fixes
* **tokenizers:** stop `TiktokenCounter.count_messages` from exploding on non-text content blocks. Its multi-part branch handled only `text` and OpenAI `image_url`; every other shape (Anthropic `image`/`tool_result`/`tool_use`, Strands blocks) fell through to `count_text(str(part))`, which json-stringified the base64 payload and tokenized it as text — a 1MB image counted as ~330K phantom tokens (~218x overcount in a standalone repro), corrupting every downstream budgeting/compression decision for multimodal OpenAI-model requests. Unknown block shapes now delegate to the base `_count_content_parts`, which prices images/documents by a bounded estimate (the overcount that helper already exists to prevent).
* **install:** don't let a host env export override the manifest in persistent-docker deployments. `build_runtime_command` emitted the manifest's pinned `--env NAME=VALUE` pairs and then, for every host var matching a passthrough prefix, a bare `--env NAME`. Docker resolves duplicate `--env` last-wins, so a stale host export (e.g. `HEADROOM_BACKEND=anyllm`) that shared a passthrough prefix with a pinned manifest value (`HEADROOM_BACKEND=anthropic`) was appended after it and silently won, diverging the container from its deployment config. The bare passthrough is now skipped for any name the manifest already pins.
* **memory:** honor explicit `store=false` on OpenAI `/v1/responses` requests by skipping Headroom memory-tool injection that depends on stored-response continuations. Memory context injection stays available, and requests no longer get rewritten to `store=true` behind the client's back ([#1944](https://github.com/headroomlabs-ai/headroom/issues/1944)).
* **proxy/batch:** stop corrupting Google `batchGenerateContent` requests whose contents interleave text turns with text-less entries (functionCall/functionResponse/images). The batch handler restored preserved (non-text) entries with the raw-index loop that #836 replaced everywhere else — indexing the shorter `optimized_contents` (text-less entries produce no message) by the original `contents[]` index, which overwrites the wrong entry and drops any preserved entry whose original index is past the optimized length. A request like `[user text, model functionCall, user functionResponse, model text]` was forwarded to Google as two entries: the model's answer overwritten by the functionCall and the functionResponse dropped. The batch handler now uses the shared `_rebuild_gemini_contents` interleaving helper, so all entries survive in order.

View file

@ -282,7 +282,15 @@ class TiktokenCounter(BaseTokenizer):
else:
total += 170 # Base for high detail
else:
total += self.count_text(str(part))
# Any other block shape (Anthropic
# image/tool_result/tool_use, Strands blocks)
# is priced by the base handler, which uses a
# bounded per-image/document estimate. Stringifying
# it here would json-serialize a base64 blob and
# count it as text — a 1MB image becomes ~330K
# phantom tokens (the exact overcount base.py
# _count_content_parts exists to prevent).
total += self._count_content_parts([part])
elif isinstance(part, str):
total += self.count_text(part)
elif key == "role":

View file

@ -103,6 +103,43 @@ class TestTiktokenCounter:
count = counter.count_messages(messages)
assert count > 0
def test_count_messages_image_block_is_not_stringified(self):
"""An Anthropic-style image block must be priced as an image, not text.
Over the wire the image arrives as a base64 string inside list content.
The old count_messages else-branch stringified any non-text/non-image_url
part and tokenized it as text, so a 1MB image counted as ~330K phantom
tokens. The base handler prices image blocks by a bounded estimate, so the
count must stay small regardless of the base64 payload size.
"""
import base64
counter = TiktokenCounter()
blob = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * 200_000).decode()
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is in this screenshot?"},
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": blob,
},
},
],
}
]
count = counter.count_messages(messages)
# The base64 blob alone would be tens of thousands of text tokens; a
# bounded image estimate keeps the whole message well under that.
assert count < 5000, count
assert count < len(blob) // 10
def test_encode_decode_roundtrip(self):
"""Test encode/decode roundtrip."""
counter = TiktokenCounter()