mirror of
https://github.com/usestrix/strix.git
synced 2026-08-17 01:29:42 +02:00
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.
This commit is contained in:
@@ -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)}"
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user