From 21afdaea9e2b192fbbb5bfb78a7d229356ee3c7f Mon Sep 17 00:00:00 2001 From: yoni Date: Fri, 14 Aug 2026 17:09:39 +0000 Subject: [PATCH] fix(llm): resolve grok/ models to xai/ for LiteLLM metadata LiteLLM maps xAI models only provider-qualified, so neither "grok/grok-4" nor bare "grok-4" resolves: subscription runs fell back to the generic 200k context window and an 8k output cap instead of Grok's 256k/256k. subscription.litellm_model_name() now owns the routing-prefix -> LiteLLM name mapping (grok/ -> xai/, chatgpt/ -> bare) and context_budget uses it. --- strix/config/subscription.py | 18 ++++++++++++++++++ strix/llm/context_budget.py | 14 +++++++++----- tests/test_context_budget.py | 11 +++++++++++ tests/test_grok_routing.py | 11 +++++++++++ 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/strix/config/subscription.py b/strix/config/subscription.py index 204c05e1..2641c3c4 100644 --- a/strix/config/subscription.py +++ b/strix/config/subscription.py @@ -22,6 +22,10 @@ _PROVIDERS: tuple[ModuleType, ...] = (codex, grok) # Human-facing provider names keyed by each module's ``PROVIDER`` constant. _DISPLAY_NAMES: dict[str, str] = {codex.PROVIDER: "ChatGPT", grok.PROVIDER: "Grok"} +# Prefix LiteLLM keys each provider's model metadata under: ChatGPT models are +# mapped bare ("gpt-5.4"), xAI's only provider-qualified ("xai/grok-4"). +_LITELLM_PREFIXES: dict[str, str] = {codex.PROVIDER: "", grok.PROVIDER: "xai/"} + def provider_for_model(model_name: str | None) -> ModuleType | None: """Return the subscription provider module that owns ``model_name``'s prefix, @@ -43,3 +47,17 @@ def provider_label(model_name: str | None) -> str | None: if provider is None: return None return _DISPLAY_NAMES.get(provider.PROVIDER) + + +def litellm_model_name(model_name: str | None) -> str | None: + """``model_name`` rewritten to the name LiteLLM maps metadata under. + + Subscription prefixes are Strix routing labels LiteLLM never maps, so a + lookup of "grok/grok-4" (or bare "grok-4") finds nothing. Non-subscription + models are returned unchanged. + """ + provider = provider_for_model(model_name) + if provider is None: + return model_name + prefix = _LITELLM_PREFIXES.get(provider.PROVIDER, "") + return f"{prefix}{provider.subscription_model(model_name)}" diff --git a/strix/llm/context_budget.py b/strix/llm/context_budget.py index b7589a9e..1c18592b 100644 --- a/strix/llm/context_budget.py +++ b/strix/llm/context_budget.py @@ -10,7 +10,7 @@ from typing import Any import litellm -from strix.config import load_settings +from strix.config import load_settings, subscription logger = logging.getLogger(__name__) @@ -19,7 +19,6 @@ logger = logging.getLogger(__name__) # ``litellm/``, ``ollama/`` ...). Strip a leading provider segment on lookup. _STRIPPABLE_PREFIXES = ( "openai/", - "chatgpt/", "litellm/", "any-llm/", "ollama/", @@ -30,6 +29,8 @@ _DEFAULT_OUTPUT_TOKENS = 8_192 def _lookup_key(model: str) -> str: + if subscription.provider_for_model(model) is not None: + return subscription.litellm_model_name(model) or model for prefix in _STRIPPABLE_PREFIXES: if model.startswith(prefix): return model[len(prefix) :] @@ -46,9 +47,12 @@ def _safe_get_model_info(model: str) -> dict[str, Any] | None: @lru_cache(maxsize=128) def _model_info(model: str) -> dict[str, int]: 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) + # Subscription prefixes are never LiteLLM keys, and a provider-qualified + # ChatGPT lookup may start a synchronous device-login poll: only ask about + # the resolved name. + candidates = ( + (lookup_key,) if subscription.provider_for_model(model) is not None else (model, lookup_key) + ) for candidate in candidates: info = _safe_get_model_info(candidate) if info is not None: diff --git a/tests/test_context_budget.py b/tests/test_context_budget.py index a9a48e0c..c6032f69 100644 --- a/tests/test_context_budget.py +++ b/tests/test_context_budget.py @@ -39,6 +39,17 @@ def test_context_window_chatgpt_prefix_skips_provider_auth( context_budget._model_info.cache_clear() +def test_context_window_grok_prefix_resolves_to_xai() -> None: + # LiteLLM maps xAI models only provider-qualified: neither "grok/grok-4" nor + # bare "grok-4" resolves, so the subscription prefix becomes "xai/". + context_budget._model_info.cache_clear() + try: + assert context_budget.context_window("grok/grok-4") == 256_000 + assert context_budget.output_limit("grok/grok-4") == 256_000 + finally: + context_budget._model_info.cache_clear() + + def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None: context_budget._model_info.cache_clear() diff --git a/tests/test_grok_routing.py b/tests/test_grok_routing.py index 581bb613..4a9c2496 100644 --- a/tests/test_grok_routing.py +++ b/tests/test_grok_routing.py @@ -47,6 +47,17 @@ def test_provider_label_names_the_subscription() -> None: assert subscription.provider_label("openai/gpt-5.4") is None +def test_litellm_model_name_maps_subscription_prefixes() -> None: + # Model metadata (context window, output cap) is keyed "xai/…" for Grok and + # bare for ChatGPT; the routing prefixes themselves are never LiteLLM keys. + assert subscription.litellm_model_name("grok/grok-4") == "xai/grok-4" + assert subscription.litellm_model_name("chatgpt/gpt-5.4") == "gpt-5.4" + # Non-subscription models pass through untouched. + assert subscription.litellm_model_name("xai/grok-4") == "xai/grok-4" + assert subscription.litellm_model_name("openai/gpt-5.4") == "openai/gpt-5.4" + assert subscription.litellm_model_name(None) is None + + def test_run_record_reports_grok_provider(monkeypatch) -> None: # type: ignore[no-untyped-def] settings = mock.MagicMock() settings.llm.model = "grok/grok-4"