diff --git a/packages/moss-cli/README.md b/packages/moss-cli/README.md index 600d59df..8fcf8eb1 100644 --- a/packages/moss-cli/README.md +++ b/packages/moss-cli/README.md @@ -13,6 +13,7 @@ Moss CLI wraps the [Moss Python SDK](https://docs.moss.dev/) so you can build an - **Local by default** — downloads indexes for on-device queries, `--cloud` to skip - **Flexible auth** — CLI flags, environment variables, or config file - **Multiple output formats** — rich tables for humans, `--json` for scripts +- **Latency benchmarking** — measure p50/p95/p99 retrieval latency with `moss bench` - **Job tracking** — poll background jobs with live progress display - **Pipe-friendly** — stdin/stdout support for composing with other tools @@ -145,6 +146,28 @@ echo "what is AI" | moss query my-index moss query my-index "query" --json | jq '.docs[0].text' ``` +### Benchmark + +```bash +# Benchmark latency with a single inline query +moss bench my-index --query "what is semantic search?" --runs 20 + +# Multiple inline queries +moss bench my-index \ + --query "what is semantic search?" \ + --query "how does embedding work?" \ + --runs 20 --warmup 5 + +# Load queries from a file (one per line) +moss bench my-index --queries-file queries.txt --runs 50 + +# JSON output for CI/scripting +moss bench my-index --query "what is AI?" --json + +# Cloud mode (skip local index download) +moss bench my-index --query "what is AI?" --cloud --runs 10 +``` + ### Job Tracking ```bash diff --git a/packages/moss-cli/src/moss_cli/commands/bench.py b/packages/moss-cli/src/moss_cli/commands/bench.py new file mode 100644 index 00000000..e8e74033 --- /dev/null +++ b/packages/moss-cli/src/moss_cli/commands/bench.py @@ -0,0 +1,184 @@ +"""moss bench command — measure retrieval latency percentiles.""" + +from __future__ import annotations + +import asyncio +import json +import time +from pathlib import Path + +import typer +from moss import MossClient, QueryOptions +from rich.console import Console +from rich.table import Table + +from .. import output +from ..completion import complete_index_name +from ..config import resolve_credentials + +console = Console() + + +def _percentile(sorted_ms: list[float], p: float) -> float: + """Return the p-th percentile (0–100) of a sorted list using linear interpolation.""" + if not sorted_ms: + return 0.0 + n = len(sorted_ms) + idx = (p / 100) * (n - 1) + lo = int(idx) + hi = min(lo + 1, n - 1) + return sorted_ms[lo] + (idx - lo) * (sorted_ms[hi] - sorted_ms[lo]) + + +def bench_command( + ctx: typer.Context, + index_name: str = typer.Argument( + ..., help="Index name to benchmark", autocompletion=complete_index_name + ), + queries: list[str] | None = typer.Option( # noqa: B008 + None, "--query", "-q", help="Query string (repeatable)" + ), + queries_file: Path | None = typer.Option( # noqa: B008 + None, "--queries-file", "-f", help="File with one query per line" + ), + runs: int = typer.Option(20, "--runs", "-n", help="Number of timed runs per query"), + top_k: int = typer.Option(10, "--top-k", "-k", help="Number of results per query"), + alpha: float = typer.Option( + 0.8, "--alpha", "-a", help="Semantic weight (0.0=keyword, 1.0=semantic)" + ), + warmup: int = typer.Option( + 3, "--warmup", help="Number of warmup runs before timing (discarded)" + ), + cloud: bool = typer.Option(False, "--cloud", "-c", help="Query via cloud API"), + profile: str | None = typer.Option( + None, "--profile", help="Credential profile name" + ), + json_output: bool = typer.Option(False, "--json", help="Output as JSON"), +) -> None: + """Benchmark retrieval latency and report p50/p95/p99 percentiles.""" + json_mode = json_output or ctx.obj.get("json_output", False) + if profile: + ctx.obj["profile"] = profile + + all_queries: list[str] = list(queries or []) + if queries_file: + try: + lines = queries_file.read_text().splitlines() + all_queries.extend(ln.strip() for ln in lines if ln.strip()) + except OSError as e: + output.print_error(f"Cannot read queries file: {e}", json_mode) + raise typer.Exit(1) + + if not all_queries: + output.print_error( + "No queries provided. Use --query or --queries-file.", json_mode + ) + raise typer.Exit(1) + + if runs < 1: + output.print_error("--runs must be >= 1.", json_mode) + raise typer.Exit(1) + + if warmup < 0: + output.print_error("--warmup must be >= 0.", json_mode) + raise typer.Exit(1) + + pid, pkey = resolve_credentials( + ctx.obj.get("project_id"), ctx.obj.get("project_key"), ctx.obj.get("profile") + ) + + client = MossClient(pid, pkey) + + async def _run() -> None: + if not cloud: + if not json_mode: + console.print(f"Loading index [cyan]{index_name}[/cyan] locally...") + await client.load_index(index_name) + + options = QueryOptions(top_k=top_k) if cloud else QueryOptions(top_k=top_k, alpha=alpha) + + if not json_mode: + q_word = "query" if len(all_queries) == 1 else "queries" + console.print( + f"Running [bold]{warmup * len(all_queries)}[/bold] warmup + " + f"[bold]{runs * len(all_queries)}[/bold] timed iterations " + f"across [bold]{len(all_queries)}[/bold] {q_word}..." + ) + + for q in all_queries: + for _ in range(warmup): + await client.query(index_name, q, options) + + # Store latencies by input position so repeated queries are each tracked separately, + # preserving weighted-workload semantics. + per_query: list[tuple[str, list[float]]] = [(q, []) for q in all_queries] + for q, latencies in per_query: + for _ in range(runs): + t0 = time.perf_counter() + await client.query(index_name, q, options) + latencies.append((time.perf_counter() - t0) * 1000) + + all_latencies: list[float] = [] + query_stats = [] + for q, latencies in per_query: + s = sorted(latencies) + all_latencies.extend(s) + query_stats.append( + { + "query": q, + "runs": len(latencies), + "p50_ms": _percentile(s, 50), + "p95_ms": _percentile(s, 95), + "p99_ms": _percentile(s, 99), + "min_ms": s[0], + "max_ms": s[-1], + } + ) + + overall_sorted = sorted(all_latencies) + overall = { + "p50_ms": _percentile(overall_sorted, 50), + "p95_ms": _percentile(overall_sorted, 95), + "p99_ms": _percentile(overall_sorted, 99), + "min_ms": overall_sorted[0], + "max_ms": overall_sorted[-1], + "total_runs": len(overall_sorted), + } + + if json_mode: + print( + json.dumps( + {"index": index_name, "queries": query_stats, "overall": overall}, + indent=2, + ) + ) + return + + if len(query_stats) > 1: + table = Table(title=f"Benchmark — {index_name}") + table.add_column("Query", max_width=40) + table.add_column("p50 (ms)", justify="right") + table.add_column("p95 (ms)", justify="right") + table.add_column("p99 (ms)", justify="right") + table.add_column("min (ms)", justify="right") + table.add_column("max (ms)", justify="right") + for qs in query_stats: + table.add_row( + qs["query"][:40], + f"{qs['p50_ms']:.2f}", + f"{qs['p95_ms']:.2f}", + f"{qs['p99_ms']:.2f}", + f"{qs['min_ms']:.2f}", + f"{qs['max_ms']:.2f}", + ) + console.print(table) + + console.print(f"\n[bold]Overall ({overall['total_runs']} runs)[/bold]") + console.print(f" p50: [cyan]{overall['p50_ms']:.2f} ms[/cyan]") + console.print(f" p95: [cyan]{overall['p95_ms']:.2f} ms[/cyan]") + console.print(f" p99: [cyan]{overall['p99_ms']:.2f} ms[/cyan]") + console.print( + f" min: {overall['min_ms']:.2f} ms max: {overall['max_ms']:.2f} ms" + ) + + asyncio.run(_run()) diff --git a/packages/moss-cli/src/moss_cli/main.py b/packages/moss-cli/src/moss_cli/main.py index eae0ef50..ad2cc957 100644 --- a/packages/moss-cli/src/moss_cli/main.py +++ b/packages/moss-cli/src/moss_cli/main.py @@ -3,10 +3,10 @@ from __future__ import annotations import logging -from typing import Optional import typer +from .commands.bench import bench_command from .commands.completions import completions_command from .commands.doc import doc_app from .commands.index import index_app @@ -33,6 +33,7 @@ app.add_typer(profile_app, name="profile", help="Manage auth profiles") # Register top-level commands +app.command(name="bench")(bench_command) app.command(name="query")(query_command) app.command(name="init")(init_command) app.command(name="version")(version_command) @@ -44,13 +45,13 @@ @app.callback() def main( ctx: typer.Context, - project_id: Optional[str] = typer.Option( + project_id: str | None = typer.Option( None, "--project-id", "-p", envvar="MOSS_PROJECT_ID", help="Project ID" ), - project_key: Optional[str] = typer.Option( + project_key: str | None = typer.Option( None, "--project-key", envvar="MOSS_PROJECT_KEY", help="Project key" ), - profile: Optional[str] = typer.Option( + profile: str | None = typer.Option( None, "--profile", envvar="MOSS_PROFILE", @@ -80,7 +81,7 @@ def run() -> None: app() except (typer.Exit, typer.Abort, SystemExit): raise - except Exception as e: + except Exception as e: # noqa: BLE001 print_error(str(e)) raise typer.Exit(1) diff --git a/packages/moss-cli/tests/test_bench.py b/packages/moss-cli/tests/test_bench.py new file mode 100644 index 00000000..97aeb8cb --- /dev/null +++ b/packages/moss-cli/tests/test_bench.py @@ -0,0 +1,257 @@ +"""Tests for moss bench command.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from moss_cli.commands.bench import _percentile +from moss_cli.main import app +from typer.testing import CliRunner + +runner = CliRunner() + + +# --------------------------------------------------------------------------- +# _percentile unit tests +# --------------------------------------------------------------------------- + + +def test_percentile_single_element() -> None: + assert _percentile([5.0], 50) == 5.0 + assert _percentile([5.0], 99) == 5.0 + + +def test_percentile_two_elements() -> None: + data = [1.0, 3.0] + assert _percentile(data, 0) == 1.0 + assert _percentile(data, 100) == 3.0 + assert _percentile(data, 50) == pytest.approx(2.0) + + +def test_percentile_p50_odd_length() -> None: + data = sorted([1.0, 2.0, 3.0, 4.0, 5.0]) + assert _percentile(data, 50) == pytest.approx(3.0) + + +def test_percentile_p95_known_values() -> None: + # 20 evenly-spaced values 1..20 + data = [float(i) for i in range(1, 21)] + # p95 index = 0.95 * 19 = 18.05 → interpolate between 19 and 20 + expected = 19.0 + 0.05 * (20.0 - 19.0) + assert _percentile(data, 95) == pytest.approx(expected) + + +def test_percentile_p99_known_values() -> None: + data = [float(i) for i in range(1, 101)] # 1..100 + # p99 index = 0.99 * 99 = 98.01 → between index 98 (99.0) and 99 (100.0) + expected = 99.0 + 0.01 * (100.0 - 99.0) + assert _percentile(data, 99) == pytest.approx(expected) + + +def test_percentile_empty_returns_zero() -> None: + assert _percentile([], 50) == 0.0 + + +# --------------------------------------------------------------------------- +# CLI validation tests (no credentials / no network needed) +# --------------------------------------------------------------------------- + + +def _make_mock_client() -> Any: + mock_result = MagicMock() + mock_result.docs = [] + mock_result.query = "q" + mock_result.index_name = "test-index" + mock_result.time_taken_ms = 1.0 + + mock_client = MagicMock() + mock_client.load_index = AsyncMock(return_value=None) + mock_client.query = AsyncMock(return_value=mock_result) + return mock_client + + +def test_bench_no_queries_error() -> None: + with patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")): + result = runner.invoke(app, ["bench", "my-index"]) + assert result.exit_code != 0 + assert "No queries provided" in result.output + + +def test_bench_zero_runs_error() -> None: + with patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")): + result = runner.invoke(app, ["bench", "my-index", "--query", "hello", "--runs", "0"]) + assert result.exit_code != 0 + assert "--runs must be >= 1" in result.output + + +def test_bench_negative_warmup_error() -> None: + with patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")): + result = runner.invoke( + app, ["bench", "my-index", "--query", "hello", "--warmup", "-1"] + ) + assert result.exit_code != 0 + assert "--warmup must be >= 0" in result.output + + +def test_bench_missing_queries_file_error(tmp_path: Path) -> None: + missing = tmp_path / "nonexistent.txt" + with patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")): + result = runner.invoke( + app, ["bench", "my-index", "--queries-file", str(missing)] + ) + assert result.exit_code != 0 + assert "Cannot read queries file" in result.output + + +# --------------------------------------------------------------------------- +# Happy-path tests (mocked MossClient) +# --------------------------------------------------------------------------- + + +def test_bench_single_query_human_output() -> None: + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + result = runner.invoke( + app, + ["bench", "my-index", "--query", "what is semantic search?", "--runs", "5", "--warmup", "1"], + ) + assert result.exit_code == 0 + assert "Overall" in result.output + assert "p50" in result.output + assert "p95" in result.output + assert "p99" in result.output + + +def test_bench_json_output() -> None: + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + result = runner.invoke( + app, + ["--json", "bench", "my-index", "--query", "hello", "--runs", "3", "--warmup", "0"], + ) + assert result.exit_code == 0 + data = json.loads(result.output) + assert data["index"] == "my-index" + assert len(data["queries"]) == 1 + assert data["queries"][0]["runs"] == 3 + assert "p50_ms" in data["queries"][0] + assert "p95_ms" in data["queries"][0] + assert "p99_ms" in data["queries"][0] + assert "overall" in data + + +def test_bench_queries_file(tmp_path: Path) -> None: + qfile = tmp_path / "queries.txt" + qfile.write_text("what is AI?\nhow does embedding work?\n") + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + result = runner.invoke( + app, + ["bench", "my-index", "--queries-file", str(qfile), "--runs", "3", "--warmup", "0"], + ) + assert result.exit_code == 0 + assert "Overall" in result.output + # Two queries → table should be rendered + assert "Benchmark" in result.output + + +def test_bench_cloud_skips_load_index() -> None: + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + result = runner.invoke( + app, + ["bench", "my-index", "--query", "hello", "--cloud", "--runs", "2", "--warmup", "0"], + ) + assert result.exit_code == 0 + mock_client.load_index.assert_not_called() + + +def test_bench_warmup_calls_are_discarded() -> None: + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + runner.invoke( + app, + ["bench", "my-index", "--query", "hello", "--runs", "4", "--warmup", "2", "--cloud"], + ) + # 2 warmup + 4 timed = 6 total calls + assert mock_client.query.call_count == 6 + + +def test_bench_cloud_omits_alpha() -> None: + """In cloud mode QueryOptions must not receive alpha.""" + mock_client = _make_mock_client() + captured: list[Any] = [] + + async def capture_query(index: str, q: str, opts: Any) -> Any: + captured.append(opts) + return mock_client.query.return_value + + mock_client.query = capture_query # type: ignore[method-assign] + + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + patch("moss_cli.commands.bench.QueryOptions") as mock_qo, + ): + mock_qo.return_value = MagicMock() + runner.invoke( + app, + ["bench", "my-index", "--query", "hello", "--cloud", "--runs", "1", "--warmup", "0"], + ) + mock_qo.assert_called_once_with(top_k=10) + + +def test_bench_duplicate_queries_preserved_as_separate_entries(tmp_path: Path) -> None: + """Repeated queries must be tracked as separate workload entries (weighted traffic).""" + qfile = tmp_path / "queries.txt" + qfile.write_text("hello\nhello\nworld\n") + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + result = runner.invoke( + app, + ["--json", "bench", "my-index", "--queries-file", str(qfile), "--runs", "2", "--warmup", "0"], + ) + assert result.exit_code == 0 + data = json.loads(result.output) + # 3 input lines → 3 separate entries, not 2 + assert len(data["queries"]) == 3 + assert data["overall"]["total_runs"] == 6 # 3 entries × 2 runs + + +def test_bench_json_flag_after_subcommand() -> None: + """--json placed after the subcommand (moss bench ... --json) must work.""" + mock_client = _make_mock_client() + with ( + patch("moss_cli.commands.bench.resolve_credentials", return_value=("pid", "pkey")), + patch("moss_cli.commands.bench.MossClient", return_value=mock_client), + ): + result = runner.invoke( + app, + ["bench", "my-index", "--query", "hello", "--runs", "3", "--warmup", "0", "--json"], + ) + assert result.exit_code == 0 + data = json.loads(result.output) + assert data["index"] == "my-index" + assert "overall" in data