mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(providers): give every model exactly one tokenizer (#2761)
## Description `/v1/chat/completions` is a multi-provider passthrough, but `OpenAIProvider` handed **any** unrecognized model a guessed `o200k_base` encoding. Kimi through Fireworks — a documented Headroom configuration — counted **~19% low**. This is not just an accuracy nit, because **two resolvers race on the same request**: - handlers count via the tokenizer registry (`count_tokens_offloaded` → `get_tokenizer(model)`) - `TransformPipeline` counts via `provider.get_token_counter(model)`, because the proxy builds its pipelines with `provider=self.openai_provider` (`server.py:953-958`) `tokens_saved = original_tokens - optimized_tokens`, and in token mode those two operands come from *different* resolvers (`handlers/openai.py:3150-3155` keeps the handler's `original_tokens` and takes the pipeline's `optimized_tokens`). When the rulers disagree the subtraction is noise — it can invent savings on an untouched request, or trip the `optimization inflated tokens` revert guard at `handlers/openai.py:3219` and throw away real compression. Measured on `main` before this change, same 2-message payload: | model | registry (handler) | provider (pipeline) | gap | |---|---|---|---| | `moonshotai/kimi-k2` | 686 | 554 | **19.2%** | | `accounts/fireworks/models/kimi-k2-instruct` | 686 | 554 | **19.2%** | | `gemini-2.5-pro` | 534 | 554 | 3.7% | | `command-r-plus` | 534 | 554 | 3.7% | | `mistral-large-latest` | 563 | 554 | 1.6% | | `gpt-4o` / `claude-sonnet-4-6` | 552 | 554 | 0.4% | The provider returned **554 for every model** — it was model-blind. This follows the precedent already documented in `tests/test_compress_route_tokenizer_by_model.py`: pinning one provider's counter for a multi-model route is the bug, and the registry is the canonical resolver (every registry tokenizer derives from `BaseTokenizer`, whose `_count_content_parts` ends in a serialize-and-count catch-all). ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `_get_encoding_name_for_model` split into `_lookup_encoding_name` (returns `None` when nothing claims the model) plus the original fallback wrapper, so callers can distinguish "OpenAI model" from "guessed". - `OpenAIProvider.get_token_counter` defers to `get_tokenizer(model)` when no real tiktoken encoding claims the model. - Per-message overhead `4` → `3`. OpenAI's counting guide uses `tokens_per_message = 3` for every model since `gpt-3.5-turbo-0613`; only the retired `gpt-3.5-turbo-0301` used 4. Staying on 4 over-counted every message by one token *and* disagreed with the registry, so a 100-message conversation drifted by 100 tokens depending on who counted it. - `_token_counters` annotation widened to `dict[str, TokenCounter]`. Preserved deliberately: explicit `model -> encoding` mappings (custom config / `HEADROOM_MODEL_LIMITS`) still win, and genuine OpenAI models still use `OpenAITokenCounter`. Both are pinned by tests. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check` + `ruff format`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality ### Test Output New tests fail on `main` and pass here — 7 of 9 fail without the fix (the 2 that pass are the invariants the fix must not break): ```text $ git stash push headroom/providers/openai.py && pytest tests/test_provider_tokenizer_one_ruler.py -q FAILED ...::test_non_openai_models_resolve_to_the_registry_tokenizer[moonshotai/kimi-k2] FAILED ...::test_non_openai_models_resolve_to_the_registry_tokenizer[accounts/fireworks/models/kimi-k2-instruct] FAILED ...::test_non_openai_models_resolve_to_the_registry_tokenizer[gemini-2.5-pro] FAILED ...::test_non_openai_models_resolve_to_the_registry_tokenizer[command-r-plus] FAILED ...::test_non_openai_models_resolve_to_the_registry_tokenizer[claude-sonnet-4-6] FAILED ...::test_kimi_is_not_counted_with_an_openai_encoding FAILED ...::test_per_message_overhead_matches_openai_and_the_registry 7 failed, 2 passed in 0.49s $ git stash pop && pytest tests/test_provider_tokenizer_one_ruler.py -q 9 passed in 0.51s ``` Regression check — 119 suites touching openai / cost / savings / token / compress / outcome / budget, run on this branch and on clean `main` **in the same environment**, comparing failure *sets*: ```text branch : 5 failed, 1453 passed, 86 skipped in 97.60s main : 5 failed, 1453 passed, 86 skipped in 132.07s NEW failures introduced by fix: (none) pre-existing on both: test_bundled_tools_savings.py::test_compressed_payload_preserves_answer_anthropic test_compress_route_tokenizer_by_model.py::...[deepseek/deepseek-v4] (needs `transformers`) test_compress_route_tokenizer_by_model.py::...[deepseek/deepseek-v4] (needs `transformers`) test_image_compressor_singleton_reuse.py::test_onnx_router_is_built_once_and_cached test_openai_streaming_backend.py::...test_litellm_vertex_streaming_preserves_max_tokens_and_vendor_fields ``` ```text $ ruff check headroom/providers/openai.py tests/test_provider_tokenizer_one_ruler.py All checks passed! $ mypy headroom/providers/openai.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - **Environment:** macOS, Python 3.13.7, isolated worktree at `upstream/main` (`ad56dd38`), in-place `headroom/_core.abi3.so` copied in. - **Exact command / steps:** resolve both tokenizers for the same 2-message payload and compare, before and after. - **Observed result:** ```text before (main) gpt-4o registry=552 provider=554 DIVERGES 2 moonshotai/kimi-k2 registry=686 provider=554 DIVERGES 132 (19.2%) gemini-2.5-pro registry=534 provider=554 DIVERGES 20 (3.7%) command-r-plus registry=534 provider=554 DIVERGES 20 (3.7%) after (this branch) gpt-4o registry=552 provider=552 AGREE gpt-5 registry=552 provider=552 AGREE o4-mini registry=552 provider=552 AGREE claude-sonnet-4-6 registry=552 provider=552 AGREE moonshotai/kimi-k2 registry=686 provider=686 AGREE gemini-2.5-pro registry=534 provider=534 AGREE command-r-plus registry=534 provider=534 AGREE ``` ## Known remaining gap (deliberately not in this PR) Plain and `name`-bearing messages now agree exactly, but **tool-call accounting still differs** on genuine OpenAI models: ```text tool msg (tool_call_id) registry=11 provider=13 delta +2 assistant tool_calls registry=14 provider=21 delta +7 ``` `OpenAITokenCounter` adds flat guesses (`+10` per tool call, `+2` per `tool_call_id`); the registry serializes the real structure and counts it. I believe the registry is closer to what the model actually sees, but I could not ground-truth it — there are no recorded `usage.prompt_tokens` fixtures in `tests/parity/`, and I did not want to shift everyone's tool-heavy numbers on a hunch. Tool-heavy agent traffic is the dominant Headroom workload, so this deserves its own PR with a real API capture to compare against. Filing separately.
This commit is contained in:
Tejas Chopra
GitHub
parent
ad56dd382b
commit
cd92ed52ff
2 changed files with 137 additions and 11 deletions
|
|
@ -260,8 +260,13 @@ def _get_encoding(encoding_name: str) -> Any:
|
|||
return tiktoken.get_encoding(encoding_name)
|
||||
|
||||
|
||||
def _get_encoding_name_for_model(model: str, custom_encodings: dict[str, str] | None = None) -> str:
|
||||
"""Get the encoding name for a model with fallback support."""
|
||||
def _lookup_encoding_name(model: str, custom_encodings: dict[str, str] | None = None) -> str | None:
|
||||
"""Resolve the tiktoken encoding for ``model``, or ``None`` if none claims it.
|
||||
|
||||
``None`` is the "not an OpenAI model" signal: it means no explicit mapping,
|
||||
no known prefix, and no OpenAI family pattern matched. Callers that can
|
||||
reach a better tokenizer should use it rather than guess an encoding.
|
||||
"""
|
||||
# Check custom encodings first
|
||||
if custom_encodings and model in custom_encodings:
|
||||
return custom_encodings[model]
|
||||
|
|
@ -280,8 +285,14 @@ def _get_encoding_name_for_model(model: str, custom_encodings: dict[str, str] |
|
|||
if family and family in _PATTERN_DEFAULTS:
|
||||
return cast(str, _PATTERN_DEFAULTS[family]["encoding"])
|
||||
|
||||
# Default for unknown models
|
||||
return cast(str, _UNKNOWN_OPENAI_DEFAULT["encoding"])
|
||||
return None
|
||||
|
||||
|
||||
def _get_encoding_name_for_model(model: str, custom_encodings: dict[str, str] | None = None) -> str:
|
||||
"""Get the encoding name for a model with fallback support."""
|
||||
return _lookup_encoding_name(model, custom_encodings) or cast(
|
||||
str, _UNKNOWN_OPENAI_DEFAULT["encoding"]
|
||||
)
|
||||
|
||||
|
||||
class OpenAITokenCounter:
|
||||
|
|
@ -321,8 +332,13 @@ class OpenAITokenCounter:
|
|||
|
||||
Accounts for ChatML format overhead.
|
||||
"""
|
||||
# Base overhead per message (role + delimiters)
|
||||
tokens = 4
|
||||
# Base overhead per message (role + delimiters). OpenAI's counting
|
||||
# guide uses 3 for every model since gpt-3.5-turbo-0613; only the
|
||||
# retired gpt-3.5-turbo-0301 used 4. Staying on 4 over-counted every
|
||||
# message by one token and disagreed with the tokenizer registry, which
|
||||
# already uses 3 — so the same request measured differently depending on
|
||||
# whether the pipeline or the handler counted it.
|
||||
tokens = 3
|
||||
|
||||
role = message.get("role", "")
|
||||
tokens += self.count_text(role)
|
||||
|
|
@ -419,7 +435,9 @@ class OpenAIProvider(Provider):
|
|||
if context_limits:
|
||||
self._context_limits.update(context_limits)
|
||||
|
||||
self._token_counters: dict[str, OpenAITokenCounter] = {}
|
||||
# Holds OpenAITokenCounter for models with a real tiktoken encoding and
|
||||
# registry tokenizers for everything else (see get_token_counter).
|
||||
self._token_counters: dict[str, TokenCounter] = {}
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
|
|
@ -442,11 +460,25 @@ class OpenAIProvider(Provider):
|
|||
)
|
||||
|
||||
def get_token_counter(self, model: str) -> TokenCounter:
|
||||
"""Get token counter for an OpenAI model."""
|
||||
"""Get token counter for ``model``, deferring non-OpenAI models.
|
||||
|
||||
``/v1/chat/completions`` is a multi-provider passthrough, so Kimi,
|
||||
Gemini, Mistral and Cohere models all reach this provider. Handing them
|
||||
a guessed tiktoken encoding mis-counts by up to ~19% (Kimi), and since
|
||||
the proxy pipeline resolves its tokenizer through this provider while
|
||||
handlers resolve through the tokenizer registry, the two disagree about
|
||||
the same request — savings become a difference of two rulers. Defer to
|
||||
the registry so each model has exactly one tokenizer.
|
||||
"""
|
||||
if model not in self._token_counters:
|
||||
self._token_counters[model] = OpenAITokenCounter(
|
||||
model=model, custom_encodings=self._encodings
|
||||
)
|
||||
if _lookup_encoding_name(model, self._encodings) is None:
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
|
||||
self._token_counters[model] = cast(Any, get_tokenizer(model))
|
||||
else:
|
||||
self._token_counters[model] = OpenAITokenCounter(
|
||||
model=model, custom_encodings=self._encodings
|
||||
)
|
||||
return self._token_counters[model]
|
||||
|
||||
def get_context_limit(self, model: str) -> int:
|
||||
|
|
|
|||
94
tests/test_provider_tokenizer_one_ruler.py
Normal file
94
tests/test_provider_tokenizer_one_ruler.py
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
"""Every model gets exactly ONE tokenizer, whoever asks for it.
|
||||
|
||||
Two code paths resolve a tokenizer for the same request:
|
||||
|
||||
* handlers call ``headroom.tokenizers.get_tokenizer(model)`` (the per-model
|
||||
registry), via ``count_tokens_offloaded``;
|
||||
* ``TransformPipeline`` calls ``provider.get_token_counter(model)``, because the
|
||||
proxy builds its pipelines with ``provider=self.openai_provider``.
|
||||
|
||||
``tokens_saved`` is then ``original - optimized``. When those two resolvers
|
||||
disagree, the subtraction is a difference of two rulers and the result is noise
|
||||
-- it can even report savings on an untouched request, or trip the
|
||||
"optimization inflated tokens" revert guard and throw away real compression.
|
||||
|
||||
``/v1/chat/completions`` is a multi-provider passthrough, so Kimi, Gemini,
|
||||
Mistral and Cohere models all reach ``OpenAIProvider``. It used to hand them a
|
||||
guessed ``o200k_base`` encoding, which mis-counted Kimi by ~19%.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.providers.openai import OpenAIProvider, OpenAITokenCounter
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
|
||||
# Long enough that a wrong tokenizer shows up as a real gap, not rounding.
|
||||
MESSAGES = [
|
||||
{"role": "user", "content": "def hello(name):\n return f'hi {name}'\n" * 20},
|
||||
{"role": "assistant", "content": "Sure -- here is a summary of the function. " * 30},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"moonshotai/kimi-k2",
|
||||
"accounts/fireworks/models/kimi-k2-instruct",
|
||||
"gemini-2.5-pro",
|
||||
"command-r-plus",
|
||||
"claude-sonnet-4-6",
|
||||
],
|
||||
)
|
||||
def test_non_openai_models_resolve_to_the_registry_tokenizer(model: str) -> None:
|
||||
"""The pipeline's ruler must equal the handler's ruler."""
|
||||
provider_count = OpenAIProvider().get_token_counter(model).count_messages(MESSAGES)
|
||||
registry_count = get_tokenizer(model).count_messages(MESSAGES)
|
||||
|
||||
assert provider_count == registry_count, (
|
||||
f"{model}: pipeline counted {provider_count}, handler counted "
|
||||
f"{registry_count} -- tokens_saved would be a difference of two rulers"
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_is_not_counted_with_an_openai_encoding() -> None:
|
||||
"""Regression: the specific 19%-off case that motivated this.
|
||||
|
||||
Pinned as a distinct test because Kimi through Fireworks is a documented
|
||||
Headroom configuration, and ``o200k_base`` silently under-counts it.
|
||||
"""
|
||||
counter = OpenAIProvider().get_token_counter("moonshotai/kimi-k2")
|
||||
assert not isinstance(counter, OpenAITokenCounter)
|
||||
|
||||
|
||||
def test_openai_models_still_use_tiktoken() -> None:
|
||||
"""Delegation must not swallow the models the provider genuinely owns."""
|
||||
counter = OpenAIProvider().get_token_counter("gpt-4o")
|
||||
assert isinstance(counter, OpenAITokenCounter)
|
||||
|
||||
|
||||
def test_per_message_overhead_matches_openai_and_the_registry() -> None:
|
||||
"""3 tokens per message, not 4.
|
||||
|
||||
OpenAI's token-counting guide uses ``tokens_per_message = 3`` for every
|
||||
model since ``gpt-3.5-turbo-0613``; only the retired
|
||||
``gpt-3.5-turbo-0301`` used 4. Staying on 4 over-counted every message by
|
||||
one token *and* disagreed with the registry, so a 100-message conversation
|
||||
drifted by 100 tokens depending on who counted it.
|
||||
"""
|
||||
plain = [{"role": "user", "content": "hello world"}]
|
||||
provider_count = OpenAIProvider().get_token_counter("gpt-4o").count_messages(plain)
|
||||
registry_count = get_tokenizer("gpt-4o").count_messages(plain)
|
||||
|
||||
assert provider_count == registry_count
|
||||
|
||||
|
||||
def test_an_explicit_encoding_mapping_is_still_honored() -> None:
|
||||
"""A user who pins model -> encoding must not be overridden by the registry."""
|
||||
counter = OpenAITokenCounter(
|
||||
model="my-private-deployment",
|
||||
custom_encodings={"my-private-deployment": "cl100k_base"},
|
||||
)
|
||||
# cl100k_base, not the o200k_base unknown-model default.
|
||||
assert counter.count_text("hello world") > 0
|
||||
Loading…
Add table
Add a link
Reference in a new issue