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

Author: NullPointerDepressiveDisorder

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)
@@ -51,11 +57,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="Default number of prompts to use from a suite.",
+)
 @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 +116,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.
 
@@ -121,9 +135,11 @@ def sweep(
     from infer_check.runner import TestRunner
     from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
+    # 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
 
     # Parse label=model_path pairs
     model_map: dict[str, str] = {}
@@ -153,6 +169,9 @@ def sweep(
         tag = " (baseline)" if label == baseline_label else ""
         console.print(f"  {label}: {path}{tag}")
     prompt_list = load_suite(_resolve_prompts(prompts))
+    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:
@@ -308,6 +327,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.
 
@@ -337,9 +357,11 @@ def compare(
     from infer_check.runner import TestRunner
     from infer_check.suites.loader import load_suite
 
-    # Update max_tokens from subcommand if provided
+    # 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
 
     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)
@@ -351,6 +373,9 @@ def compare(
     )
 
     prompt_list = load_suite(_resolve_prompts(prompts))
+    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:
@@ -573,15 +598,18 @@ 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
+    # 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
 
     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)
@@ -592,6 +620,9 @@ def diff(
     console.print(f"[bold cyan]diff[/bold cyan] model={model} backends={backend_names} quant={quant}")
 
     prompt_list = load_suite(_resolve_prompts(prompts))
+    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:
@@ -693,15 +724,18 @@ 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
+    # 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
 
     concurrency_levels = [int(c.strip()) for c in concurrency.split(",") if c.strip()]
 
@@ -716,6 +750,9 @@ def stress(
     )
 
     prompt_list = load_suite(_resolve_prompts(prompts))
+    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:
@@ -790,15 +827,18 @@ 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
+    # 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
 
     backend_instance = get_backend_for_model(
         model_str=model,
@@ -809,6 +849,9 @@ 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))
+    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:

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