Harden provider response handling

mp1/pluto/dispatcher.py

@@ -73,10 +73,62 @@ def _select_nvidia_env_var(model_id: str) -> str:
 def _resolve_nvidia_api_key(model_id: str) -> tuple[str, str]:
     """Resolve model-specific NVIDIA credentials with a global fallback."""
     env_var = _select_nvidia_env_var(model_id)
-    api_key = os.getenv(env_var) or os.getenv("NVIDIA_API_KEY", "")
+    api_key = _clean_api_key(os.getenv(env_var) or os.getenv("NVIDIA_API_KEY", ""))
     return env_var, api_key
 
 
+def _clean_api_key(api_key: str | None) -> str:
+    """Normalize secret values pasted with optional auth prefixes or whitespace."""
+    cleaned = str(api_key or "").strip().strip('"').strip("'")
+    if cleaned.lower().startswith("bearer "):
+        cleaned = cleaned[7:].strip()
+    return cleaned
+
+
+def _normalize_model_text(value: Any) -> str:
+    """Turn provider message content into a safe string."""
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    if isinstance(value, list):
+        parts: list[str] = []
+        for item in value:
+            if isinstance(item, str):
+                parts.append(item)
+            elif isinstance(item, dict):
+                text = item.get("text") or item.get("content") or item.get("value")
+                if text is not None:
+                    parts.append(str(text))
+        return "\n".join(part for part in parts if part)
+    return str(value)
+
+
+def _chat_response_text(data: dict[str, Any], provider: str) -> str:
+    """Extract assistant text from OpenAI-compatible chat response JSON."""
+    choices = data.get("choices")
+    if not choices:
+        raise RuntimeError(f"{provider} response did not include choices")
+
+    message = choices[0].get("message") if isinstance(choices[0], dict) else None
+    if not isinstance(message, dict):
+        raise RuntimeError(f"{provider} response did not include a message")
+
+    text = _normalize_model_text(message.get("content"))
+    if not text:
+        text = _normalize_model_text(message.get("reasoning_content"))
+    if not text:
+        text = _normalize_model_text(message.get("reasoning"))
+    return text
+
+
+def _require_model_text(text: str, provider: str) -> str:
+    normalized = _normalize_model_text(text)
+    if not normalized.strip():
+        raise RuntimeError(f"{provider} returned an empty response")
+    return normalized
+
+
 def dispatch(
     mode_name: str,
     prompt: str,
@@ -84,6 +136,7 @@ def dispatch(
     images: list[bytes] | None = None,
 ) -> str:
     """Route prompt to the configured provider for the selected mode."""
+    prompt = str(prompt or "")
     cfg = get_mode(mode_name)
 
     if tracer:
@@ -102,11 +155,11 @@ def dispatch(
 
     try:
         if cfg.provider == "nvidia":
-            text = _call_nvidia(cfg, prompt)
+            text = _require_model_text(_call_nvidia(cfg, prompt), "NVIDIA")
         elif cfg.provider == "groq":
-            text = _call_groq(cfg, prompt)
+            text = _require_model_text(_call_groq(cfg, prompt), "Groq")
         elif cfg.provider == "mistral":
-            text = _call_mistral(cfg, prompt)
+            text = _require_model_text(_call_mistral(cfg, prompt), "Mistral")
         else:
             raise ValueError(f"Unknown provider: {cfg.provider}")
 
@@ -114,9 +167,9 @@ def dispatch(
         print(f"  [WARNING] {cfg.provider} failed: {e}")
 
         if cfg.provider == "groq":
-            text = _call_best_fallback(cfg, prompt, allow_groq=False)
+            text = _require_model_text(_call_best_fallback(cfg, prompt, allow_groq=False), "fallback")
         else:
-            text = _call_best_fallback(cfg, prompt, allow_groq=True)
+            text = _require_model_text(_call_best_fallback(cfg, prompt, allow_groq=True), "fallback")
 
     elapsed = time.perf_counter() - t0
     if tracer:
@@ -133,7 +186,7 @@ def dispatch(
 
 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()
+    groq_key = _clean_api_key(os.getenv("GROQ_API_KEY", ""))
     if allow_groq and groq_key:
         print("  [FALLBACK] Trying Groq...")
         fb = ModeConfig(
@@ -146,7 +199,7 @@ def _call_best_fallback(cfg: ModeConfig, prompt: str, allow_groq: bool) -> str:
         )
         return _call_groq(fb, prompt)
 
-    mistral_key = os.getenv("MISTRAL_API_KEY", "").strip()
+    mistral_key = _clean_api_key(os.getenv("MISTRAL_API_KEY", ""))
     if mistral_key and not _looks_like_nvidia_key(mistral_key):
         print("  [FALLBACK] Trying Mistral...")
         fb = ModeConfig(
@@ -159,7 +212,7 @@ def _call_best_fallback(cfg: ModeConfig, prompt: str, allow_groq: bool) -> str:
         )
         return _call_mistral(fb, prompt)
 
-    nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip()
+    nvidia_key = _clean_api_key(os.getenv("NVIDIA_API_KEY", ""))
     if nvidia_key:
         print("  [FALLBACK] Trying NVIDIA-hosted Mistral...")
         return _call_nvidia_hosted_mistral(cfg, prompt, nvidia_key)
@@ -170,7 +223,7 @@ def _call_best_fallback(cfg: ModeConfig, prompt: str, allow_groq: bool) -> str:
 
 
 def _looks_like_nvidia_key(api_key: str) -> bool:
-    return api_key.strip().startswith("nvapi-")
+    return _clean_api_key(api_key).startswith("nvapi-")
 
 
 def _call_groq(cfg: ModeConfig, prompt: str) -> str:
@@ -188,7 +241,7 @@ def _call_groq(cfg: ModeConfig, prompt: str) -> str:
                     temperature=cfg.temperature,
                     max_tokens=cfg.max_tokens,
                 )
-                return resp.choices[0].message.content or ""
+                return _normalize_model_text(resp.choices[0].message.content)
 
             except Exception as e:
                 err = str(e).lower()
@@ -204,9 +257,11 @@ def _call_groq(cfg: ModeConfig, prompt: str) -> str:
 
 def _call_mistral(cfg: ModeConfig, prompt: str) -> str:
     """Call Mistral AI REST API with retry on rate-limit."""
-    api_key = os.getenv("MISTRAL_API_KEY", "")
+    api_key = _clean_api_key(os.getenv("MISTRAL_API_KEY", ""))
     if not api_key:
         raise ValueError("MISTRAL_API_KEY not set")
+    if _looks_like_nvidia_key(api_key):
+        raise ValueError("MISTRAL_API_KEY contains an NVIDIA key; move it to NVIDIA_API_KEY")
 
     headers = {
         "Authorization": f"Bearer {api_key}",
@@ -228,7 +283,7 @@ def _call_mistral(cfg: ModeConfig, prompt: str) -> str:
                 timeout=60,
             )
             if response.status_code == 200:
-                return response.json()["choices"][0]["message"]["content"]
+                return _chat_response_text(response.json(), "Mistral")
             if response.status_code == 429:
                 delay = 10 * (attempt + 1)
                 print(f"  [RETRY] Mistral rate-limit — waiting {delay}s (attempt {attempt + 1})")
@@ -248,7 +303,10 @@ def _call_nvidia(cfg: ModeConfig, prompt: str) -> str:
     env_var, api_key = _resolve_nvidia_api_key(cfg.model_id)
     if not api_key:
         raise ValueError(f"{env_var} or NVIDIA_API_KEY not set")
-    has_fallback = bool(os.getenv("GROQ_API_KEY", "") or os.getenv("MISTRAL_API_KEY", ""))
+    has_fallback = bool(
+        _clean_api_key(os.getenv("GROQ_API_KEY", ""))
+        or _clean_api_key(os.getenv("MISTRAL_API_KEY", ""))
+    )
     max_retries = NVIDIA_MAX_RETRIES_WITH_FALLBACK if has_fallback else NVIDIA_MAX_RETRIES_DEFAULT
     request_timeout = NVIDIA_TIMEOUT_WITH_FALLBACK if has_fallback else NVIDIA_TIMEOUT_DEFAULT
 
@@ -279,8 +337,7 @@ def _call_nvidia(cfg: ModeConfig, prompt: str) -> str:
                 timeout=request_timeout,
             )
             if response.status_code == 200:
-                data = response.json()
-                return data["choices"][0]["message"]["content"]
+                return _chat_response_text(response.json(), "NVIDIA")
             raise Exception(f"NVIDIA {response.status_code}: {response.text[:300]}")
         except Exception as e:
             err = str(e).lower()
@@ -318,6 +375,7 @@ def _call_nvidia(cfg: ModeConfig, prompt: str) -> str:
 
 def _call_nvidia_hosted_mistral(cfg: ModeConfig, prompt: str, api_key: str) -> str:
     """Call Mistral served through NVIDIA NIM using NVIDIA credentials."""
+    api_key = _clean_api_key(api_key)
     payload = {
         "model": NVIDIA_MISTRAL_FALLBACK_MODEL,
         "messages": [{"role": "user", "content": prompt}],
@@ -335,8 +393,7 @@ def _call_nvidia_hosted_mistral(cfg: ModeConfig, prompt: str, api_key: str) -> s
     )
     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"]
+    return _chat_response_text(response.json(), "NVIDIA-hosted Mistral")
 
 
 def rerank(query: str, passages: list[str]) -> list[float]:
@@ -346,7 +403,7 @@ def rerank(query: str, passages: list[str]) -> list[float]:
     Returns scores in the same order as the input passages and falls back to
     uniform scores if the reranker is unavailable.
     """
-    api_key = os.getenv("NVIDIA_API_KEY_RERANK") or os.getenv("NVIDIA_API_KEY", "")
+    api_key = _clean_api_key(os.getenv("NVIDIA_API_KEY_RERANK") or os.getenv("NVIDIA_API_KEY", ""))
     if not api_key or not passages:
         return [1.0] * len(passages)
 

mp1/pluto/modes.py

@@ -23,6 +23,17 @@ from dotenv import load_dotenv
 load_dotenv()
 
 
+def _clean_api_key(api_key: str | None) -> str:
+    cleaned = str(api_key or "").strip().strip('"').strip("'")
+    if cleaned.lower().startswith("bearer "):
+        cleaned = cleaned[7:].strip()
+    return cleaned
+
+
+def _looks_like_nvidia_key(api_key: str) -> bool:
+    return _clean_api_key(api_key).startswith("nvapi-")
+
+
 @dataclass(frozen=True)
 class ModeConfig:
     """Concrete model configuration for a single processing mode."""
@@ -66,9 +77,9 @@ def _build_registry() -> dict[str, ModeConfig]:
         "NVIDIA_API_KEY_VL", "NVIDIA_API_KEY_EMBED", "NVIDIA_API_KEY_RERANK",
         "NVIDIA_API_KEY_ULTRA"
     ]
-    nvidia_ready = any(os.getenv(k) for k in nvidia_keys)
-    groq_key = os.getenv("GROQ_API_KEY", "").strip()
-    mistral_key = os.getenv("MISTRAL_API_KEY", "").strip()
+    nvidia_ready = any(_clean_api_key(os.getenv(k)) for k in nvidia_keys)
+    groq_key = _clean_api_key(os.getenv("GROQ_API_KEY", ""))
+    mistral_key = _clean_api_key(os.getenv("MISTRAL_API_KEY", ""))
 
     if nvidia_ready:
         return {
@@ -158,7 +169,7 @@ def _build_registry() -> dict[str, ModeConfig]:
                 provider="groq",
             ),
         }
-    if mistral_key:
+    if mistral_key and not _looks_like_nvidia_key(mistral_key):
         return _build_mistral_registry()
     return _build_unconfigured_registry()