diff --git a/strix/llm/context_budget.py b/strix/llm/context_budget.py index 3f153d91..b7589a9e 100644 --- a/strix/llm/context_budget.py +++ b/strix/llm/context_budget.py @@ -17,7 +17,14 @@ logger = logging.getLogger(__name__) # LiteLLM keys models without the routing prefix users type (``openai/``, # ``litellm/``, ``ollama/`` ...). Strip a leading provider segment on lookup. -_STRIPPABLE_PREFIXES = ("openai/", "litellm/", "any-llm/", "ollama/", "ollama_chat/") +_STRIPPABLE_PREFIXES = ( + "openai/", + "chatgpt/", + "litellm/", + "any-llm/", + "ollama/", + "ollama_chat/", +) _DEFAULT_OUTPUT_TOKENS = 8_192 @@ -38,7 +45,11 @@ def _safe_get_model_info(model: str) -> dict[str, Any] | None: @lru_cache(maxsize=128) def _model_info(model: str) -> dict[str, int]: - for candidate in (model, _lookup_key(model)): + lookup_key = _lookup_key(model) + # Provider-qualified ChatGPT lookups may start a synchronous device-login + # poll. LiteLLM keys the metadata by the underlying model slug. + candidates = (lookup_key,) if model.startswith("chatgpt/") else (model, lookup_key) + for candidate in candidates: info = _safe_get_model_info(candidate) if info is not None: return { diff --git a/tests/test_context_budget.py b/tests/test_context_budget.py index 8b34b483..a9a48e0c 100644 --- a/tests/test_context_budget.py +++ b/tests/test_context_budget.py @@ -21,6 +21,24 @@ def test_context_window_strips_provider_prefix() -> None: assert context_budget.context_window("openai/gpt-4o") == 128_000 +def test_context_window_chatgpt_prefix_skips_provider_auth( + monkeypatch: pytest.MonkeyPatch, +) -> None: + context_budget._model_info.cache_clear() + calls: list[str] = [] + + def _model_info(model: str) -> dict[str, int]: + calls.append(model) + return {"max_input_tokens": 1_050_000, "max_output_tokens": 128_000} + + monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _model_info) + try: + assert context_budget.context_window("chatgpt/gpt-5.6-luna") == 1_050_000 + assert calls == ["gpt-5.6-luna"] + finally: + context_budget._model_info.cache_clear() + + def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None: context_budget._model_info.cache_clear()