Use NVIDIA hosted Mistral fallback

mp1/pluto/dispatcher.py

@@ -23,6 +23,7 @@ from pluto.tracer import Tracer
 load_dotenv()
 
 NVIDIA_CHAT_URL = "https://integrate.api.nvidia.com/v1/chat/completions"
+NVIDIA_MISTRAL_FALLBACK_MODEL = "mistralai/mistral-medium-3.5-128b"
 NVIDIA_RERANK_URLS = (
     "https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-nemotron-rerank-1b-v2/reranking",
     "https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking",
@@ -113,42 +114,9 @@ def dispatch(
         print(f"  [WARNING] {cfg.provider} failed: {e}")
 
         if cfg.provider == "groq":
-            print("  [FALLBACK] Trying Mistral...")
-            fb = ModeConfig(
-                mode_name=cfg.mode_name,
-                model_id="mistral-small-latest",
-                temperature=cfg.temperature,
-                max_tokens=cfg.max_tokens,
-                compute_profile="fallback",
-                provider="mistral",
-            )
-            text = _call_mistral(fb, prompt)
+            text = _call_best_fallback(cfg, prompt, allow_groq=False)
         else:
-            groq_key = os.getenv("GROQ_API_KEY", "")
-            if groq_key:
-                print("  [FALLBACK] Trying Groq...")
-                fb = ModeConfig(
-                    mode_name=cfg.mode_name,
-                    model_id="llama-3.1-8b-instant",
-                    temperature=cfg.temperature,
-                    max_tokens=cfg.max_tokens,
-                    compute_profile="fallback",
-                    provider="groq",
-                )
-                text = _call_groq(fb, prompt)
-            elif os.getenv("MISTRAL_API_KEY", ""):
-                print("  [FALLBACK] Trying Mistral...")
-                fb = ModeConfig(
-                    mode_name=cfg.mode_name,
-                    model_id="mistral-small-latest",
-                    temperature=cfg.temperature,
-                    max_tokens=cfg.max_tokens,
-                    compute_profile="fallback",
-                    provider="mistral",
-                )
-                text = _call_mistral(fb, prompt)
-            else:
-                raise
+            text = _call_best_fallback(cfg, prompt, allow_groq=True)
 
     elapsed = time.perf_counter() - t0
     if tracer:
@@ -163,6 +131,48 @@ def dispatch(
     return text
 
 
+def _call_best_fallback(cfg: ModeConfig, prompt: str, allow_groq: bool) -> str:
+    """Try the configured fallback providers without confusing NVIDIA and Mistral keys."""
+    groq_key = os.getenv("GROQ_API_KEY", "").strip()
+    if allow_groq and groq_key:
+        print("  [FALLBACK] Trying Groq...")
+        fb = ModeConfig(
+            mode_name=cfg.mode_name,
+            model_id="llama-3.1-8b-instant",
+            temperature=cfg.temperature,
+            max_tokens=cfg.max_tokens,
+            compute_profile="fallback",
+            provider="groq",
+        )
+        return _call_groq(fb, prompt)
+
+    mistral_key = os.getenv("MISTRAL_API_KEY", "").strip()
+    if mistral_key and not _looks_like_nvidia_key(mistral_key):
+        print("  [FALLBACK] Trying Mistral...")
+        fb = ModeConfig(
+            mode_name=cfg.mode_name,
+            model_id="mistral-small-latest",
+            temperature=cfg.temperature,
+            max_tokens=cfg.max_tokens,
+            compute_profile="fallback",
+            provider="mistral",
+        )
+        return _call_mistral(fb, prompt)
+
+    nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip()
+    if nvidia_key:
+        print("  [FALLBACK] Trying NVIDIA-hosted Mistral...")
+        return _call_nvidia_hosted_mistral(cfg, prompt, nvidia_key)
+
+    if mistral_key:
+        raise ValueError("MISTRAL_API_KEY appears to be an NVIDIA nvapi key. Put it in NVIDIA_API_KEY instead.")
+    raise ValueError("No fallback provider is configured.")
+
+
+def _looks_like_nvidia_key(api_key: str) -> bool:
+    return api_key.strip().startswith("nvapi-")
+
+
 def _call_groq(cfg: ModeConfig, prompt: str) -> str:
     """Call Groq with smart retry and rate-limit backoff."""
     groq_max_retries = 8
@@ -306,6 +316,29 @@ def _call_nvidia(cfg: ModeConfig, prompt: str) -> str:
     return ""
 
 
+def _call_nvidia_hosted_mistral(cfg: ModeConfig, prompt: str, api_key: str) -> str:
+    """Call Mistral served through NVIDIA NIM using NVIDIA credentials."""
+    payload = {
+        "model": NVIDIA_MISTRAL_FALLBACK_MODEL,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": cfg.temperature,
+        "max_tokens": cfg.max_tokens,
+    }
+    response = requests.post(
+        NVIDIA_CHAT_URL,
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        },
+        json=payload,
+        timeout=NVIDIA_TIMEOUT_WITH_FALLBACK,
+    )
+    if response.status_code != 200:
+        raise RuntimeError(f"NVIDIA-hosted Mistral {response.status_code}: {response.text[:300]}")
+    data = response.json()
+    return data["choices"][0]["message"]["content"]
+
+
 def rerank(query: str, passages: list[str]) -> list[float]:
     """
     Score each passage for relevance to the query using the NVIDIA reranker.