Add --num-prompts option to CLI and tests (#12)

* feature: add --num-prompts option to limit number of prompts used in CLI and tests

* docs: update README and CHANGELOG to document --num-prompts CLI option

* refactor: centralize prompt loading and limiting logic in _load_prompts helper

Author: NullPointerDepressiveDisorder

CHANGELOG.md

@@ -1,5 +1,6 @@
-# 0.2.0 (2026-04-02)
+# 0.2.3 (2026-04-07)
 
+- Added `--num-prompts` option to all CLI subcommands to limit the number of prompts used in a task.
 - Added global `--max-tokens` flag (defaults to 1024) to the main CLI.
 - Increased default `max_tokens` for all prompts from 256 to 1024.
 

README.md

@@ -68,7 +68,7 @@ pip install "infer-check[mlx]"
 
 ### Quantization sweep
 
-Compare pre-quantized models against a baseline. Each model is a separate HuggingFace repo. Use `--max-tokens` to control generation length (defaults to 1024).
+Compare pre-quantized models against a baseline. Each model is a separate HuggingFace repo. Use `--max-tokens` to control generation length (defaults to 1024) and `--num-prompts` to limit the number of prompts used.
 
 ```
 infer-check sweep \
@@ -78,6 +78,7 @@ infer-check sweep \
   --backend mlx-lm \
   --prompts reasoning \
   --max-tokens 512 \
+  --num-prompts 10 \
   --output ./results/sweep/
 ```
 

pyproject.toml

@@ -76,7 +76,8 @@ select = ["E", "F", "I", "N", "W", "UP", "B", "SIM"]
 "html.py" = ["E501"]
 
 [tool.bandit]
-targets= ["src"]
+targets = ["src"]
+exclude_dirs = ["tests"]
 
 [tool.mypy]
 python_version = "3.11"

src/infer_check/cli.py

@@ -27,6 +27,12 @@ def common_options(f: F) -> F:
             type=click.IntRange(min=1, clamp=True),
             help="Override default max tokens for generation.",
         ),
+        click.option(
+            "--num-prompts",
+            default=None,
+            type=click.IntRange(min=1, clamp=True),
+            help="Limit number of prompts to use.",
+        ),
     ]
     for option in reversed(options):
         f = option(f)
@@ -43,6 +49,30 @@ def _resolve_prompts(prompts: str) -> Path:
         raise click.BadParameter(str(exc)) from exc
 
 
+def _load_prompts(ctx: click.Context, prompts: str, max_tokens: int | None, num_prompts: int | None) -> list[Any]:
+    """Load and prepare a prompt suite with overrides applied."""
+    from infer_check.suites.loader import load_suite
+
+    # Update defaults from subcommand if provided
+    if max_tokens is not None:
+        ctx.obj["max_tokens"] = max_tokens
+    if num_prompts is not None:
+        ctx.obj["num_prompts"] = num_prompts
+
+    prompt_list = load_suite(_resolve_prompts(prompts))
+
+    # Apply num_prompts limit
+    if ctx.obj["num_prompts"] is not None:
+        prompt_list = prompt_list[: ctx.obj["num_prompts"]]
+
+    # Apply global max_tokens only if not explicitly set in the prompt JSONL
+    for p in prompt_list:
+        if "max_tokens" not in p.model_fields_set:
+            p.max_tokens = ctx.obj["max_tokens"]
+
+    return prompt_list
+
+
 @click.group()
 @click.version_option(package_name="infer-check")
 @click.option(
@@ -51,11 +81,18 @@ def _resolve_prompts(prompts: str) -> Path:
     show_default=True,
     help="Default max tokens for generation (applies to all prompts unless they specify their own).",
 )
+@click.option(
+    "--num-prompts",
+    default=None,
+    type=click.IntRange(min=1, clamp=True),
+    help="Limit the number of prompts to use from a suite; if omitted, all prompts are used.",
+)
 @click.pass_context
-def main(ctx: click.Context, max_tokens: int) -> None:
+def main(ctx: click.Context, max_tokens: int, num_prompts: int | None) -> None:
     """infer-check: correctness and reliability testing for LLM inference engines."""
     ctx.ensure_object(dict)
     ctx.obj["max_tokens"] = max_tokens
+    ctx.obj["num_prompts"] = num_prompts
 
 
 # ---------------------------------------------------------------------------
