rename main and run to cli for consistency (#1148)

Author: rogthefrog

pyproject.toml

@@ -117,8 +117,8 @@ amazon = ["modelgauge_amazon"]
 all_plugins = ["modelgauge_anthropic", "modelgauge_azure", "modelgauge_baseten", "modelgauge_demo_plugin", "modelgauge_nvidia", "modelgauge_perspective_api", "modelgauge_google", "modelgauge_vertexai", "modelgauge_mistral", "modelgauge_amazon"]
 
 [tool.poetry.scripts]
-modelbench = "modelbench.run:cli"
-modelgauge = "modelgauge.main:main"
+modelbench = "modelbench.cli:cli"
+modelgauge = "modelgauge.cli:cli"
 
 [tool.pytest.ini_options]
 addopts = [

src/modelbench/run.py src/modelbench/cli.pysrc/modelbench/run.py renamed to src/modelbench/cli.py

@@ -16,7 +16,7 @@
     create_sut_options,
     display_header,
     display_list_item,
-    modelgauge_cli,
+    cli,
     sut_options_options,
     validate_uid,
 )
@@ -38,7 +38,7 @@
 logger = logging.getLogger(__name__)
 
 
-@modelgauge_cli.command(name="list")
+@cli.command(name="list")
 @LOCAL_PLUGIN_DIR_OPTION
 def list_command() -> None:
     """Overview of Plugins, Tests, and SUTs."""
@@ -85,7 +85,7 @@ def format_missing_secrets(missing):
     click.echo()
 
 
-@modelgauge_cli.command()
+@cli.command()
 @LOCAL_PLUGIN_DIR_OPTION
 def list_tests() -> None:
     """List details about all registered tests."""
@@ -94,7 +94,7 @@ def list_tests() -> None:
         _display_factory_entry(test_uid, test_entry, secrets)
 
 
-@modelgauge_cli.command()
+@cli.command()
 @LOCAL_PLUGIN_DIR_OPTION
 def list_suts():
     """List details about all registered SUTs (System Under Test)."""
@@ -103,7 +103,7 @@ def list_suts():
         _display_factory_entry(sut_uid, sut, secrets)
 
 
-@modelgauge_cli.command()
+@cli.command()
 @LOCAL_PLUGIN_DIR_OPTION
 def list_annotators():
     """List details about all registered SUTs (System Under Test)."""
@@ -112,7 +112,7 @@ def list_annotators():
         _display_factory_entry(annotator_uid, annotator, secrets)
 
 
-@modelgauge_cli.command()
+@cli.command()
 @LOCAL_PLUGIN_DIR_OPTION
 def list_secrets() -> None:
     """List details about secrets modelgauge might need."""
@@ -124,7 +124,7 @@ def list_secrets() -> None:
         display_header("No secrets used by any installed plugin.")
 
 
-@modelgauge_cli.command()
+@cli.command()
 @LOCAL_PLUGIN_DIR_OPTION
 @click.option("--sut", "-s", help="Which SUT to run.", required=True)
 @sut_options_options
@@ -162,7 +162,7 @@ def run_sut(
     click.echo(f"Normalized response: {result.model_dump_json(indent=2)}\n")
 
 
-@modelgauge_cli.command()
+@cli.command()
 @click.option("--test", "-t", help="Which registered TEST to run.", required=True, callback=validate_uid)
 @LOCAL_PLUGIN_DIR_OPTION
 @click.option("--sut", "-s", help="Which SUT to run.", required=True, multiple=False)
@@ -233,7 +233,7 @@ def run_test(
     print("Full TestRecord json written to", output_file)
 
 
-@modelgauge_cli.command()
+@cli.command()
 @sut_options_options
 @click.option(
     "sut_uid",
@@ -352,9 +352,5 @@ def show_progress(data):
         pipeline_runner.run(show_progress, debug)
 
 
-def main():
-    modelgauge_cli()
-
-
 if __name__ == "__main__":
-    main()
+    cli()

src/modelgauge/main.py src/modelgauge/cli.pysrc/modelgauge/main.py renamed to src/modelgauge/cli.py

@@ -13,9 +13,9 @@
 
 
 @click.group(name="modelgauge")
-def modelgauge_cli():
+def cli():
     """Run the ModelGauge library from the command line."""
-    # To add a command, decorate your function with @modelgauge_cli.command().
+    # To add a command, decorate your function with @cli.command().
 
     # Always create the config directory if it doesn't already exist.
     write_default_config()

src/modelgauge/command_line.py

@@ -1,11 +1,11 @@
 import together  # type: ignore
 from collections import defaultdict
-from modelgauge.command_line import display_header, display_list_item, modelgauge_cli
+from modelgauge.command_line import display_header, display_list_item, cli
 from modelgauge.config import load_secrets_from_config
 from modelgauge.suts.together_client import TogetherApiKey
 
 
-@modelgauge_cli.command()
+@cli.command()
 def list_together():
     """List all models available in together.ai."""
 

src/modelgauge/suts/together_cli.py

@@ -14,7 +14,7 @@
 from modelbench.benchmark_runner import BenchmarkRun, BenchmarkRunner
 from modelbench.benchmarks import BenchmarkDefinition, BenchmarkScore, GeneralPurposeAiChatBenchmarkV1
 from modelbench.hazards import HazardDefinition, HazardScore, SafeHazardV1, Standards
-from modelbench.run import benchmark, cli, get_benchmark
+from modelbench.cli import benchmark, cli, get_benchmark
 from modelbench.scoring import ValueEstimate
 from modelgauge.base_test import PromptResponseTest
 from modelgauge.preflight import make_sut
@@ -136,18 +136,18 @@ def wrapper(*args, **kwargs):
     @pytest.fixture(autouse=False)
     def mock_run_benchmarks(self, sut, monkeypatch, tmp_path):
         mock = MagicMock(return_value=fake_benchmark_run(AHazard(), sut, tmp_path))
-        monkeypatch.setattr(modelbench.run, "run_benchmarks_for_sut", mock)
+        monkeypatch.setattr(modelbench.cli, "run_benchmarks_for_sut", mock)
         return mock
 
     @pytest.fixture(autouse=False)
     def mock_score_benchmarks(self, sut, monkeypatch):
         mock = MagicMock(return_value=[self.mock_score(sut)])
-        monkeypatch.setattr(modelbench.run, "score_benchmarks", mock)
+        monkeypatch.setattr(modelbench.cli, "score_benchmarks", mock)
         return mock
 
     @pytest.fixture(autouse=True)
     def do_print_summary(self, monkeypatch):
-        monkeypatch.setattr(modelbench.run, "print_summary", MagicMock())
+        monkeypatch.setattr(modelbench.cli, "print_summary", MagicMock())
 
     @pytest.fixture
     def runner(self):
@@ -225,7 +225,7 @@ def test_benchmark_multiple_suts_produces_json(
         )
 
         mock = MagicMock(return_value=[self.mock_score(sut_uid, benchmark), self.mock_score("demo_yes_no", benchmark)])
-        monkeypatch.setattr(modelbench.run, "score_benchmarks", mock)
+        monkeypatch.setattr(modelbench.cli, "score_benchmarks", mock)
 
         result = runner.invoke(
             cli,

tests/modelbench_tests/test_run.py

@@ -8,7 +8,7 @@
 import pytest
 from click.testing import CliRunner, Result
 
-from modelgauge import main
+from modelgauge import cli
 from modelgauge.annotator_registry import ANNOTATORS
 from modelgauge.annotator_set import AnnotatorSet
 from modelgauge.command_line import validate_uid
@@ -34,7 +34,7 @@
 
 def run_cli(*args) -> Result:
     # noinspection PyTypeChecker
-    result = CliRunner().invoke(main.modelgauge_cli, args, catch_exceptions=False)
+    result = CliRunner().invoke(cli.cli, args, catch_exceptions=False)
     return result
 
 
@@ -96,7 +96,7 @@ def test_run_sut_invalid_uid():
 def test_run_sut_with_options(mock_translate_text_prompt):
     runner = CliRunner()
     result = runner.invoke(
-        main.modelgauge_cli,
+        cli.cli,
         [
             "run-sut",
             "--sut",
@@ -189,7 +189,7 @@ def test_run_job_sut_only_output_name(caplog, tmp_path, prompts_file):
     caplog.set_level(logging.INFO)
     runner = CliRunner()
     result = runner.invoke(
-        main.modelgauge_cli,
+        cli.cli,
         ["run-job", "--sut", "demo_yes_no", "--output-dir", tmp_path, str(prompts_file)],
         catch_exceptions=False,
     )
@@ -210,7 +210,7 @@ def test_run_job_with_tag_output_name(caplog, tmp_path, prompts_file):
     caplog.set_level(logging.INFO)
     runner = CliRunner()
     result = runner.invoke(
-        main.modelgauge_cli,
+        cli.cli,
         ["run-job", "--sut", "demo_yes_no", "--output-dir", tmp_path, "--tag", "test", str(prompts_file)],
         catch_exceptions=False,
     )
@@ -226,7 +226,7 @@ def test_run_job_sut_and_annotator_output_name(caplog, tmp_path, prompts_file):
     caplog.set_level(logging.INFO)
     runner = CliRunner()
     result = runner.invoke(
-        main.modelgauge_cli,
+        cli.cli,
         [
             "run-job",
             "--sut",
@@ -257,7 +257,7 @@ def test_run_job_annotators_only_output_name(caplog, tmp_path, prompt_responses_
     caplog.set_level(logging.INFO)
     runner = CliRunner()
     result = runner.invoke(
-        main.modelgauge_cli,
+        cli.cli,
         ["run-job", "--annotator", "demo_annotator", "--output-dir", tmp_path, str(prompt_responses_file)],
         catch_exceptions=False,
     )
@@ -291,7 +291,7 @@ def evaluate(self, item):
     with patch.dict(sys.modules, {"modelgauge.private_ensemble_annotator_set": dummy_module}):
         runner = CliRunner()
         result = runner.invoke(
-            main.modelgauge_cli,
+            cli.cli,
             [
                 "run-job",
                 "--ensemble",
@@ -318,7 +318,7 @@ def evaluate(self, item):
 def test_run_missing_ensemble_raises_error(tmp_path, prompt_responses_file):
     runner = CliRunner()
     result = runner.invoke(
-        main.modelgauge_cli,
+        cli.cli,
         [
             "run-job",
             "--ensemble",