cli commands refactored

Author: DzmitryPihulski

.gitignore

@@ -3,6 +3,7 @@ dataset/sqlite_tables.db
 *__pycache__
 .env
 dist/
+testenv
 
 *.egg-info/
 .pdm-python

llmsql/__main__.py

@@ -1,153 +1,11 @@
-import argparse
-import inspect
-import json
+from llmsql._cli import LLMSQLCLI
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(prog="llmsql", description="LLMSQL CLI")
-    subparsers = parser.add_subparsers(dest="command")
-
-    # ================================================================
-    # Inference command
-    # ================================================================
-    inference_examples = r"""
-Examples:
-
-  # 1️⃣ Run inference with Transformers backend
-  llmsql inference --method transformers \
-      --model-or-model-name-or-path Qwen/Qwen2.5-1.5B-Instruct \
-      --output-file outputs/preds_transformers.jsonl \
-      --batch-size 8 \
-      --num-fewshots 5
-
-  # 2️⃣ Run inference with vLLM backend
-  llmsql inference --method vllm \
-      --model-name Qwen/Qwen2.5-1.5B-Instruct \
-      --output-file outputs/preds_vllm.jsonl \
-      --batch-size 8 \
-      --num-fewshots 5
-
-  # 3️⃣ Pass model-specific kwargs (for Transformers)
-  llmsql inference --method transformers \
-      --model-or-model-name-or-path meta-llama/Llama-3-8b-instruct \
-      --output-file outputs/llama_preds.jsonl \
-      --model-args '{"attn_implementation": "flash_attention_2", "torch_dtype": "bfloat16"}'
-
-  # 4️⃣ Pass LLM init kwargs (for vLLM)
-  llmsql inference --method vllm \
-      --model-name mistralai/Mixtral-8x7B-Instruct-v0.1 \
-      --output-file outputs/mixtral_preds.jsonl \
-      --llm-kwargs '{"max_model_len": 4096, "gpu_memory_utilization": 0.9}'
-
-  # 5️⃣ Override generation parameters dynamically
-  llmsql inference --method transformers \
-      --model-or-model-name-or-path Qwen/Qwen2.5-1.5B-Instruct \
-      --output-file outputs/temp_0.9.jsonl \
-      --temperature 0.9 \
-      --generate-kwargs '{"do_sample": true, "top_p": 0.9, "top_k": 40}'
-"""
-
-    inf_parser = subparsers.add_parser(
-        "inference",
-        help="Run inference using either Transformers or vLLM backend.",
-        description="Run SQL generation using a chosen inference method "
-        "(either 'transformers' or 'vllm').",
-        epilog=inference_examples,
-        formatter_class=argparse.RawTextHelpFormatter,
-    )
-
-    inf_parser.add_argument(
-        "--method",
-        type=str,
-        required=True,
-        choices=["transformers", "vllm"],
-        help="Inference backend to use ('transformers' or 'vllm').",
-    )
-
-    # ================================================================
-    # Parse CLI
-    # ================================================================
-    args, extra = parser.parse_known_args()
-
-    # ------------------------------------------------
-    # Inference
-    # ------------------------------------------------
-    if args.command == "inference":
-        if args.method == "vllm":
-            from llmsql import inference_vllm as inference_fn
-        elif args.method == "transformers":
-            from llmsql import inference_transformers as inference_fn  # type: ignore
-        else:
-            raise ValueError(f"Unknown inference method: {args.method}")
-
-        # Dynamically create parser from the function signature
-        fn_parser = argparse.ArgumentParser(
-            prog=f"llmsql inference --method {args.method}",
-            description=f"Run inference using {args.method} backend",
-        )
-
-        sig = inspect.signature(inference_fn)
-        for name, param in sig.parameters.items():
-            if param.kind == inspect.Parameter.VAR_KEYWORD:
-                fn_parser.add_argument(
-                    "--llm-kwargs",
-                    default="{}",
-                    help="Additional LLM kwargs as a JSON string, e.g. '{\"top_p\": 0.9}'",
-                )
-                fn_parser.add_argument(
-                    "--generate-kwargs",
-                    default="{}",
-                    help="",
-                )
-                continue
-            arg_name = f"--{name.replace('_', '-')}"
-            default = param.default
-            if default is inspect.Parameter.empty:
-                fn_parser.add_argument(arg_name, required=True)
-            else:
-                if isinstance(default, bool):
-                    fn_parser.add_argument(
-                        arg_name,
-                        action="store_true" if not default else "store_false",
-                        help=f"(default: {default})",
-                    )
-                elif default is None:
-                    fn_parser.add_argument(arg_name, type=str, default=None)
-                else:
-                    fn_parser.add_argument(
-                        arg_name, type=type(default), default=default
-                    )
-
-        fn_args = fn_parser.parse_args(extra)
-        fn_kwargs = vars(fn_args)
-
-        if "llm_kwargs" in fn_kwargs and isinstance(fn_kwargs["llm_kwargs"], str):
-            try:
-                fn_kwargs["llm_kwargs"] = json.loads(fn_kwargs["llm_kwargs"])
-            except json.JSONDecodeError:
-                print("⚠️  Could not parse --llm-kwargs JSON, passing as string.")
-
-        if fn_kwargs.get("model_args") is not None:
-            try:
-                fn_kwargs["model_args"] = json.loads(fn_kwargs["model_args"])
-            except json.JSONDecodeError:
-                raise
-
-        if fn_kwargs.get("generate_kwargs") is not None:
-            try:
-                fn_kwargs["generate_kwargs"] = json.loads(fn_kwargs["generate_kwargs"])
-            except json.JSONDecodeError:
-                raise
-
-        print(f"🔹 Running {args.method} inference with arguments:")
-        for k, v in fn_kwargs.items():
-            print(f"  {k}: {v}")
-
-        results = inference_fn(**fn_kwargs)
-        print(f"✅ Inference complete. Generated {len(results)} results.")
-
-    else:
-        parser.print_help()
+    """Main CLI entry point."""
+    parser = LLMSQLCLI()
+    args = parser.parse_args()
+    parser.execute(args)
 
 
 if __name__ == "__main__":

