Skip to content

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

Open
Arvuno wants to merge 1 commit into
ideogram-oss:mainfrom
Arvuno:fix/cli-magic-prompt-preflight
Open

fix(cli): pre-flight check for magic-prompt API key#23
Arvuno wants to merge 1 commit into
ideogram-oss:mainfrom
Arvuno:fix/cli-magic-prompt-preflight

Conversation

@Arvuno

@Arvuno Arvuno commented Jun 7, 2026

Copy link
Copy Markdown

Summary

  • Adds a pre-flight check in main() that calls parser.error(...) immediately when --magic-prompt is on but no key is resolvable (flag, $MAGIC_PROMPT_API_KEY, or $IDEOGRAM_API_KEY).
  • The key-resolution helper and boolean check 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).
  • 4 new unit tests for check_magic_prompt_key and _resolve_magic_prompt_key.

Why

The previous behaviour was a late failure: a user who forgot the API key would first trigger a Hive text-moderation call (if configured) and then the model load would start before the script bailed with ERROR: .... The pre-flight check saves a real network call and any partial model load in the common "I forgot the key" case, and the failure message is now formatted as a proper argparse error (exit code 2 with a usage hint).

Testing

  • pytest tests/test_cli_preflight.py → 4 passed.
  • The helper module is torch-free; tests run on a plain Python interpreter.
  • python3 -c "import ast; ast.parse(open('run_inference.py').read())" → OK.

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).
Copilot AI review requested due to automatic review settings June 7, 2026 22:35

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a lightweight, testable “magic prompt” API-key preflight check and wires it into the CLI to fail fast before heavier work.

Changes:

  • Introduces cli_preflight.py with key-resolution + preflight validation logic.
  • Refactors run_inference.py to use a shared argparser builder and run the preflight check early.
  • Adds pytest coverage for the preflight helper logic.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File Description
tests/test_cli_preflight.py New unit tests for key resolution and preflight behavior.
run_inference.py Adds build_argparser() and enforces a preflight check before inference.
cli_preflight.py New module encapsulating key resolution + preflight check.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread run_inference.py
Comment on lines +155 to +160
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."
)
Comment thread run_inference.py
Comment on lines +178 to 181
# 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]
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
Comment on lines +22 to +27
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants