FEATURE: add --prompts-file option to bench command for custom test suites (#30)

* feat(bench): add --prompts-file option with path traversal security

* test(bench): add unit tests for custom prompts and path traversal

* docs(bench): document --prompts-file usage in README.md

* feat(bench): add --prompts-file support with path validation

* test(bench): add unit tests for custom prompts and security checks

* style: remove trailing whitespace to pass ruff linting

* test: fix mock patch targets for local imports in bench command

* refactor(bench): simplify prompts-file logic and clean up comment

* test(bench): update assertions to match new prompts-file semantics

README.md

@@ -871,6 +871,10 @@ soup bench ./output
 
 # Customizing benchmarking parameters
 soup bench ./output --num-prompts 5 --max-tokens 256
+
+# Use custom prompts from a text file (one per line) or JSONL
+soup bench ./output --prompts-file my_prompts.txt
+soup bench ./output --prompts-file bench_suite.jsonl
 ```
 
 This acts as a built-in "speedometer," outputting Tokens-Per-Second (TPS), Total Latency, and Peak VRAM allocations into a clean status table.

soup_cli/commands/bench.py

@@ -32,7 +32,12 @@ def bench(
         3,
         "--num-prompts",
         "-n",
-        help="Number of prompts to run for averaging",
+        help="Number of prompts to run for averaging (ignored when --prompts-file is set)",
+    ),
+    prompts_file: Optional[str] = typer.Option(
+        None,
+        "--prompts-file",
+        help="Path to custom prompts file (.txt or .jsonl)",
     ),
 ) -> None:
     """Run an inference benchmark (speed and memory) on a loaded model."""
@@ -54,11 +59,63 @@ def bench(
             "10-100x slower than GPU -- results will not reflect production TPS."
         )
 
+    if prompts_file:
+        import json
+        p_path = Path(prompts_file).resolve()
+
+        try:
+            p_path.relative_to(Path.cwd())
+        except ValueError:
+            console.print(
+                f"[red]Security Error:[/] Path {p_path} is outside the current working directory."
+            )
+            raise typer.Exit(1)
+
+        if not p_path.is_file():
+            console.print(f"[red]Prompts file not found:[/] {p_path}")
+            raise typer.Exit(1)
+
+        prompts = []
+        try:
+            if p_path.suffix == ".jsonl":
+                with open(p_path, "r", encoding="utf-8") as f:
+                    for line in f:
+                        line = line.strip()
+                        if line:
+                            data = json.loads(line)
+                            if "prompt" in data:
+                                prompts.append(data["prompt"])
+            else:
+                with open(p_path, "r", encoding="utf-8") as f:
+                    for line in f:
+                        line = line.strip()
+                        if line:
+                            prompts.append(line)
+        except Exception as e:
+            console.print(f"[red]Failed to read prompts file:[/] {e}")
+            raise typer.Exit(1)
+
+        if not prompts:
+            console.print("[red]No prompts found in file.[/]")
+            raise typer.Exit(1)
+
+        # If --prompts-file is provided, we use all prompts in the file and ignore num_prompts.
+    else:
+        prompts = [
+            "Explain the theory of relativity briefly.",
+            "Write a short Python function to calculate fibonacci numbers.",
+            "What are the main consequences of the Industrial Revolution?",
+            "Compose a poem about a wandering space traveler.",
+            "Describe how a database index works under the hood.",
+        ]
+
+    actual_num_prompts = len(prompts) if prompts_file else num_prompts
+
     console.print(
         Panel(
             f"Model:    [bold]{model_path}[/]\n"
             f"Device:   [bold]{device}[/]\n"
-            f"Prompts:  [bold]{num_prompts}[/]\n"
+            f"Prompts:  [bold]{actual_num_prompts}[/]\n"
             f"Tokens/P: [bold]{max_tokens}[/]",
             title="Benchmarking Configuration",
         )
@@ -81,14 +138,7 @@ def bench(
     load_time = time.time() - start_load
     console.print(f"[green]Model loaded in {load_time:.2f}s.[/]\n")
 
-    prompts = [
-        "Explain the theory of relativity briefly.",
-        "Write a short Python function to calculate fibonacci numbers.",
-        "What are the main consequences of the Industrial Revolution?",
-        "Compose a poem about a wandering space traveler.",
-        "Describe how a database index works under the hood.",
-    ]
-    test_prompts = (prompts * (num_prompts // len(prompts) + 1))[:num_prompts]
+    test_prompts = (prompts * (actual_num_prompts // len(prompts) + 1))[:actual_num_prompts]
 
     # Warmup run: first inference includes CUDA kernel JIT compilation,
     # which would skew the average. Discarded from timing.
@@ -102,7 +152,7 @@ def bench(
     total_tokens = 0
     total_latency = 0.0
 
-    console.print(f"[bold]Running {num_prompts} test inferences...[/]")
+    console.print(f"[bold]Running {len(test_prompts)} test inferences...[/]")
 
     for i, prompt_text in enumerate(test_prompts):
         messages = [{"role": "user", "content": prompt_text}]

tests/test_bench.py

@@ -12,3 +12,47 @@ def test_bench_model_not_found():
     result = runner.invoke(app, ["bench", "nonexistent_model_path"])
     assert result.exit_code == 1
     assert "not found" in result.output.lower()
+
+def test_bench_custom_prompts(tmp_path, monkeypatch):
+    """Test using custom prompts from a text file and JSONL."""
+    monkeypatch.chdir(tmp_path)
+
+    dummy_model = tmp_path / "dummy_model"
+    dummy_model.mkdir()
+
+    # Text file
+    prompts_txt = tmp_path / "prompts.txt"
+    prompts_txt.write_text("Custom prompt 1\nCustom prompt 2\n")
+
+    # JSONL file
+    prompts_jsonl = tmp_path / "prompts.jsonl"
+    prompts_jsonl.write_text('{"prompt": "JSON prompt 1"}\n{"prompt": "JSON prompt 2"}\n')
+
+    # Path traversal
+    outside_file = tmp_path.parent / "outside.txt"
+    outside_file.write_text("Outside\n")
+
+    from unittest.mock import patch
+
+    with patch("soup_cli.commands.infer._load_model") as mock_load, \
+         patch("soup_cli.commands.infer._generate") as mock_generate:
+
+        mock_load.return_value = ("mock_model", "mock_tokenizer")
+        mock_generate.return_value = (None, 10)
+
+        # Test 1: TXT
+        result = runner.invoke(app, ["bench", str(dummy_model), "--prompts-file", "prompts.txt"])
+        assert result.exit_code == 0
+        assert "Running 2 test inferences" in result.output
+
+        # Test 2: JSONL
+        result = runner.invoke(app, ["bench", str(dummy_model), "--prompts-file", "prompts.jsonl"])
+        assert result.exit_code == 0
+        assert "Running 2 test inferences" in result.output
+
+        # Test 3: Path outside CWD
+        result = runner.invoke(
+            app, ["bench", str(dummy_model), "--prompts-file", str(outside_file)]
+        )
+        assert result.exit_code == 1
+        assert "Security Error" in result.output