From 0e551de9d81021bb7f0dde1857a2341408606969 Mon Sep 17 00:00:00 2001
From: Mubashir R <112580905+Mubashirrrr@users.noreply.github.com>
Date: Thu, 4 Jun 2026 01:12:16 +0500
Subject: [PATCH] fix: correct tiktoken encoding for unknown gpt-4 model
 snapshots (#552)

get_encoding_for_model() resolved an unknown model to an encoding by
scanning MODEL_TO_ENCODING for the first key that starts with the
matched prefix. Because the gpt-4o entries are defined before the
plain gpt-4 entries, the "gpt-4" prefix matched "gpt-4o" first and
returned o200k_base for any gpt-4 snapshot not already in the table
(e.g. a future dated build like gpt-4-2025-01-01). The gpt-4 family
uses cl100k_base, so token counts for those models were computed with
the wrong encoding, skewing every downstream budget/truncation
decision.

Map each prefix directly to its encoding (still ordered most-specific
first) so the result is deterministic and independent of dict
insertion order.

Regression test in tests/test_tokenizers.py asserts unknown gpt-4 /
gpt-4-turbo snapshots resolve to cl100k_base while gpt-4o snapshots
stay on o200k_base. It fails before this change and passes after.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
---
 headroom/tokenizers/tiktoken_counter.py | 21 +++++++++++++++------
 tests/test_tokenizers.py                | 17 +++++++++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/headroom/tokenizers/tiktoken_counter.py b/headroom/tokenizers/tiktoken_counter.py
index 6808ed3fa..a3ccab025 100644
--- a/headroom/tokenizers/tiktoken_counter.py
+++ b/headroom/tokenizers/tiktoken_counter.py
@@ -97,13 +97,22 @@ def get_encoding_for_model(model: str) -> str:
     if model in MODEL_TO_ENCODING:
         return MODEL_TO_ENCODING[model]
 
-    # Try prefix matching for versioned models
-    for prefix in ["gpt-4o", "gpt-4-turbo", "gpt-4", "gpt-3.5", "o1", "o3"]:
+    # Try prefix matching for versioned models. Ordered most-specific first
+    # so that, e.g., "gpt-4o-*" resolves before "gpt-4-*". Each prefix maps
+    # directly to its encoding: scanning MODEL_TO_ENCODING for the first key
+    # that merely starts with the prefix is order-dependent and wrong — the
+    # "gpt-4" prefix would match the "gpt-4o" dict entry first and return
+    # o200k_base instead of cl100k_base for unknown gpt-4 snapshots.
+    for prefix, encoding in (
+        ("gpt-4o", "o200k_base"),
+        ("gpt-4-turbo", "cl100k_base"),
+        ("gpt-4", "cl100k_base"),
+        ("gpt-3.5", "cl100k_base"),
+        ("o1", "o200k_base"),
+        ("o3", "o200k_base"),
+    ):
         if model.startswith(prefix):
-            # Find any model with this prefix
-            for known_model, encoding in MODEL_TO_ENCODING.items():
-                if known_model.startswith(prefix):
-                    return encoding
+            return encoding
 
     return DEFAULT_ENCODING
 
diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py
index 99de6d096..742d5aefb 100644
--- a/tests/test_tokenizers.py
+++ b/tests/test_tokenizers.py
@@ -34,6 +34,23 @@ class TestTiktokenCounter:
         assert counter.model == "gpt-4"
         assert counter.encoding_name == "cl100k_base"
 
+    def test_unknown_gpt4_snapshot_uses_cl100k(self):
+        """Unknown gpt-4 (non-o, non-turbo) snapshots must use cl100k_base.
+
+        Regression: the prefix matcher scanned MODEL_TO_ENCODING for the
+        first key starting with the prefix. For prefix "gpt-4" that matched
+        the "gpt-4o" entry first and wrongly returned o200k_base for any
+        gpt-4 snapshot not in the table (e.g. a future dated build).
+        """
+        from headroom.tokenizers.tiktoken_counter import get_encoding_for_model
+
+        assert get_encoding_for_model("gpt-4-2025-01-01") == "cl100k_base"
+        assert get_encoding_for_model("gpt-4-future") == "cl100k_base"
+        # gpt-4o snapshots still resolve to o200k_base (most-specific first).
+        assert get_encoding_for_model("gpt-4o-2099-12-31") == "o200k_base"
+        # gpt-4-turbo snapshots use cl100k_base.
+        assert get_encoding_for_model("gpt-4-turbo-2099") == "cl100k_base"
+
     def test_count_text_empty(self):
         """Test counting empty text."""
         counter = TiktokenCounter()