Add files via upload

Author: jty1128

sweagent/run/run.py

@@ -27,42 +27,14 @@
 [bold][green]remove-unfinished[/green][/bold] or [bold][green]ru[/green][/bold]: Remove unfinished trajectories
 [bold][green]quick-stats[/green][/bold] or [bold][green]qs[/green][/bold]: Calculate quick stats from a directory of trajectories
 """
-
-import argparse
 import sys
-
 import rich
+import argparse
 
 
 def get_cli():
     parser = argparse.ArgumentParser(add_help=False)
-    parser.add_argument(
-        "command",
-        choices=[
-            "run",
-            "run-batch",
-            "run-replay",
-            "traj-to-demo",
-            "run-api",
-            "merge-preds",
-            "inspect",
-            "inspector",
-            "r",
-            "b",
-            "i",
-            "I",
-            "extract-pred",
-            "compare-runs",
-            "cr",
-            "remove-unfinished",
-            "ru",
-            "quick-stats",
-            "qs",
-            "shell",
-            "sh",
-        ],
-        nargs="?",
-    )
+    parser.add_argument("command", choices=["run", "run-batch", "run-replay", "traj-to-demo", "run-api", "merge-preds", "inspect", "inspector", "r", "b", "i", "I", "extract-pred", "compare-runs", "cr", "remove-unfinished", "ru", "quick-stats", "qs", "shell", "sh"], nargs="?")
     parser.add_argument("-h", "--help", action="store_true", help="Show this help message and exit")
     return parser
 
@@ -88,60 +60,50 @@ def main(args: list[str] | None = None):
     # Defer imports to avoid unnecessary long loading times
     if command in ["run", "r"]:
         from sweagent.run.run_single import run_from_cli as run_single_main
-
         run_single_main(remaining_args)
     elif command in ["run-batch", "b"]:
         from sweagent.run.run_batch import run_from_cli as run_batch_main
-
         run_batch_main(remaining_args)
+    elif command in ["run-locate"]:
+        from sweagent.run.run_locate import run_from_cli as run_locate_main
+        run_locate_main(remaining_args)
     elif command == "run-replay":
         from sweagent.run.run_replay import run_from_cli as run_replay_main
-
         run_replay_main(remaining_args)
     elif command == "traj-to-demo":
         from sweagent.run.run_traj_to_demo import run_from_cli as convert_traj_to_demo_main
-
         convert_traj_to_demo_main(remaining_args)
     elif command == "run-api":
         from sweagent.api.server import run_from_cli as run_api_main
-
         run_api_main(remaining_args)
     elif command == "merge-preds":
         from sweagent.run.merge_predictions import run_from_cli as merge_predictions_main
-
         merge_predictions_main(remaining_args)
     elif command in ["inspector", "I"]:
         from sweagent.inspector.server import run_from_cli as inspector_main
-
         inspector_main(remaining_args)
     elif command in ["inspect", "i"]:
         from sweagent.run.inspector_cli import main as inspect_main
-
         inspect_main(remaining_args)
     elif command == "extract-pred":
         from sweagent.run.extract_pred import run_from_cli as extract_pred_main
-
         extract_pred_main(remaining_args)
     elif command in ["compare-runs", "cr"]:
         from sweagent.run.compare_runs import run_from_cli as compare_runs_main
-
         compare_runs_main(remaining_args)
     elif command in ["remove-unfinished", "ru"]:
         from sweagent.run.remove_unfinished import run_from_cli as remove_unfinished_main
-
         remove_unfinished_main(remaining_args)
     elif command in ["quick-stats", "qs"]:
         from sweagent.run.quick_stats import run_from_cli as quick_stats_main
-
         quick_stats_main(remaining_args)
     elif command in ["shell", "sh"]:
         from sweagent.run.run_shell import run_from_cli as run_shell_main
-
         run_shell_main(remaining_args)
     else:
         msg = f"Unknown command: {command}"
         raise ValueError(msg)
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    sys.exit(main())