llmsql/_cli/__init__.py

@@ -0,0 +1,7 @@
+"""
+CLI subcommands to run from the terminal.
+"""
+
+from .llmsql_cli import LLMSQLCLI
+
+__all__ = ["LLMSQLCLI"]

llmsql/_cli/evaluation.py

@@ -0,0 +1,106 @@
+import argparse
+import json
+
+from llmsql.config.config import DEFAULT_WORKDIR_PATH
+from llmsql.evaluation.evaluate import evaluate
+
+
+class EvaluationCommand:
+    """CLI wrapper for the `evaluate()` function."""
+
+    @staticmethod
+    def register(subparsers: argparse._SubParsersAction) -> None:
+        eval_parser = subparsers.add_parser(
+            "evaluate",
+            help="Evaluate SQL predictions against the LLMSQL benchmark.",
+            description="Evaluate predicted SQL against the LLMSQL benchmark.",
+            formatter_class=argparse.RawTextHelpFormatter,
+        )
+
+        eval_parser.add_argument(
+            "--outputs",
+            required=True,
+            help=(
+                "Path to predictions JSONL file OR inline JSON list.\n"
+                "Examples:\n"
+                "  --outputs outputs/preds.jsonl\n"
+                '  --outputs \'[{"id":1,"sql":"SELECT ..."}]\''
+            ),
+        )
+
+        eval_parser.add_argument(
+            "--workdir-path",
+            type=str,
+            default=None,
+            help=(
+                f"Optional. Where to store help .db files. Default: {DEFAULT_WORKDIR_PATH}"
+            ),
+        )
+
+        eval_parser.add_argument(
+            "--questions-path",
+            type=str,
+            default=None,
+            help=(
+                "Optional. Where is the questions of the benchmark stored. If not provided, questions will be downloaded."
+            ),
+        )
+
+        eval_parser.add_argument(
+            "--db-path",
+            type=str,
+            default=None,
+            help=(
+                "Optional. Where is the db file of the benchmark stored. If not provided, db file will be downloaded."
+            ),
+        )
+
+        eval_parser.add_argument(
+            "--save-report",
+            type=str,
+            default=None,
+            help=(
+                "Optional. Manual save path. If None → auto-generated with name 'evaluation_results_{uuid}.json'."
+            ),
+        )
+
+        # Boolean toggle
+        eval_parser.add_argument(
+            "--show-mismatches",
+            type=bool,
+            default=True,
+            help="Optional. Show mismatches during evaluation. Default: True.",
+        )
+
+        eval_parser.add_argument(
+            "--max-mismatches",
+            type=int,
+            default=5,
+            help="Optional. Number of mismatches to print. Default: 5",
+        )
+
+        # Dispatcher
+        eval_parser.set_defaults(func=EvaluationCommand.execute)
+
+    # ----------------------------------------------------------
+    @staticmethod
+    def execute(args: argparse.Namespace) -> None:
+        """Run evaluation function."""
+        # Try inline JSON first; fallback to path
+        try:
+            parsed = json.loads(args.outputs)
+            outputs = parsed if isinstance(parsed, list) else args.outputs
+        except Exception:
+            outputs = args.outputs
+
+        result = evaluate(
+            outputs=outputs,
+            workdir_path=args.workdir_path,
+            questions_path=args.questions_path,
+            db_path=args.db_path,
+            save_report=args.save_report,
+            show_mismatches=args.show_mismatches,
+            max_mismatches=args.max_mismatches,
+        )
+
+        print(json.dumps(result, indent=2))