fix(cli): pre-flight check for magic-prompt API key

cli_preflight.py

@@ -0,0 +1,53 @@
+"""Magic-prompt pre-flight check, isolated so it can be tested without torch
+or the ideogram4 model code.
+
+Lives at the repo root (next to run_inference.py) so the helper is
+importable by both the CLI and the tests.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+
+def _resolve_magic_prompt_key(explicit: str | None) -> str | None:
+  """Resolve the magic-prompt API key.
+
+  Precedence:
+    1. ``explicit`` (the ``--magic-prompt-key`` value).
+    2. ``$MAGIC_PROMPT_API_KEY``.
+    3. ``$IDEOGRAM_API_KEY``.
+  Returns the first non-empty match, or None if all are missing.
+  """
+  if explicit:
+    return explicit
+  env_key = os.environ.get("MAGIC_PROMPT_API_KEY")
+  if env_key:
+    return env_key
+  return os.environ.get("IDEOGRAM_API_KEY") or None
+
+
+def check_magic_prompt_key(args: argparse.Namespace) -> str | None:
+  """Resolve and return the magic-prompt API key, or None if not available.
+
+  Returns:
+    - ``None`` if ``--magic-prompt`` is disabled (caller should skip the
+      pre-flight entirely).
+    - A non-empty string with the resolved key (``--magic-prompt-key``,
+      ``$MAGIC_PROMPT_API_KEY``, or ``$IDEOGRAM_API_KEY``) when pre-flight
+      passes.
+    - ``None`` if pre-flight fails (caller is expected to call
+      ``parser.error(...)`` with a user-facing message).
+
+  As a convenience for callers that re-use ``args.magic_prompt_key`` after
+  the pre-flight, the resolved key is also written back onto ``args``.
+  """
+  if not getattr(args, "magic_prompt", False):
+    return None
+  resolved = _resolve_magic_prompt_key(getattr(args, "magic_prompt_key", None))
+  if resolved is not None:
+    # Populate args so downstream code can rely on args.magic_prompt_key
+    # being set whenever the pre-flight passes.
+    args.magic_prompt_key = resolved
+  return resolved

run_inference.py

@@ -19,6 +19,8 @@
   moderate_prompt,
 )
 
+from cli_preflight import check_magic_prompt_key
+
 QUANTIZATION_REPOS = {
   "nf4": "ideogram-ai/ideogram-4-nf4",
   "fp8": "ideogram-ai/ideogram-4-fp8",
@@ -45,7 +47,7 @@ def _print_flags(label: str, flags: list[tuple[str, float]]) -> None:
     print(f"  {name}: {score:.3f}", file=sys.stderr)
 
 
-def main() -> None:
+def build_argparser() -> argparse.ArgumentParser:
   parser = argparse.ArgumentParser()
   parser.add_argument("--prompt", required=True)
   parser.add_argument("--output", default="output.png")
@@ -140,8 +142,23 @@ def main() -> None:
       "create a Visual Content Moderation project."
     ),
   )
+  return parser
+
+
+def main() -> None:
+  parser = build_argparser()
   args = parser.parse_args()
 
+  # Pre-flight: if magic-prompt is on but no key is set (via flag or env),
+  # fail fast before any network or model-load work. Saves a real Hive call
+  # and any partial model load in the common "I forgot the key" case.
+  if args.magic_prompt and not check_magic_prompt_key(args):
+    parser.error(
+      "--magic-prompt is on but no API key was set. Provide one via "
+      "--magic-prompt-key, $MAGIC_PROMPT_API_KEY, or $IDEOGRAM_API_KEY; or "
+      "pass --no-magic-prompt to disable expansion."
+    )
+
   if args.hive_text_key:
     flags = moderate_prompt(args.prompt, args.hive_text_key)
     if flags:
@@ -158,14 +175,8 @@ def main() -> None:
 
   prompt = args.prompt
   if args.magic_prompt:
-    if not args.magic_prompt_key:
-      print(
-        "ERROR: magic prompt is enabled but no API key was found. Set "
-        "MAGIC_PROMPT_API_KEY, pass --magic-prompt-key, or disable expansion "
-        "with --no-magic-prompt.",
-        file=sys.stderr,
-      )
-      sys.exit(2)
+    # The pre-flight check above guarantees args.magic_prompt_key is set
+    # when args.magic_prompt is true.
     aspect_ratio = aspect_ratio_from_size(args.width, args.height)
     magic = MAGIC_PROMPTS[args.magic_prompt_model](api_key=args.magic_prompt_key)  # type: ignore[call-arg]
     print(

tests/test_cli_preflight.py

@@ -0,0 +1,66 @@
+"""Tests for the magic-prompt pre-flight check (no torch, no ideogram4)."""
+
+from __future__ import annotations
+
+import argparse
+
+import pytest
+
+from cli_preflight import check_magic_prompt_key, _resolve_magic_prompt_key
+
+
+def test_check_returns_true_when_magic_prompt_disabled() -> None:
+  args = argparse.Namespace(magic_prompt=False, magic_prompt_key=None)
+  assert check_magic_prompt_key(args) is None
+
+
+def test_check_returns_true_when_flag_key_set() -> None:
+  args = argparse.Namespace(magic_prompt=True, magic_prompt_key="sk-abc")
+  # Pre-flight passes AND the resolved key is written back onto args.
+  assert check_magic_prompt_key(args) == "sk-abc"
+  assert args.magic_prompt_key == "sk-abc"
+
+
+def test_check_returns_false_when_no_key_anywhere() -> None:
+  args = argparse.Namespace(magic_prompt=True, magic_prompt_key=None)
+  with pytest.MonkeyPatch.context() as mp:
+    mp.delenv("MAGIC_PROMPT_API_KEY", raising=False)
+    mp.delenv("IDEOGRAM_API_KEY", raising=False)
+    assert check_magic_prompt_key(args) is None
+
+
+def test_check_resolves_via_magic_prompt_api_key_env(monkeypatch) -> None:
+  """When no flag key is set but $MAGIC_PROMPT_API_KEY is, pre-flight
+  passes and args.magic_prompt_key is populated from the env."""
+  monkeypatch.delenv("IDEOGRAM_API_KEY", raising=False)
+  monkeypatch.setenv("MAGIC_PROMPT_API_KEY", "sk-from-env")
+  args = argparse.Namespace(magic_prompt=True, magic_prompt_key=None)
+  assert check_magic_prompt_key(args) == "sk-from-env"
+  assert args.magic_prompt_key == "sk-from-env"
+
+
+def test_check_resolves_via_ideogram_api_key_env(monkeypatch) -> None:
+  """When no flag key and no $MAGIC_PROMPT_API_KEY, but $IDEOGRAM_API_KEY
+  is set, pre-flight passes and args.magic_prompt_key is populated from
+  the env."""
+  monkeypatch.delenv("MAGIC_PROMPT_API_KEY", raising=False)
+  monkeypatch.setenv("IDEOGRAM_API_KEY", "sk-ideo")
+  args = argparse.Namespace(magic_prompt=True, magic_prompt_key=None)
+  assert check_magic_prompt_key(args) == "sk-ideo"
+  assert args.magic_prompt_key == "sk-ideo"
+
+
+def test_resolve_prefers_explicit_then_magic_then_ideogram(monkeypatch) -> None:
+  monkeypatch.delenv("MAGIC_PROMPT_API_KEY", raising=False)
+  monkeypatch.delenv("IDEOGRAM_API_KEY", raising=False)
+  assert _resolve_magic_prompt_key("sk-explicit") == "sk-explicit"
+  monkeypatch.setenv("MAGIC_PROMPT_API_KEY", "sk-magic")
+  assert _resolve_magic_prompt_key(None) == "sk-magic"
+  monkeypatch.setenv("IDEOGRAM_API_KEY", "sk-ideo")
+  monkeypatch.delenv("MAGIC_PROMPT_API_KEY")
+  # An explicit=None (not just "") falls through to env lookup.
+  assert _resolve_magic_prompt_key(None) == "sk-ideo"
+  # But an empty env var is treated as missing.
+  monkeypatch.setenv("MAGIC_PROMPT_API_KEY", "")
+  monkeypatch.delenv("IDEOGRAM_API_KEY")
+  assert _resolve_magic_prompt_key(None) is None