diff --git a/cli_preflight.py b/cli_preflight.py new file mode 100644 index 0000000..452f67f --- /dev/null +++ b/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 diff --git a/run_inference.py b/run_inference.py index dc747f1..15a2d59 100644 --- a/run_inference.py +++ b/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( diff --git a/tests/test_cli_preflight.py b/tests/test_cli_preflight.py new file mode 100644 index 0000000..4519896 --- /dev/null +++ b/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