1 commit

Author	SHA1	Message	Date
Tejas Chopra	cd92ed52ff	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.	2026-08-03 23:46:27 -07:00