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()