fix(cli): pre-flight check for magic-prompt API key#23
Conversation
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).
There was a problem hiding this comment.
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.pywith key-resolution + preflight validation logic. - Refactors
run_inference.pyto 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.
| 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." | ||
| ) |
| # 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 |
| 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 |
|
Summary
main()that callsparser.error(...)immediately when--magic-promptis on but no key is resolvable (flag,$MAGIC_PROMPT_API_KEY, or$IDEOGRAM_API_KEY).cli_preflight.pymodule so they can be unit-tested without importing torch or the rest of the model code.main()is removed (the pre-flight already guarantees the key is set).check_magic_prompt_keyand_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.python3 -c "import ast; ast.parse(open('run_inference.py').read())"→ OK.