diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index d21073bff..56f2a04db 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -64,6 +64,10 @@ _MODEL_ENCODINGS: dict[str, str] = { "o1-mini": "o200k_base", "o3": "o200k_base", "o3-mini": "o200k_base", + "o4": "o200k_base", + "o4-mini": "o200k_base", + "gpt-4.1": "o200k_base", + "gpt-5": "o200k_base", # GPT-4 and GPT-3.5 use cl100k_base "gpt-4": "cl100k_base", "gpt-4-turbo": "cl100k_base", @@ -78,6 +82,18 @@ _CONTEXT_LIMITS: dict[str, int] = { "gpt-4o-2024-11-20": 128000, "gpt-4o-2024-08-06": 128000, "gpt-4o-2024-05-13": 128000, + # GPT-4.1 series (~1M input). LiteLLM is still consulted first in + # get_context_limit; these are the manual fallback for installs without it + # (the litellm dep is gated python_version < '3.14'). + "gpt-4.1": 1_047_576, + "gpt-4.1-mini": 1_047_576, + "gpt-4.1-nano": 1_047_576, + # GPT-5 series. This table is an INPUT budget (get_context_limit returns + # litellm's max_input_tokens when available), so these are 272K input -- + # not the 400K total window, which is 272K in + 128K out. + "gpt-5": 272000, + "gpt-5-mini": 272000, + "gpt-5-nano": 272000, # GPT-4 Turbo "gpt-4-turbo": 128000, "gpt-4-turbo-preview": 128000, @@ -94,6 +110,7 @@ _CONTEXT_LIMITS: dict[str, int] = { "o1-mini": 128000, "o3": 200000, "o3-mini": 200000, + "o4-mini": 200000, # DeepSeek (often accessed via OpenAI-compatible API). Values verified # against api-docs.deepseek.com (V4) and LiteLLM model_cost (deprecated # aliases). LiteLLM lookup is still attempted first in get_context_limit; @@ -275,10 +292,12 @@ def _lookup_encoding_name(model: str, custom_encodings: dict[str, str] | None = if model in _MODEL_ENCODINGS: return _MODEL_ENCODINGS[model] - # Prefix match for versioned models - for prefix, encoding in _MODEL_ENCODINGS.items(): + # Prefix match for versioned models, longest prefix first. Plain dict order + # let a shorter family shadow a longer one: "gpt-4.1" hit the "gpt-4" entry + # and got cl100k_base instead of o200k_base, which over-counts CJK by ~33%. + for prefix in sorted(_MODEL_ENCODINGS, key=len, reverse=True): if model.startswith(prefix): - return encoding + return _MODEL_ENCODINGS[prefix] # Pattern-based inference family = _infer_model_family(model) @@ -525,10 +544,12 @@ class OpenAIProvider(Provider): if model in self._context_limits: return self._context_limits[model] - # Prefix match - for prefix, limit in self._context_limits.items(): + # Prefix match, longest prefix first. Plain dict order let a shorter + # family shadow a longer one: "gpt-4.1" hit the "gpt-4" entry and got + # 8192 instead of ~1M, and "gpt-4-32k-0613" got 8192 instead of 32768. + for prefix in sorted(self._context_limits, key=len, reverse=True): if model.startswith(prefix): - return limit + return self._context_limits[prefix] # Pattern-based inference family = _infer_model_family(model) diff --git a/tests/test_openai_model_table_resolution.py b/tests/test_openai_model_table_resolution.py new file mode 100644 index 000000000..7148fc294 --- /dev/null +++ b/tests/test_openai_model_table_resolution.py @@ -0,0 +1,84 @@ +"""A shorter model family must not shadow a longer one. + +``_MODEL_ENCODINGS`` and ``_CONTEXT_LIMITS`` are matched by prefix. Iterating +them in plain dict order meant the first *inserted* prefix won, not the most +specific one, so ``gpt-4.1`` matched the ``gpt-4`` entry: + +* context limit 8192 instead of ~1M -- a 128x under-estimate, which makes the + proxy think a 1M-context model is nearly full and compress accordingly; +* encoding ``cl100k_base`` instead of ``o200k_base``, which over-counts CJK + text by ~33%. + +``gpt-4-32k-0613`` had the same problem (8192 instead of 32768). + +``get_context_limit`` consults LiteLLM before this table, so the limit half only +surfaces where LiteLLM is missing or does not know the model -- notably any +install on Python >= 3.14, where the ``litellm`` dependency is excluded by its +``python_version < '3.14'`` marker. The encoding half has no such fallback and +was always wrong. +""" + +from __future__ import annotations + +import pytest + +from headroom.providers.openai import ( + OpenAIProvider, + _get_encoding_name_for_model, +) + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + # The shadowing cases. + ("gpt-4.1", 1_047_576), + ("gpt-4.1-mini", 1_047_576), + ("gpt-4.1-nano", 1_047_576), + ("gpt-4.1-2025-04-14", 1_047_576), + ("gpt-4-32k-0613", 32768), + # Newer families that fell through to the unknown-model default. + ("gpt-5", 272_000), + ("gpt-5-mini", 272_000), + ("o4-mini", 200_000), + # Must not regress. + ("gpt-4", 8192), + ("gpt-4-turbo", 128_000), + ("gpt-4o", 128_000), + ("o3", 200_000), + ("gpt-3.5-turbo", 16385), + ], +) +def test_context_limit_prefers_the_most_specific_prefix(model: str, expected: int) -> None: + assert OpenAIProvider()._get_context_limit_manual(model) == expected + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("gpt-4.1", "o200k_base"), + ("gpt-4.1-mini", "o200k_base"), + ("gpt-4.1-2025-04-14", "o200k_base"), + ("gpt-5", "o200k_base"), + ("gpt-5-mini", "o200k_base"), + ("o4-mini", "o200k_base"), + # Must not regress: these genuinely are cl100k_base. + ("gpt-4", "cl100k_base"), + ("gpt-4-turbo", "cl100k_base"), + ("gpt-3.5-turbo", "cl100k_base"), + ("gpt-4o", "o200k_base"), + ], +) +def test_encoding_prefers_the_most_specific_prefix(model: str, expected: str) -> None: + assert _get_encoding_name_for_model(model) == expected + + +def test_cjk_is_not_over_counted_for_gpt_41() -> None: + """The concrete cost of picking cl100k_base for a gpt-4.1 request.""" + tiktoken = pytest.importorskip("tiktoken") + text = "这是一个测试文档,用于验证分词器的差异。" * 30 + + chosen = _get_encoding_name_for_model("gpt-4.1") + assert len(tiktoken.get_encoding(chosen).encode(text)) == len( + tiktoken.get_encoding("o200k_base").encode(text) + ) diff --git a/tests/test_provider_model_fallback.py b/tests/test_provider_model_fallback.py index 9de7bc1e1..b19b4e49b 100644 --- a/tests/test_provider_model_fallback.py +++ b/tests/test_provider_model_fallback.py @@ -276,10 +276,24 @@ class TestOpenAIModelFallback: """Test fallback for unknown models.""" provider = OpenAIProvider() - # Unknown model should get 128K default - limit = provider.get_context_limit("gpt-5-future") + # Unknown model should get 128K default. Deliberately a name that + # matches no known family prefix -- this used to say "gpt-5-future", + # which stopped being unknown once gpt-5 was added to _CONTEXT_LIMITS. + limit = provider.get_context_limit("gpt-9-imaginary") assert limit == 128000 + def test_unknown_variant_inherits_its_family_limit(self): + """An unrecognized variant of a *known* family takes that family's limit. + + This is the same prefix inheritance that gives "gpt-4o-2024-11-20" the + gpt-4o limit, and it is strictly better than dropping such a model to + the generic 128K default. + """ + provider = OpenAIProvider() + + assert provider.get_context_limit("gpt-5-future") == 272000 + assert provider.get_context_limit("gpt-4.1-preview") == 1_047_576 + def test_no_exception_for_unknown_model(self): """Test that unknown models don't raise exceptions.""" provider = OpenAIProvider()