cli.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from src.config import load_config
from src.evaluators import ExperimentRunner
from src.llm import OllamaClient, OllamaConfig, OllamaUnavailableError
from src.metrics import compute_summary, load_jsonl, write_analysis_artifacts
from src.plots import write_summary_figure


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="split-inference-bench CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    run_parser = subparsers.add_parser("run", help="Run an experiment config")
    run_parser.add_argument("--config", required=True)
    run_parser.add_argument("--output-dir", required=True)

    summarize_parser = subparsers.add_parser("summarize", help="Summarize a run directory")
    summarize_parser.add_argument("--runs-jsonl", required=True)
    summarize_parser.add_argument("--output", required=True)

    plot_parser = subparsers.add_parser("plot", help="Create a summary figure")
    plot_parser.add_argument("--summary", required=True)
    plot_parser.add_argument("--output", required=True)

    check_parser = subparsers.add_parser("check-ollama", help="Check local Ollama availability")
    check_parser.add_argument("--base-url", default="http://localhost:11434/api")
    check_parser.add_argument("--model", default="gemma3:1b")
    check_parser.add_argument("--timeout-s", type=float, default=30.0)
    check_parser.add_argument("--max-retries", type=int, default=1)
    return parser


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()
    if args.command == "run":
        config = load_config(args.config).raw
        output_dir = Path(args.output_dir)
        if output_dir.exists() and any(output_dir.iterdir()):
            raise FileExistsError(f"Refusing to overwrite existing output directory: {output_dir}")
        summary = ExperimentRunner(config=config, output_dir=output_dir).run()
        print(json.dumps(summary, indent=2, sort_keys=True))
        return
    if args.command == "summarize":
        records = load_jsonl(args.runs_jsonl)
        summary = compute_summary(records)
        output = Path(args.output)
        output.parent.mkdir(parents=True, exist_ok=True)
        with output.open("w", encoding="utf-8") as handle:
            json.dump(summary, handle, indent=2, sort_keys=True)
        write_analysis_artifacts(records, output.parent, summary)
        print(json.dumps(summary, indent=2, sort_keys=True))
        return
    if args.command == "plot":
        output = write_summary_figure(args.summary, args.output)
        print(output)
        return
    if args.command == "check-ollama":
        config = OllamaConfig(
            base_url=args.base_url,
            model=args.model,
            timeout_s=args.timeout_s,
            max_retries=args.max_retries,
        )
        client = OllamaClient(config)
        try:
            result = client.check_connection()
        except OllamaUnavailableError as exc:
            raise SystemExit(str(exc)) from exc
        print(json.dumps(result, indent=2, sort_keys=True))
        return
    raise ValueError(f"Unknown command {args.command}")


if __name__ == "__main__":
    main()