@@ -103,6 +140,7 @@ def sweep(
     baseline: str | None,
     base_url: str | None,
     max_tokens: int | None,
+    num_prompts: int | None,
 ) -> None:
     """Run a quantization sweep: compare pre-quantized models against a baseline.
 
@@ -119,11 +157,8 @@ def sweep(
     """
     from infer_check.backends.base import get_backend_for_model
     from infer_check.runner import TestRunner
-    from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
-    if max_tokens is not None:
-        ctx.obj["max_tokens"] = max_tokens
+    prompt_list = _load_prompts(ctx, prompts, max_tokens, num_prompts)
 
     # Parse label=model_path pairs
     model_map: dict[str, str] = {}
@@ -152,11 +187,6 @@ def sweep(
     for label, path in model_map.items():
         tag = " (baseline)" if label == baseline_label else ""
         console.print(f"  {label}: {path}{tag}")
-    prompt_list = load_suite(_resolve_prompts(prompts))
-    # Apply global max_tokens only if not explicitly set in the prompt JSONL
-    for p in prompt_list:
-        if "max_tokens" not in p.model_fields_set:
-            p.max_tokens = ctx.obj["max_tokens"]
 
     # Build a separate backend for each model
     backend_map: dict[str, Any] = {}
@@ -308,6 +338,7 @@ def compare(
     label_b: str | None,
     report: bool,
     max_tokens: int | None,
+    num_prompts: int | None,
 ) -> None:
     """Compare two quantizations of the same model.
 
@@ -335,11 +366,8 @@ def compare(
     # ── Resolve both model specs ─────────────────────────────────────
     from infer_check.resolve import resolve_model
     from infer_check.runner import TestRunner
-    from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
-    if max_tokens is not None:
-        ctx.obj["max_tokens"] = max_tokens
+    prompt_list = _load_prompts(ctx, prompts, max_tokens, num_prompts)
 
     resolved_a = resolve_model(model_a, base_url=base_url, label=label_a)
     resolved_b = resolve_model(model_b, base_url=base_url, label=label_b)
@@ -350,12 +378,6 @@ def compare(
         f"vs B={resolved_b.label} ({resolved_b.backend})"
     )
 
-    prompt_list = load_suite(_resolve_prompts(prompts))
-    # Apply global max_tokens only if not explicitly set in the prompt JSONL
-    for p in prompt_list:
-        if "max_tokens" not in p.model_fields_set:
-            p.max_tokens = ctx.obj["max_tokens"]
-
     console.print(f"  prompts: {len(prompt_list)} from '{prompts}'")
 
     # ── Build backends ───────────────────────────────────────────────
@@ -573,15 +595,13 @@ def diff(
     base_urls: str | None,
     chat: bool,
     max_tokens: int | None,
+    num_prompts: int | None,
 ) -> None:
     """Compare outputs across different backends for the same model and prompts."""
     from infer_check.backends.base import BackendConfig, get_backend
     from infer_check.runner import TestRunner
-    from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
-    if max_tokens is not None:
-        ctx.obj["max_tokens"] = max_tokens
+    prompt_list = _load_prompts(ctx, prompts, max_tokens, num_prompts)
 
     backend_names = [b.strip() for b in backends.split(",") if b.strip()]
     url_list: list[str | None] = [u.strip() for u in base_urls.split(",")] if base_urls else [None] * len(backend_names)
@@ -591,12 +611,6 @@ def diff(
 
     console.print(f"[bold cyan]diff[/bold cyan] model={model} backends={backend_names} quant={quant}")
 
-    prompt_list = load_suite(_resolve_prompts(prompts))
-    # Apply global max_tokens only if not explicitly set in the prompt JSONL
-    for p in prompt_list:
-        if "max_tokens" not in p.model_fields_set:
-            p.max_tokens = ctx.obj["max_tokens"]
-
     backend_instances = []
     for name, url in zip(backend_names, url_list, strict=True):
         config = BackendConfig(
@@ -693,15 +707,13 @@ def stress(
     concurrency: str,
     base_url: str | None,
     max_tokens: int | None,
+    num_prompts: int | None,
 ) -> None:
     """Stress-test a backend with varying concurrency levels."""
     from infer_check.backends.base import get_backend_for_model
     from infer_check.runner import TestRunner
-    from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
-    if max_tokens is not None:
-        ctx.obj["max_tokens"] = max_tokens
+    prompt_list = _load_prompts(ctx, prompts, max_tokens, num_prompts)
 
     concurrency_levels = [int(c.strip()) for c in concurrency.split(",") if c.strip()]
 
@@ -715,12 +727,6 @@ def stress(
         f"[bold cyan]stress[/bold cyan] model={model} backend={backend_instance.name} concurrency={concurrency_levels}"
     )
 
-    prompt_list = load_suite(_resolve_prompts(prompts))
-    # Apply global max_tokens only if not explicitly set in the prompt JSONL
-    for p in prompt_list:
-        if "max_tokens" not in p.model_fields_set:
-            p.max_tokens = ctx.obj["max_tokens"]
-
     runner = TestRunner()
     stress_results = asyncio.run(
         runner.stress(
@@ -790,15 +796,13 @@ def determinism(
     runs: int,
     base_url: str | None,
     max_tokens: int | None,
+    num_prompts: int | None,
 ) -> None:
     """Test whether a backend produces identical outputs across repeated runs at temperature=0."""
     from infer_check.backends.base import get_backend_for_model
     from infer_check.runner import TestRunner
-    from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
-    if max_tokens is not None:
-        ctx.obj["max_tokens"] = max_tokens
+    prompt_list = _load_prompts(ctx, prompts, max_tokens, num_prompts)
 
     backend_instance = get_backend_for_model(
         model_str=model,
@@ -808,12 +812,6 @@ def determinism(
 
     console.print(f"[bold cyan]determinism[/bold cyan] model={model} backend={backend_instance.name} runs={runs}")
 
-    prompt_list = load_suite(_resolve_prompts(prompts))
-    # Apply global max_tokens only if not explicitly set in the prompt JSONL
-    for p in prompt_list:
-        if "max_tokens" not in p.model_fields_set:
-            p.max_tokens = ctx.obj["max_tokens"]
-
     runner = TestRunner()
     det_results = asyncio.run(
         runner.determinism(

tests/unit/test_cli_num_prompts.py

@@ -0,0 +1,76 @@
+import json
+from datetime import UTC, datetime
+from pathlib import Path
+from unittest.mock import patch
+
+from click.testing import CliRunner
+
+from infer_check.cli import main
+from infer_check.runner import TestRunner
+from infer_check.types import SweepResult
+
+
+def test_num_prompts_limit() -> None:
+    """Test that --num-prompts correctly limits the number of prompts used."""
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        prompt_path = Path("test_prompts.jsonl")
+        prompts = [
+            {"text": "Prompt 1"},
+            {"text": "Prompt 2"},
+            {"text": "Prompt 3"},
+            {"text": "Prompt 4"},
+        ]
+        prompt_path.write_text("\n".join(json.dumps(p) for p in prompts), encoding="utf-8")
+
+        mock_result = SweepResult(
+            model_id="m1",
+            backend_name="mlx",
+            quantization_levels=["a", "b"],
+            comparisons=[],
+            timestamp=datetime.now(UTC),
+            summary={},
+        )
+
+        with (
+            patch("infer_check.cli._resolve_prompts", return_value=prompt_path),
+            patch("infer_check.backends.base.get_backend_for_model"),
+            patch.object(TestRunner, "sweep", return_value=mock_result) as mock_sweep,
+        ):
+            # 1. Test global --num-prompts
+            result = runner.invoke(
+                main, ["--num-prompts", "2", "sweep", "--models", "a=m1,b=m2", "--prompts", "test_prompts.jsonl"]
+            )
+            assert result.exit_code == 0
+            args, kwargs = mock_sweep.call_args
+            captured_prompts = kwargs.get("prompts") or args[2]
+            assert len(captured_prompts) == 2
+            assert captured_prompts[0].text == "Prompt 1"
+            assert captured_prompts[1].text == "Prompt 2"
+
+            # 2. Test subcommand --num-prompts override
+            result = runner.invoke(
+                main,
+                [
+                    "--num-prompts",
+                    "2",
+                    "sweep",
+                    "--models",
+                    "a=m1,b=m2",
+                    "--prompts",
+                    "test_prompts.jsonl",
+                    "--num-prompts",
+                    "3",
+                ],
+            )
+            assert result.exit_code == 0
+            args, kwargs = mock_sweep.call_args
+            captured_prompts = kwargs.get("prompts") or args[2]
+            assert len(captured_prompts) == 3
+
+            # 3. Test without --num-prompts (should use all)
+            result = runner.invoke(main, ["sweep", "--models", "a=m1,b=m2", "--prompts", "test_prompts.jsonl"])
+            assert result.exit_code == 0
+            args, kwargs = mock_sweep.call_args
+            captured_prompts = kwargs.get("prompts") or args[2]
+            assert len(captured_prompts) == 4