mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
b097ef3e25
commit
ae10d6c99d
3 changed files with 47 additions and 1 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue