diff --git a/index/src/pixelrag_index/pipelines.py b/index/src/pixelrag_index/pipelines.py index f943ec8..0f4800f 100644 --- a/index/src/pixelrag_index/pipelines.py +++ b/index/src/pixelrag_index/pipelines.py @@ -217,7 +217,10 @@ def main(): ) parser.add_argument("--output", "-o", default=None, help="Output directory") parser.add_argument( - "--device", default=None, choices=["cpu", "cuda"], help="Embedding device" + "--device", + default=None, + choices=["auto", "cpu", "mps", "cuda"], + help="Embedding device (default: from config; auto-detects cuda/mps/cpu)", ) parser.add_argument( "--limit", "-n", type=int, default=None, help="Max documents to process" diff --git a/tests/test_cli.py b/tests/test_cli.py index 1c43df2..2d32016 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -32,6 +32,14 @@ def test_pixelrag_umbrella_help(): assert stage in out +def test_index_build_device_choices(): + # CLI must offer every device the embedder supports (auto/mps were missing). + # Assert the argparse-rendered choices token, not loose words in the help prose. + r = _run("pixelrag", "index", "build", "--help") + assert r.returncode == 0 + assert "{auto,cpu,mps,cuda}" in r.stdout + + def test_pixelrag_unknown_stage_errors(): r = _run("pixelrag", "definitely-not-a-stage") assert r.returncode != 0