diff --git a/src/model_context.py b/src/model_context.py index dd32a7b..c812041 100644 --- a/src/model_context.py +++ b/src/model_context.py @@ -184,14 +184,22 @@ def get_context_length(endpoint_url: str, model: str) -> int: def _lookup_known(model: str) -> Optional[int]: - """Check known context windows by substring match.""" + """Check known context windows by substring match. + + Picks the LONGEST matching key so a short key never shadows a more specific + one. Without this, 'o1' (200k) precedes 'o1-mini' (128k) in the table and a + first-match return would report o1-mini's window as 200k. + """ name = model.lower() basename = name.split("/")[-1] if "/" in name else name basename = basename.split(":")[0] # strip :free, :extended etc. + best_key: Optional[str] = None + best_ctx: Optional[int] = None for key, ctx in KNOWN_CONTEXT_WINDOWS.items(): if key in basename or key in name: - return ctx - return None + if best_key is None or len(key) > len(best_key): + best_key, best_ctx = key, ctx + return best_ctx def _query_context_length(endpoint_url: str, model: str) -> int: diff --git a/tests/test_model_context.py b/tests/test_model_context.py index 9067b8c..20aae51 100644 --- a/tests/test_model_context.py +++ b/tests/test_model_context.py @@ -109,6 +109,19 @@ class TestLookupKnown: result = _lookup_known("deepseek-r1:free") assert result == 64000 + def test_o1_mini_not_shadowed_by_o1(self): + """'o1' (200k) precedes 'o1-mini' (128k) in the table; longest match wins.""" + assert _lookup_known("o1-mini") == 128000 + + def test_o1_full(self): + assert _lookup_known("o1") == 200000 + + def test_gpt4o_mini_not_shadowed_by_gpt4(self): + assert _lookup_known("gpt-4o-mini") == 128000 + + def test_gpt4_base(self): + assert _lookup_known("gpt-4") == 8192 + class TestGetContextLength: def setup_method(self):