scripts.py

#!/usr/bin/env python3
"""Utility scripts for seeding, benchmarking, exporting, and resetting the demo."""
from __future__ import annotations

import argparse, csv, json, shutil, sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
from backend.main import compare  # noqa: E402

DATA = ROOT / "data"
OUT = ROOT / "output"


def seed_demo_data():
    (DATA / "sources").mkdir(parents=True, exist_ok=True)
    (DATA / "benchmarks").mkdir(parents=True, exist_ok=True)
    if not (DATA / "sources" / "public_transit.json").exists():
        raise SystemExit("Expected bundled data/sources/public_transit.json to exist.")
    return "Demo data is already seeded."


def generate_benchmark_runs(model: str):
    OUT.mkdir(exist_ok=True)
    payload = compare(model)
    path = OUT / "benchmark_run.json"
    path.write_text(json.dumps(payload, indent=2))
    return str(path)


def export_comparison_results(model: str):
    OUT.mkdir(exist_ok=True)
    payload = compare(model)
    rows = []
    for key in ["python", "other", "weft"]:
        t = payload[key]["telemetry"]
        rows.append({"mode": key, "model": t["model"], "input_tokens": t["input_tokens"], "output_tokens": t["output_tokens"], "estimated_cost_usd": t["estimated_cost_usd"]})
    rows.append({"mode": "weft_savings_vs_python", "model": model, "input_tokens": payload["savings"]["input_tokens_saved_by_weft"], "output_tokens": payload["savings"]["output_tokens_saved_by_weft"], "estimated_cost_usd": payload["savings"]["cost_saved_by_weft_usd"]})
    rows.append({"mode": "weft_savings_vs_other", "model": model, "input_tokens": payload["weft_vs_other_savings"]["input_tokens_saved_by_weft_vs_other"], "output_tokens": payload["weft_vs_other_savings"]["output_tokens_saved_by_weft_vs_other"], "estimated_cost_usd": payload["weft_vs_other_savings"]["cost_saved_by_weft_vs_other_usd"]})
    path = OUT / "comparison_results.csv"
    with path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=rows[0].keys())
        writer.writeheader(); writer.writerows(rows)
    return str(path)


def reset_local_caches():
    for folder in [DATA / "cached", OUT]:
        if folder.exists(): shutil.rmtree(folder)
        folder.mkdir(parents=True, exist_ok=True)
    return "Local caches and output files reset."


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["seed-demo-data", "benchmark", "export-comparison", "reset-caches"])
    parser.add_argument("--model", default="gpt-4o")
    args = parser.parse_args()
    if args.command == "seed-demo-data": print(seed_demo_data())
    elif args.command == "benchmark": print(generate_benchmark_runs(args.model))
    elif args.command == "export-comparison": print(export_comparison_results(args.model))
    elif args.command == "reset-caches": print(reset_local_caches())

if __name__ == "__main__": main()