diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index 73a953c21..d21073bff 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -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: diff --git a/tests/test_provider_tokenizer_one_ruler.py b/tests/test_provider_tokenizer_one_ruler.py new file mode 100644 index 000000000..f9159b5f2 --- /dev/null +++ b/tests/test_provider_tokenizer_one_ruler.py @@ -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