scripts/benchmark.py

import argparse
import json
import time
import urllib.error
import urllib.request


def request_json(url: str, method: str = "GET", payload: dict | None = None) -> dict:
    body = None
    headers = {}
    if payload is not None:
        body = json.dumps(payload).encode("utf-8")
        headers["Content-Type"] = "application/json"
    request = urllib.request.Request(url, data=body, headers=headers, method=method)
    try:
        with urllib.request.urlopen(request, timeout=30) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8")
        raise RuntimeError(f"{exc.code}: {detail}") from exc


def main() -> None:
    parser = argparse.ArgumentParser(description="Run an AI Clip Studio API benchmark.")
    parser.add_argument("--api", default="http://localhost:8000")
    parser.add_argument("--youtube-url", required=True)
    parser.add_argument("--language", default="Thai")
    parser.add_argument("--style", default="informative")
    parser.add_argument("--niche", default="education")
    parser.add_argument("--clip-length", type=int, default=60)
    args = parser.parse_args()

    payload = {
        "youtube_url": args.youtube_url,
        "profile": {
            "niche": args.niche,
            "clip_style": args.style,
            "clip_length_seconds": args.clip_length,
            "primary_language": args.language,
            "target_platform": "tiktok",
        },
    }
    started = time.perf_counter()
    job = request_json(f"{args.api}/api/jobs/youtube", "POST", payload)
    while job["status"] in {"queued", "running"}:
        time.sleep(2)
        job = request_json(f"{args.api}/api/jobs/{job['id']}")

    elapsed = round(time.perf_counter() - started, 3)
    print(
        json.dumps(
            {
                "job_id": job["id"],
                "status": job["status"],
                "elapsed_wall_seconds": elapsed,
                "clips": len(job.get("clips", [])),
                "timings": job.get("timings", {}),
                "error": job.get("error"),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()