From 50f218c7bcc7ef5941617c023cbf8ebda51bf73c Mon Sep 17 00:00:00 2001 From: Arvuno Date: Sun, 7 Jun 2026 22:34:58 +0000 Subject: [PATCH 1/2] fix(cli): pre-flight check for magic-prompt API key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When `--magic-prompt` is left at its default (on) and no key is resolvable (via `--magic-prompt-key`, `$MAGIC_PROMPT_API_KEY`, or `$IDEOGRAM_API_KEY`), the script used to fail late — after the Hive prompt-screening call and after at least some model load had started — with a bare "ERROR: ..." line. The new pre-flight check in `main()` calls `parser.error(...)` immediately, before any network or model work. The key-resolution helper (`_resolve_magic_prompt_key`) and the boolean check (`check_magic_prompt_key`) are extracted into a new `cli_preflight.py` module so they can be unit-tested without importing torch or the rest of the model code. The post-load duplicate check in `main()` is removed (the pre-flight already guarantees the key is set when `--magic-prompt` is on). New tests cover: magic-prompt disabled, magic-prompt enabled with flag key, magic-prompt enabled with no key (and both env vars unset), and the resolution precedence (explicit flag, then `$MAGIC_PROMPT_API_KEY`, then `$IDEOGRAM_API_KEY`, with empty env vars treated as missing). --- cli_preflight.py | 43 +++++++++++++++++++++++++++++++++++++ run_inference.py | 29 +++++++++++++++++-------- tests/test_cli_preflight.py | 42 ++++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 9 deletions(-) create mode 100644 cli_preflight.py create mode 100644 tests/test_cli_preflight.py diff --git a/cli_preflight.py b/cli_preflight.py new file mode 100644 index 0000000..2803f0b --- /dev/null +++ b/cli_preflight.py @@ -0,0 +1,43 @@ +"""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) -> bool: + """Return True if the pre-flight check passes. + + The pre-flight check requires that either ``--magic-prompt`` is disabled + or that an API key is resolvable (via ``--magic-prompt-key`` or the + ``MAGIC_PROMPT_API_KEY`` / ``IDEOGRAM_API_KEY`` environment variables). + + When the check fails the caller is expected to call ``parser.error(...)`` + with a user-facing message, so this function is side-effect-free. + """ + if not getattr(args, "magic_prompt", False): + return True + return _resolve_magic_prompt_key(getattr(args, "magic_prompt_key", None)) is not None diff --git a/run_inference.py b/run_inference.py index dc747f1..b079e9b 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 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..e0248b4 --- /dev/null +++ b/tests/test_cli_preflight.py @@ -0,0 +1,42 @@ +"""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 True + + +def test_check_returns_true_when_flag_key_set() -> None: + args = argparse.Namespace(magic_prompt=True, magic_prompt_key="sk-abc") + assert check_magic_prompt_key(args) is True + + +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 False + + +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") + assert _resolve_magic_prompt_key(None) == "sk-ideo" # empty explicit falls through to env + # 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 From e97744eccda2663ce0c889de8babe023e51b3fea Mon Sep 17 00:00:00 2001 From: kartalops Date: Wed, 1 Jul 2026 12:44:16 +0000 Subject: [PATCH 2/2] fix(preflight): populate args.magic_prompt_key from env when pre-flight passes Copilot review flagged that check_magic_prompt_key returned only a boolean while run_inference.py later relied on args.magic_prompt_key being a non-None string. With the old signature, args.magic_prompt_key could be None even when the pre-flight succeeded via $MAGIC_PROMPT_API_KEY or $IDEOGRAM_API_KEY, and the downstream MAGIC_PROMPTS[...]() call would crash on api_key=None. Changes: - check_magic_prompt_key now returns the resolved key string (or None) and writes it back onto args.magic_prompt_key as a convenience for callers. - run_inference.py short-circuits the pre-flight when --magic-prompt is off so we no longer rely on a sentinel-True return. - tests: add coverage for both env-var resolution paths and assert that args.magic_prompt_key is populated in those cases. --- cli_preflight.py | 32 +++++++++++++++++++++----------- run_inference.py | 2 +- tests/test_cli_preflight.py | 34 +++++++++++++++++++++++++++++----- 3 files changed, 51 insertions(+), 17 deletions(-) diff --git a/cli_preflight.py b/cli_preflight.py index 2803f0b..452f67f 100644 --- a/cli_preflight.py +++ b/cli_preflight.py @@ -28,16 +28,26 @@ def _resolve_magic_prompt_key(explicit: str | None) -> str | None: return os.environ.get("IDEOGRAM_API_KEY") or None -def check_magic_prompt_key(args: argparse.Namespace) -> bool: - """Return True if the pre-flight check passes. - - The pre-flight check requires that either ``--magic-prompt`` is disabled - or that an API key is resolvable (via ``--magic-prompt-key`` or the - ``MAGIC_PROMPT_API_KEY`` / ``IDEOGRAM_API_KEY`` environment variables). - - When the check fails the caller is expected to call ``parser.error(...)`` - with a user-facing message, so this function is side-effect-free. +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 True - return _resolve_magic_prompt_key(getattr(args, "magic_prompt_key", None)) is not None + 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 b079e9b..15a2d59 100644 --- a/run_inference.py +++ b/run_inference.py @@ -152,7 +152,7 @@ def main() -> None: # 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 not check_magic_prompt_key(args): + 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 " diff --git a/tests/test_cli_preflight.py b/tests/test_cli_preflight.py index e0248b4..4519896 100644 --- a/tests/test_cli_preflight.py +++ b/tests/test_cli_preflight.py @@ -11,12 +11,14 @@ 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 True + 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") - assert check_magic_prompt_key(args) is True + # 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: @@ -24,7 +26,28 @@ def test_check_returns_false_when_no_key_anywhere() -> 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 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: @@ -35,8 +58,9 @@ def test_resolve_prefers_explicit_then_magic_then_ideogram(monkeypatch) -> None: assert _resolve_magic_prompt_key(None) == "sk-magic" monkeypatch.setenv("IDEOGRAM_API_KEY", "sk-ideo") monkeypatch.delenv("MAGIC_PROMPT_API_KEY") - assert _resolve_magic_prompt_key(None) == "sk-ideo" # empty explicit falls through to env - # But an empty env var is treated as missing + # 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