diff --git a/.github/workflows/benchmark-generate.yml b/.github/workflows/benchmark-generate.yml new file mode 100644 index 0000000000..cf3c3c0319 --- /dev/null +++ b/.github/workflows/benchmark-generate.yml @@ -0,0 +1,209 @@ +name: "Benchmark: Generate" +run-name: "Benchmark: ${{ inputs.library }} for ${{ inputs.specification_id }} (${{ inputs.models }})" + +# LLM benchmark generation (issue #6913): generate the same (spec, library) +# implementation with several models served by Google Vertex AI — Gemini, +# Claude on Vertex, and Model Garden partner models — under the repo's +# existing Workload Identity Federation auth. Adding a model to the +# comparison is just another id in the `models` input; no new secrets or +# integrations. Results (code, renders, raw responses, result.yaml) are +# uploaded as a workflow artifact; nothing is committed to the repo and the +# plots/ catalog is untouched. + +on: + workflow_dispatch: + inputs: + specification_id: + description: "Specification ID (e.g., scatter-basic)" + required: true + type: string + library: + description: "Library to benchmark (Python libraries only in v1)" + required: true + type: choice + default: 'matplotlib' + options: + - matplotlib + - seaborn + - plotly + - bokeh + - altair + - plotnine + - pygal + - letsplot + models: + description: "Comma-separated Vertex AI model ids (gemini-*, claude-*@, meta/*, mistralai/*, …)" + required: true + type: string + default: 'google/gemini-2.5-pro,anthropic/claude-sonnet-4-5@20250929' + max_attempts: + description: "Generation attempts per model (render errors are fed back)" + required: false + type: string + default: '3' + location: + description: "Vertex region for Gemini / Model Garden models" + required: false + type: string + default: 'us-central1' + anthropic_location: + description: "Vertex region for Claude models" + required: false + type: string + default: 'us-east5' + +# One benchmark per (spec, library) at a time; a second dispatch queues. +concurrency: + group: benchmark-generate-${{ inputs.specification_id }}-${{ inputs.library }} + cancel-in-progress: false + +jobs: + benchmark: + runs-on: ubuntu-latest + timeout-minutes: 120 + permissions: + contents: read + id-token: write + + steps: + - name: Checkout repository + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - name: Validate specification exists + env: + SPEC_ID: ${{ inputs.specification_id }} + run: | + if [ ! -f "plots/${SPEC_ID}/specification.md" ]; then + echo "::error::Specification not found: plots/${SPEC_ID}/specification.md" + exit 1 + fi + + - name: Set up Python + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 + with: + python-version: '3.13' + + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + + - name: Install dependencies + env: + LIBRARY: ${{ inputs.library }} + run: | + uv venv .venv + source .venv/bin/activate + # Same per-library extras impl-generate uses, plus the Vertex client's auth dep. + uv pip install -e ".[lib-${LIBRARY}]" pillow pyyaml google-auth + + - name: Authenticate to GCP + uses: google-github-actions/auth@7c6bc770dae815cd3e89ee6cdf493a5fab2cc093 # v3 + with: + project_id: anyplot + workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} + + - name: Run benchmark generation per model + env: + SPEC_ID: ${{ inputs.specification_id }} + LIBRARY: ${{ inputs.library }} + MODELS: ${{ inputs.models }} + MAX_ATTEMPTS: ${{ inputs.max_attempts }} + LOCATION: ${{ inputs.location }} + ANTHROPIC_LOCATION: ${{ inputs.anthropic_location }} + GOOGLE_CLOUD_PROJECT: anyplot + run: | + set -u + succeeded=0 + failed=0 + + IFS=',' read -ra MODEL_LIST <<< "$MODELS" + for MODEL in "${MODEL_LIST[@]}"; do + MODEL=$(echo "$MODEL" | xargs) # trim whitespace + [ -z "$MODEL" ] && continue + echo "::group::Model: ${MODEL}" + if .venv/bin/python -m automation.benchmark.generate \ + --spec-id "$SPEC_ID" \ + --library "$LIBRARY" \ + --model "$MODEL" \ + --location "$LOCATION" \ + --anthropic-location "$ANTHROPIC_LOCATION" \ + --max-attempts "$MAX_ATTEMPTS" \ + --output-dir benchmark-results; then + succeeded=$((succeeded + 1)) + else + echo "::warning::Benchmark generation failed for model ${MODEL}" + failed=$((failed + 1)) + fi + echo "::endgroup::" + done + + echo "::notice::Benchmark done: ${succeeded} succeeded, ${failed} failed" + # Fail the job only when NO model produced a working implementation — + # individual model failures are a benchmark signal, not a CI error. + if [ "$succeeded" -eq 0 ]; then + exit 1 + fi + + - name: Write job summary + if: always() + run: | + # Guard: on early failure (e.g. invalid spec id, install failure) the + # venv or results tree may not exist — don't obscure the real error. + if [ ! -x .venv/bin/python ] || [ ! -d benchmark-results ]; then + echo "## Benchmark results" >> "$GITHUB_STEP_SUMMARY" + echo "" >> "$GITHUB_STEP_SUMMARY" + echo "Workflow failed before any benchmark results were produced." >> "$GITHUB_STEP_SUMMARY" + exit 0 + fi + .venv/bin/python - <<'PY' + import os + import pathlib + + import yaml + + rows = [] + for result_file in sorted(pathlib.Path("benchmark-results").rglob("result.yaml")): + data = yaml.safe_load(result_file.read_text(encoding="utf-8")) or {} + rows.append(data) + + lines = ["## Benchmark results", ""] + if rows: + lines += [ + "| Model | Provider | Success | Attempts | Canvas OK | Gen time (s) | In tokens | Out tokens |", + "|---|---|---|---|---|---|---|---|", + ] + for r in rows: + lines.append( + "| {model} | {provider} | {success} | {attempts}/{max_attempts} | {canvas_ok} " + "| {generation_seconds} | {input_tokens} | {output_tokens} |".format(**{ + k: r.get(k, "?") + for k in ( + "model", "provider", "success", "attempts", "max_attempts", + "canvas_ok", "generation_seconds", "input_tokens", "output_tokens", + ) + }) + ) + errors = [(r.get("model"), r.get("error")) for r in rows if r.get("error")] + if errors: + lines += ["", "### Errors", ""] + for model, error in errors: + # Errors can be multi-line tracebacks; keep each bullet on + # one line so the Markdown list survives, and cap length. + flat = " ".join(str(error).split()) + if len(flat) > 500: + flat = flat[:500] + "…" + lines.append(f"- **{model}**: {flat}") + else: + lines.append("No result.yaml files were produced.") + + with open(os.environ["GITHUB_STEP_SUMMARY"], "a", encoding="utf-8") as handle: + handle.write("\n".join(lines) + "\n") + PY + + - name: Upload benchmark artifact + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: benchmark-${{ inputs.specification_id }}-${{ inputs.library }} + path: benchmark-results/ + if-no-files-found: warn + retention-days: 30 diff --git a/CHANGELOG.md b/CHANGELOG.md index 3840b5adab..4cd36679fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,18 @@ aggregate instead: an italic *Catalog* line at the end of the version section an ## [Unreleased] +### Added + +- **LLM benchmark generation on Vertex AI (foundation for #6913)** — new + `benchmark-generate.yml` workflow plus `automation/benchmark/` module generate the same + (spec, library) implementation with multiple models served by Google Vertex AI — Gemini, + Claude on Vertex, and Model Garden partners (Llama, Mistral, …) — under the existing + Workload Identity Federation auth, so adding a model to the comparison is just another + model id. Single-shot generation with render-error feedback keeps the comparison + model-neutral; per-model `result.yaml` records provider, attempts, latency, tokens, and + the canvas gate. Results are workflow artifacts only (nothing lands in `plots/`); Python + libraries in v1. Docs: `docs/workflows/benchmark.md` (#9638). + ### Fixed - **Global keyboard shortcuts no longer hijack focused elements on `/plots`** — the diff --git a/automation/benchmark/__init__.py b/automation/benchmark/__init__.py new file mode 100644 index 0000000000..d9862dbbc0 --- /dev/null +++ b/automation/benchmark/__init__.py @@ -0,0 +1,7 @@ +"""Vertex AI multi-model benchmark generation (issue #6913). + +Generates plot implementations for a (spec, library) pair with arbitrary +models served by Google Vertex AI — Gemini, Claude on Vertex, and Model +Garden partner models — under a single GCP authentication, so many models +can be benchmarked against the same specification and rubric. +""" diff --git a/automation/benchmark/generate.py b/automation/benchmark/generate.py new file mode 100644 index 0000000000..9bf12f22a3 --- /dev/null +++ b/automation/benchmark/generate.py @@ -0,0 +1,187 @@ +"""CLI: generate one benchmark implementation with a Vertex AI model. + +Usage (from the repo root, with the target library's extras installed): + + python -m automation.benchmark.generate \\ + --spec-id scatter-basic \\ + --library matplotlib \\ + --model google/gemini-2.5-pro \\ + --output-dir benchmark-results + +Writes ``////`` containing the +generated code, both theme renders, each raw model response, and a +``result.yaml`` with the accounting (provider, attempts, latency, tokens, +canvas gate) that the future /benchmark leaderboard aggregates. +""" + +from __future__ import annotations + +import argparse +import datetime +import os +import re +import sys +from pathlib import Path + +import yaml + +from automation.benchmark.prompting import build_generation_prompt, build_repair_prompt, extract_code_block +from automation.benchmark.runner import run_python_implementation +from automation.benchmark.vertex_client import VertexClient, resolve_provider +from core.constants import LIBRARIES_METADATA + + +BENCHMARK_RESULT_VERSION = 1 + +PYTHON_LIBRARIES = sorted(lib["id"] for lib in LIBRARIES_METADATA if lib["language_id"] == "python") + + +def slugify_model(model: str) -> str: + """Filesystem-safe directory name for a Vertex model id.""" + return re.sub(r"[^a-z0-9]+", "-", model.lower()).strip("-") + + +def accumulate_tokens(current: int | None, reported: int | None) -> int | None: + """Add a reported token count, keeping ``None`` when nothing was reported. + + ``None`` means "the provider never reported usage" and must stay ``None`` + across attempts — recording 0 would look like a free run and skew the + benchmark comparison. + """ + if reported is None: + return current + return (current or 0) + reported + + +def _positive_int(value: str) -> int: + number = int(value) + if number < 1: + raise argparse.ArgumentTypeError(f"must be >= 1, got {number}") + return number + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate one plot implementation with a Vertex AI model") + parser.add_argument("--spec-id", required=True, help="Specification id, e.g. scatter-basic") + parser.add_argument("--library", required=True, choices=PYTHON_LIBRARIES, help="Target library (Python only in v1)") + parser.add_argument( + "--model", + required=True, + help="Vertex model id: gemini-*, claude-*@, or publisher-prefixed (meta/…, mistralai/…)", + ) + parser.add_argument( + "--project", + default=os.getenv("GOOGLE_CLOUD_PROJECT", "anyplot"), + help="GCP project id (default: $GOOGLE_CLOUD_PROJECT or 'anyplot')", + ) + parser.add_argument("--location", default="us-central1", help="Vertex region for Gemini / Model Garden") + parser.add_argument("--anthropic-location", default="us-east5", help="Vertex region for Claude models") + parser.add_argument("--output-dir", default="benchmark-results", help="Root directory for benchmark output") + parser.add_argument( + "--max-attempts", type=_positive_int, default=3, help="Generation attempts (errors are fed back)" + ) + parser.add_argument("--max-tokens", type=_positive_int, default=16000, help="Max output tokens per model call") + parser.add_argument("--temperature", type=float, default=0.2, help="Sampling temperature") + parser.add_argument("--render-timeout", type=_positive_int, default=300, help="Seconds allowed per theme render") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + repo_root = Path(__file__).resolve().parents[2] + + provider, normalized_model = resolve_provider(args.model) + # Slug on the normalized id so 'gemini-2.5-pro' and 'google/gemini-2.5-pro' + # (or 'anthropic/claude-*' and bare 'claude-*') land in the same folder. + out_dir = Path(args.output_dir) / args.spec_id / args.library / slugify_model(normalized_model) + out_dir.mkdir(parents=True, exist_ok=True) + + system, base_prompt = build_generation_prompt(repo_root, args.spec_id, args.library) + client = VertexClient(project=args.project, location=args.location, anthropic_location=args.anthropic_location) + + result: dict = { + "benchmark_result_version": BENCHMARK_RESULT_VERSION, + "specification_id": args.spec_id, + "library": args.library, + "language": "python", + "model": args.model, + "normalized_model": normalized_model, + "provider": provider, + "project": args.project, + "created": datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ"), + "success": False, + "attempts": 0, + "max_attempts": args.max_attempts, + "generation_seconds": 0.0, + # null until a provider actually reports usage — 0 would read as + # "free" and skew comparisons when a publisher omits usage data. + "input_tokens": None, + "output_tokens": None, + "canvas_ok": None, + # Whether the final response honored the "one fenced code block" + # contract; extraction is deliberately lenient, so this is recorded + # instead of enforced. + "code_fenced": None, + "error": None, + } + + prompt = base_prompt + code = "" + for attempt in range(1, args.max_attempts + 1): + result["attempts"] = attempt + print(f"[benchmark] {args.spec_id}/{args.library} · {args.model} · attempt {attempt}/{args.max_attempts}") + + try: + generation = client.generate( + normalized_model, system, prompt, max_tokens=args.max_tokens, temperature=args.temperature + ) + except Exception as exc: # noqa: BLE001 — any provider error ends this model's run, recorded in result.yaml + result["error"] = f"generation failed: {exc}" + break + + result["generation_seconds"] = round(result["generation_seconds"] + generation.latency_seconds, 3) + result["input_tokens"] = accumulate_tokens(result["input_tokens"], generation.input_tokens) + result["output_tokens"] = accumulate_tokens(result["output_tokens"], generation.output_tokens) + (out_dir / f"attempt-{attempt}-response.md").write_text(generation.text, encoding="utf-8") + + try: + code = extract_code_block(generation.text) + result["code_fenced"] = "```" in generation.text + except ValueError as exc: + result["error"] = str(exc) + # Feed back THIS attempt's raw response — there is no code block to + # show, and reusing an earlier attempt's extracted code here would + # mislead the repair prompt. + prompt = build_repair_prompt( + base_prompt, + generation.text.strip()[-4000:] or "(empty response)", + f"{exc} — your previous response is shown above in place of code", + ) + continue + + code_file = out_dir / f"{args.library}.py" + code_file.write_text(code + "\n", encoding="utf-8") + + render = run_python_implementation(code_file, workdir=out_dir, timeout=args.render_timeout) + if render.success: + result["success"] = True + result["canvas_ok"] = render.canvas_ok + result["error"] = None + break + + result["error"] = render.error + prompt = build_repair_prompt(base_prompt, code, render.error or "unknown render error") + + result_file = out_dir / "result.yaml" + with result_file.open("w", encoding="utf-8") as handle: + handle.write(f"# Benchmark result for {args.spec_id} / {args.library} / {args.model}\n") + handle.write("# Auto-generated by automation.benchmark.generate\n\n") + yaml.dump(result, handle, default_flow_style=False, sort_keys=False, allow_unicode=True) + + status = "OK" if result["success"] else "FAILED" + print(f"[benchmark] {status}: {result_file} (attempts={result['attempts']}, canvas_ok={result['canvas_ok']})") + return 0 if result["success"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/automation/benchmark/prompting.py b/automation/benchmark/prompting.py new file mode 100644 index 0000000000..d89f014628 --- /dev/null +++ b/automation/benchmark/prompting.py @@ -0,0 +1,124 @@ +"""Prompt assembly and response parsing for benchmark generation. + +Every model gets exactly the same context the Claude pipeline reads for a +real implementation — base generation rules, the Imprint style guide, the +library-specific rules, and the spec — so quality scores stay comparable +across models. Unlike the agentic pipeline, benchmark generation is +single-shot: the model must answer with one self-contained code block. +""" + +from __future__ import annotations + +import re +from pathlib import Path + + +# The single-shot contract replaces the agentic workflow steps (run, self-check, +# commit) that impl-generate-claude.md drives interactively. +_OUTPUT_CONTRACT = """ +## Output contract (benchmark mode — read carefully) + +You are running in single-shot benchmark mode. There is no shell and no +second turn: your entire answer must be ONE fenced code block containing a +complete, self-contained {library} implementation in Python. The script will +be executed for you, once per theme, in a working directory it may write to. + +Hard requirements: + +1. Answer with exactly one ```python code block and nothing else — no prose + before or after, no second block. +2. The script reads the theme from the environment: + `theme = os.getenv("ANYPLOT_THEME", "light")` and renders that theme. +3. The script saves its output as `plot-{{theme}}.png` (i.e. `plot-light.png` + or `plot-dark.png`) in the current working directory. +4. The saved PNG must be exactly 3200x1800 px (landscape) or 2400x2400 px + (square) — follow the "Canvas" rules in the library section below. +5. Only use {library}, the Python standard library, and numpy / pandas / + scipy / scikit-learn / statsmodels. No network access; do not read any + local files — the PNG you save is the only file system interaction. +6. Keep the KISS structure: imports -> data -> plot -> save. +""" + +_REPAIR_SECTION = """ +## Previous attempt failed + +Your previous implementation (below) failed to render. Fix the error and +answer again with ONE complete ```python code block following the same +output contract. + +### Previous code + +```python +{previous_code} +``` + +### Error output + +``` +{error_text} +``` +""" + + +def build_generation_prompt(repo_root: Path, spec_id: str, library: str) -> tuple[str, str]: + """Return ``(system, user)`` prompts for one (spec, library) pair. + + Raises ``FileNotFoundError`` if the spec or any prompt file is missing. + """ + base_rules = _read(repo_root / "prompts" / "plot-generator.md") + style_guide = _read(repo_root / "prompts" / "default-style-guide.md") + library_rules = _read(repo_root / "prompts" / "library" / f"{library}.md") + specification = _read(repo_root / "plots" / spec_id / "specification.md") + + system = ( + "You are an expert data-visualization engineer. You write clean, self-contained, " + "gallery-quality plot scripts that follow the provided style rules exactly." + ) + user = "\n\n---\n\n".join( + [ + f"# Task\n\nImplement the plot specification below using **{library}** (Python).", + f"# Base generation rules\n\n{base_rules}", + f"# Style guide\n\n{style_guide}", + f"# Library rules: {library}\n\n{library_rules}", + f"# Specification: {spec_id}\n\n{specification}", + _OUTPUT_CONTRACT.format(library=library), + ] + ) + return system, user + + +def build_repair_prompt(base_user_prompt: str, previous_code: str, error_text: str) -> str: + """Append the failed attempt + error to the base prompt for a retry.""" + return ( + base_user_prompt + + "\n\n---\n" + + _REPAIR_SECTION.format(previous_code=previous_code.strip(), error_text=error_text.strip()[-4000:]) + ) + + +def extract_code_block(text: str) -> str: + """Extract the implementation code from a model response. + + Takes the longest fenced code block (models occasionally emit a short + usage snippet next to the real implementation). Falls back to the raw + text when there is no fence but the response already looks like code — + deliberate leniency: the benchmark measures plot quality, not markdown + discipline, so runnable code is never discarded over a missing fence. + Contract compliance is still measurable: ``result.yaml`` records whether + the response was fenced (``code_fenced``). + """ + blocks = re.findall(r"```(?:[a-zA-Z0-9_+-]*)\n(.*?)```", text, flags=re.DOTALL) + if blocks: + return max(blocks, key=len).strip() + + stripped = text.strip() + if "import " in stripped and "```" not in stripped: + return stripped + + raise ValueError("No fenced code block found in model response") + + +def _read(path: Path) -> str: + if not path.is_file(): + raise FileNotFoundError(f"Required prompt input missing: {path}") + return path.read_text(encoding="utf-8").strip() diff --git a/automation/benchmark/runner.py b/automation/benchmark/runner.py new file mode 100644 index 0000000000..2726d8bfe8 --- /dev/null +++ b/automation/benchmark/runner.py @@ -0,0 +1,112 @@ +"""Render benchmark implementations and validate their output. + +v1 supports Python libraries only — they run with the same interpreter that +runs the benchmark, so a plain ``uv pip install -e ".[lib-]"`` is +the whole runtime. R / Julia / JavaScript need their own runtimes plus the +browser render harness; they reuse the existing setup actions when the +benchmark grows past Python. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path + + +THEMES = ("light", "dark") + +# Same targets and tolerance as the post-render canvas gate in impl-review. +CANVAS_TARGETS = ((3200, 1800), (2400, 2400)) +CANVAS_TOLERANCE_PX = 16 + + +@dataclass +class RenderResult: + """Outcome of rendering one implementation in both themes.""" + + success: bool + error: str | None = None + canvas_ok: bool | None = None + images: dict[str, Path] = field(default_factory=dict) + + +# The render subprocess executes LLM-generated code, so it must NOT inherit +# the caller's environment (in CI that includes the GCP/WIF credentials and +# GitHub token). Only these harmless variables pass through; HOME is remapped +# to a scratch directory so ~/.config (gcloud credentials, tool configs) is +# out of reach too. +_RENDER_ENV_ALLOWLIST = ("PATH", "LANG", "LC_ALL", "TMPDIR", "TEMP", "TMP") + + +def render_env(theme: str, home: Path) -> dict[str, str]: + """Minimal, credential-free environment for running generated code.""" + env = {key: value for key in _RENDER_ENV_ALLOWLIST if (value := os.environ.get(key)) is not None} + env["HOME"] = str(home) + env["ANYPLOT_THEME"] = theme + env["MPLBACKEND"] = "Agg" + return env + + +def run_python_implementation(code_file: Path, workdir: Path, timeout: int = 300) -> RenderResult: + """Run ``code_file`` once per theme inside ``workdir``. + + Success requires both runs to exit 0 and both theme PNGs to exist. + ``canvas_ok`` records (but does not fail on) the canvas-dimension gate — + the benchmark wants to *measure* how often models hit the contract. + """ + images: dict[str, Path] = {} + scratch_home = workdir / ".render-home" + scratch_home.mkdir(exist_ok=True) + for theme in THEMES: + # Attempts share the workdir — a stale PNG from a previous attempt + # must never make this run look like it rendered successfully. + (workdir / f"plot-{theme}.png").unlink(missing_ok=True) + for theme in THEMES: + env = render_env(theme, home=scratch_home) + try: + # -I (isolated mode): ignore PYTHON* env vars, user site-packages, + # and the script directory on sys.path — defense in depth on top of + # the env allowlist, since this runs LLM-generated code. The venv's + # site-packages still resolve via pyvenv.cfg next to the executable. + completed = subprocess.run( + [sys.executable, "-I", str(code_file)], + cwd=workdir, + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired: + return RenderResult(success=False, error=f"{theme} render timed out after {timeout}s", images=images) + + if completed.returncode != 0: + error = (completed.stderr or completed.stdout or "").strip()[-4000:] + return RenderResult( + success=False, error=f"{theme} render exited {completed.returncode}:\n{error}", images=images + ) + + png = workdir / f"plot-{theme}.png" + if not png.is_file(): + return RenderResult( + success=False, error=f"{theme} render exited 0 but wrote no plot-{theme}.png", images=images + ) + images[theme] = png + + # Both themes must hit the canvas contract — a dark render with layout + # drift is a contract miss even when the light render is on target. + return RenderResult(success=True, canvas_ok=all(check_canvas(png) for png in images.values()), images=images) + + +def check_canvas(png_path: Path) -> bool: + """True when the PNG is within tolerance of a canonical canvas size.""" + from PIL import Image + + with Image.open(png_path) as image: + width, height = image.size + return any( + abs(width - target_w) <= CANVAS_TOLERANCE_PX and abs(height - target_h) <= CANVAS_TOLERANCE_PX + for target_w, target_h in CANVAS_TARGETS + ) diff --git a/automation/benchmark/vertex_client.py b/automation/benchmark/vertex_client.py new file mode 100644 index 0000000000..d39d72e46e --- /dev/null +++ b/automation/benchmark/vertex_client.py @@ -0,0 +1,219 @@ +"""Vertex AI client with provider routing for multi-model benchmarking. + +One GCP authentication (Application Default Credentials — locally via +``gcloud auth application-default login``, in GitHub Actions via Workload +Identity Federation) reaches every model family available on Vertex AI: + +- **Anthropic Claude** models (``claude-*`` / ``anthropic/claude-*``) go + through the Anthropic SDK's Vertex transport (``AnthropicVertex``), + because Claude on Vertex speaks the Anthropic message schema via + rawPredict rather than the shared chat-completions surface. +- **Everything else** — Google Gemini (``gemini-*`` / ``google/gemini-*``) + and Model Garden model-as-a-service partners (``meta/*``, ``mistralai/*``, + ``qwen/*``, ``deepseek-ai/*``, ``openai/*`` …) — goes through Vertex AI's + OpenAI-compatible chat-completions endpoint, which accepts the + publisher-prefixed model id directly. + +This is what makes the benchmark cheap to extend: adding a model is a new +model id string, not a new integration. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import httpx + + +PROVIDER_ANTHROPIC = "anthropic" +PROVIDER_OPENAPI = "openapi" + +_CLOUD_PLATFORM_SCOPE = "https://www.googleapis.com/auth/cloud-platform" + + +@dataclass +class GenerationResult: + """One model response plus the accounting the benchmark persists.""" + + text: str + model: str + provider: str + latency_seconds: float + input_tokens: int | None = None + output_tokens: int | None = None + + +def resolve_provider(model: str) -> tuple[str, str]: + """Route a model id to its Vertex provider. + + Returns ``(provider, normalized_model)`` where ``normalized_model`` is the + id the provider endpoint expects: + + - ``anthropic/claude-sonnet-4-5@20250929`` → ``("anthropic", "claude-sonnet-4-5@20250929")`` + - ``claude-sonnet-4-5@20250929`` → ``("anthropic", "claude-sonnet-4-5@20250929")`` + - ``gemini-2.5-pro`` → ``("openapi", "google/gemini-2.5-pro")`` + - ``google/gemini-2.5-pro`` → ``("openapi", "google/gemini-2.5-pro")`` + - ``meta/llama-3.3-70b-instruct-maas`` → ``("openapi", "meta/llama-3.3-70b-instruct-maas")`` + + Any ``publisher/model`` id is accepted and routed to the OpenAI-compatible + endpoint — Model Garden's publisher list keeps growing, so the prefix is + deliberately not validated against a fixed set; an unknown publisher fails + at call time with Vertex's own error. Only bare ids without a recognizable + family (``claude-*`` / ``gemini-*``) or publisher prefix raise ``ValueError``. + """ + normalized = model.strip() + if not normalized: + raise ValueError("Model id must not be empty") + + if normalized.startswith("anthropic/"): + return PROVIDER_ANTHROPIC, normalized.removeprefix("anthropic/") + if normalized.startswith("claude-"): + return PROVIDER_ANTHROPIC, normalized + if normalized.startswith("gemini-"): + return PROVIDER_OPENAPI, f"google/{normalized}" + if "/" in normalized: + return PROVIDER_OPENAPI, normalized + + raise ValueError( + f"Cannot route model id '{model}': use 'claude-*' / 'anthropic/*' for Claude on Vertex, " + "'gemini-*' / 'google/*' for Gemini, or a publisher-prefixed Model Garden id like " + "'meta/llama-3.3-70b-instruct-maas'" + ) + + +class VertexClient: + """Minimal text-generation client over Vertex AI with provider routing. + + ``location`` serves the OpenAI-compatible endpoint (Gemini + Model Garden); + ``anthropic_location`` serves Claude, which Google hosts in a different + set of regions (default ``us-east5``). + """ + + def __init__( + self, project: str, location: str = "us-central1", anthropic_location: str = "us-east5", timeout: float = 600.0 + ) -> None: + self.project = project + self.location = location + self.anthropic_location = anthropic_location + self.timeout = timeout + self._credentials = None + + def generate( + self, model: str, system: str, prompt: str, max_tokens: int = 16000, temperature: float = 0.2 + ) -> GenerationResult: + """Send one system+user exchange to ``model`` and return the reply.""" + provider, normalized = resolve_provider(model) + started = time.monotonic() + if provider == PROVIDER_ANTHROPIC: + result = self._generate_anthropic(normalized, system, prompt, max_tokens, temperature) + else: + result = self._generate_openapi(normalized, system, prompt, max_tokens, temperature) + result.latency_seconds = round(time.monotonic() - started, 3) + return result + + # ------------------------------------------------------------------ + # Anthropic Claude on Vertex (rawPredict via the Anthropic SDK) + # ------------------------------------------------------------------ + + def _generate_anthropic( + self, model: str, system: str, prompt: str, max_tokens: int, temperature: float + ) -> GenerationResult: + from anthropic import AnthropicVertex + + client = AnthropicVertex(project_id=self.project, region=self.anthropic_location, timeout=self.timeout) + message = client.messages.create( + model=model, + max_tokens=max_tokens, + temperature=temperature, + system=system, + messages=[{"role": "user", "content": prompt}], + ) + text = "".join(block.text for block in message.content if block.type == "text") + return GenerationResult( + text=text, + model=model, + provider=PROVIDER_ANTHROPIC, + latency_seconds=0.0, + input_tokens=getattr(message.usage, "input_tokens", None), + output_tokens=getattr(message.usage, "output_tokens", None), + ) + + # ------------------------------------------------------------------ + # Gemini + Model Garden via the OpenAI-compatible endpoint + # ------------------------------------------------------------------ + + def _openapi_url(self) -> str: + # The "global" location drops the regional host prefix. + host = ( + "aiplatform.googleapis.com" if self.location == "global" else f"{self.location}-aiplatform.googleapis.com" + ) + return f"https://{host}/v1/projects/{self.project}/locations/{self.location}/endpoints/openapi/chat/completions" + + def _access_token(self) -> str: + try: + import google.auth + from google.auth.transport.requests import Request + except ImportError as exc: # google-auth ships transitively (google-cloud-storage); guard minimal envs + raise RuntimeError( + "google-auth is required for Vertex AI calls but is not installed. " + 'Install it with `uv pip install google-auth` (or `uv pip install -e ".[lib-]" google-auth`).' + ) from exc + + if self._credentials is None: + self._credentials, _ = google.auth.default(scopes=[_CLOUD_PLATFORM_SCOPE]) + if not self._credentials.valid: + self._credentials.refresh(Request()) + token = self._credentials.token + if not token: + raise RuntimeError("Could not obtain a GCP access token from Application Default Credentials") + return str(token) + + def _generate_openapi( + self, model: str, system: str, prompt: str, max_tokens: int, temperature: float + ) -> GenerationResult: + payload = { + "model": model, + "max_tokens": max_tokens, + "temperature": temperature, + "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}], + } + response = httpx.post( + self._openapi_url(), + json=payload, + headers={"Authorization": f"Bearer {self._access_token()}", "Content-Type": "application/json"}, + timeout=self.timeout, + ) + if response.status_code != 200: + raise RuntimeError( + f"Vertex chat-completions call for '{model}' failed with HTTP {response.status_code}: " + f"{response.text[:2000]}" + ) + body = response.json() + text = _extract_openapi_text(body) + usage = body.get("usage") or {} + return GenerationResult( + text=text, + model=model, + provider=PROVIDER_OPENAPI, + latency_seconds=0.0, + input_tokens=usage.get("prompt_tokens"), + output_tokens=usage.get("completion_tokens"), + ) + + +def _extract_openapi_text(body: dict) -> str: + """Pull the assistant text out of an OpenAI-style chat-completions body. + + ``content`` is normally a string, but some Model Garden publishers return + the content-parts list form — accept both. + """ + choices = body.get("choices") or [] + if not choices: + raise RuntimeError(f"Vertex chat-completions response contains no choices: {str(body)[:500]}") + content = (choices[0].get("message") or {}).get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join(part.get("text", "") for part in content if isinstance(part, dict)) + raise RuntimeError(f"Unexpected content shape in Vertex chat-completions response: {type(content).__name__}") diff --git a/docs/workflows/benchmark.md b/docs/workflows/benchmark.md new file mode 100644 index 0000000000..7502b77517 --- /dev/null +++ b/docs/workflows/benchmark.md @@ -0,0 +1,89 @@ +# LLM Benchmark Generation (Vertex AI) + +> Foundation for [#6913 — LLM Benchmark: Compare model performance across specs and libraries](https://github.com/MarkusNeusinger/anyplot/issues/6913). + +`benchmark-generate.yml` generates the same (spec, library) implementation with several +different LLMs served by **Google Vertex AI**, so models can be compared on identical +inputs. One GCP authentication (the repo's existing Workload Identity Federation) covers +every model family — adding a model to the comparison is just another id in the `models` +input, no new secrets or SDK integrations: + +| Model family | Example id | Vertex surface | +|---|---|---| +| Google Gemini | `gemini-2.5-pro` or `google/gemini-2.5-pro` | OpenAI-compatible chat-completions endpoint | +| Anthropic Claude | `claude-sonnet-4-5@20250929` or `anthropic/claude-sonnet-4-5@20250929` | Anthropic SDK Vertex transport (`AnthropicVertex`) | +| Model Garden MaaS partners | `meta/llama-3.3-70b-instruct-maas`, `mistralai/mistral-large-2411`, `qwen/…`, `deepseek-ai/…` | OpenAI-compatible chat-completions endpoint | + +## How it differs from impl-generate + +`impl-generate` runs Claude Code **agentically** (it iterates in a shell, renders, self-checks, +commits). Benchmark generation is deliberately **single-shot with error feedback**: every model +gets the exact same context the pipeline reads (base rules, Imprint style guide, library rules, +spec) and must answer with one self-contained code block. Render errors are fed back for up to +`max_attempts` tries. That keeps the comparison model-neutral — no model gets an agent harness +another lacks. + +Benchmark output is **never** committed to `plots/` and never touches the catalog, metadata, or +GCS. Results are uploaded as a workflow artifact and summarized in the job summary. + +## Running a benchmark + +```bash +gh workflow run benchmark-generate.yml \ + -f specification_id=scatter-basic \ + -f library=matplotlib \ + -f models='google/gemini-2.5-pro,anthropic/claude-sonnet-4-5@20250929,meta/llama-3.3-70b-instruct-maas' +``` + +Inputs: + +| Input | Default | Notes | +|---|---|---| +| `specification_id` | — | Any spec under `plots/` | +| `library` | `matplotlib` | **Python libraries only in v1** (they run on the workflow's own interpreter); R/Julia/JS reuse the existing setup actions in a later iteration | +| `models` | Gemini 2.5 Pro + Claude Sonnet 4.5 | Comma-separated Vertex model ids | +| `max_attempts` | `3` | Render errors are fed back to the model between attempts | +| `location` | `us-central1` | Region for Gemini / Model Garden | +| `anthropic_location` | `us-east5` | Claude on Vertex lives in its own region set | + +Locally (with `gcloud auth application-default login`): + +```bash +uv pip install -e ".[lib-matplotlib]" google-auth # google-auth: ADC tokens for the Vertex client +python -m automation.benchmark.generate \ + --spec-id scatter-basic --library matplotlib \ + --model gemini-2.5-pro --output-dir benchmark-results +``` + +(`google-auth` usually arrives transitively via `google-cloud-storage`, but the +Vertex client imports it directly — install it explicitly to be safe.) + +## Output layout + +``` +benchmark-results//// +├── .py # last generated implementation +├── plot-light.png # theme renders (when successful) +├── plot-dark.png +├── attempt-N-response.md # raw model response per attempt +└── result.yaml # accounting: provider, success, attempts, + # latency, tokens, canvas gate (both themes), + # fenced-output contract compliance, error +``` + +`result.yaml` is the record the future `/benchmark` leaderboard aggregates +(per issue #6913: quality scoring via the existing review rubric, persistence, +and site integration are follow-up phases). + +## GCP prerequisites (one-time) + +The workflow reuses the existing `google-github-actions/auth` WIF setup. For it to reach +Vertex AI, the project additionally needs: + +1. **Vertex AI API enabled**: `gcloud services enable aiplatform.googleapis.com` +2. **IAM**: the WIF service account needs `roles/aiplatform.user` +3. **Partner models enabled once in the console** (Claude, Llama, Mistral, … each have an + "Enable" step in Model Garden; Gemini needs no enablement) + +Model availability is regional — if a call returns 404 for a model you enabled, check the +region inputs against the model's supported regions. diff --git a/docs/workflows/overview.md b/docs/workflows/overview.md index d42f3ba920..b4495b67d0 100644 --- a/docs/workflows/overview.md +++ b/docs/workflows/overview.md @@ -228,3 +228,14 @@ $env:CLI_MODEL_COPILOT_MEDIUM = "gpt-4-turbo" ``` All environment variable names follow the pattern: `CLI_MODEL_{CLI}_{TIER}` + +--- + +## LLM Benchmark (experimental) + +`benchmark-generate.yml` generates the same (spec, library) implementation with multiple +models served by Google Vertex AI — Gemini, Claude on Vertex, and Model Garden partner +models — under the repo's existing Workload Identity Federation auth, so models can be +compared side-by-side on identical inputs (issue #6913). Results are uploaded as workflow +artifacts and summarized in the job summary; nothing is committed to `plots/`. +See [benchmark.md](benchmark.md) for model id routing, inputs, and GCP prerequisites. diff --git a/tests/unit/automation/benchmark/__init__.py b/tests/unit/automation/benchmark/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/automation/benchmark/test_generate.py b/tests/unit/automation/benchmark/test_generate.py new file mode 100644 index 0000000000..88cc0890bf --- /dev/null +++ b/tests/unit/automation/benchmark/test_generate.py @@ -0,0 +1,156 @@ +"""Tests for automation.benchmark.generate CLI helpers and main loop.""" + +import pytest +import yaml + +from automation.benchmark import generate as generate_mod +from automation.benchmark.generate import PYTHON_LIBRARIES, accumulate_tokens, main, parse_args, slugify_model +from automation.benchmark.vertex_client import GenerationResult + + +# Self-contained "generated implementation" that satisfies the runner contract +# without importing any plotting library. +WORKING_CODE = """ +import os +from PIL import Image +theme = os.getenv("ANYPLOT_THEME", "light") +Image.new("RGB", (64, 64)).save(f"plot-{theme}.png") +""" + + +class FakeVertexClient: + """Stands in for VertexClient; yields scripted responses per attempt.""" + + responses: list[str] = [] + + def __init__(self, **kwargs): + self.init_kwargs = kwargs + self._iter = iter(type(self).responses) + + def generate(self, model, system, prompt, max_tokens, temperature): + item = next(self._iter) + if isinstance(item, Exception): + raise item + return GenerationResult( + text=item, model=model, provider="openapi", latency_seconds=0.5, input_tokens=10, output_tokens=5 + ) + + +class TestSlugifyModel: + def test_publisher_prefixed_id(self): + assert slugify_model("google/gemini-2.5-pro") == "google-gemini-2-5-pro" + + def test_versioned_claude_id(self): + assert slugify_model("claude-sonnet-4-5@20250929") == "claude-sonnet-4-5-20250929" + + def test_uppercase_is_lowered(self): + assert slugify_model("MistralAI/Mistral-Large-2411") == "mistralai-mistral-large-2411" + + +class TestPythonLibraries: + def test_only_python_libraries_are_allowed(self): + assert "matplotlib" in PYTHON_LIBRARIES + assert "seaborn" in PYTHON_LIBRARIES + # Non-Python catalogue entries are excluded from benchmark v1 + for non_python in ("ggplot2", "makie", "chartjs", "d3", "echarts", "muix"): + assert non_python not in PYTHON_LIBRARIES + + +class TestAccumulateTokens: + def test_unreported_usage_stays_none(self): + assert accumulate_tokens(None, None) is None + + def test_first_report_starts_from_zero(self): + assert accumulate_tokens(None, 120) == 120 + + def test_reports_accumulate(self): + assert accumulate_tokens(120, 80) == 200 + + def test_missing_report_keeps_prior_total(self): + assert accumulate_tokens(120, None) == 120 + + def test_zero_is_a_real_report_not_unknown(self): + assert accumulate_tokens(None, 0) == 0 + + +class TestParseArgs: + def test_defaults(self): + args = parse_args(["--spec-id", "scatter-basic", "--library", "matplotlib", "--model", "gemini-2.5-pro"]) + assert args.max_attempts == 3 + assert args.location == "us-central1" + assert args.anthropic_location == "us-east5" + assert args.output_dir == "benchmark-results" + + def test_non_python_library_rejected(self): + with pytest.raises(SystemExit): + parse_args(["--spec-id", "scatter-basic", "--library", "ggplot2", "--model", "gemini-2.5-pro"]) + + def test_non_positive_numeric_args_rejected(self): + base = ["--spec-id", "scatter-basic", "--library", "matplotlib", "--model", "gemini-2.5-pro"] + for flag in ("--max-attempts", "--max-tokens", "--render-timeout"): + for bad in ("0", "-1"): + with pytest.raises(SystemExit): + parse_args([*base, flag, bad]) + + +class TestMain: + """End-to-end main() runs against the real repo prompts/spec with a fake client.""" + + def _run(self, tmp_path, monkeypatch, responses, max_attempts=3): + FakeVertexClient.responses = responses + monkeypatch.setattr(generate_mod, "VertexClient", FakeVertexClient) + exit_code = main( + [ + "--spec-id", + "scatter-basic", + "--library", + "matplotlib", + "--model", + "google/gemini-2.5-pro", + "--output-dir", + str(tmp_path), + "--max-attempts", + str(max_attempts), + ] + ) + result_file = tmp_path / "scatter-basic" / "matplotlib" / "google-gemini-2-5-pro" / "result.yaml" + return exit_code, (yaml.safe_load(result_file.read_text(encoding="utf-8")) if result_file.is_file() else None) + + def test_successful_run_writes_result_yaml(self, tmp_path, monkeypatch): + exit_code, result = self._run(tmp_path, monkeypatch, [f"```python\n{WORKING_CODE}\n```"]) + assert exit_code == 0 + assert result["success"] is True + assert result["attempts"] == 1 + assert result["code_fenced"] is True + assert result["canvas_ok"] is False # 64x64 misses the canonical canvas + assert result["input_tokens"] == 10 + assert result["output_tokens"] == 5 + assert result["provider"] == "openapi" + assert result["error"] is None + impl = tmp_path / "scatter-basic" / "matplotlib" / "google-gemini-2-5-pro" / "matplotlib.py" + assert impl.is_file() + + def test_extraction_failure_retries_then_succeeds(self, tmp_path, monkeypatch): + exit_code, result = self._run( + tmp_path, monkeypatch, ["I cannot find any code to write here.", f"```python\n{WORKING_CODE}\n```"] + ) + assert exit_code == 0 + assert result["success"] is True + assert result["attempts"] == 2 + # Token usage accumulated across both attempts + assert result["input_tokens"] == 20 + + def test_provider_error_is_recorded_and_fails(self, tmp_path, monkeypatch): + exit_code, result = self._run(tmp_path, monkeypatch, [RuntimeError("Vertex says no")]) + assert exit_code == 1 + assert result["success"] is False + assert "Vertex says no" in result["error"] + assert result["attempts"] == 1 + + def test_render_failure_exhausts_attempts(self, tmp_path, monkeypatch): + broken = "```python\nraise SystemExit('render boom')\n```" + exit_code, result = self._run(tmp_path, monkeypatch, [broken, broken], max_attempts=2) + assert exit_code == 1 + assert result["success"] is False + assert result["attempts"] == 2 + assert "render boom" in result["error"] diff --git a/tests/unit/automation/benchmark/test_prompting.py b/tests/unit/automation/benchmark/test_prompting.py new file mode 100644 index 0000000000..40437807c3 --- /dev/null +++ b/tests/unit/automation/benchmark/test_prompting.py @@ -0,0 +1,75 @@ +"""Tests for automation.benchmark.prompting.""" + +import pytest + +from automation.benchmark.prompting import build_generation_prompt, build_repair_prompt, extract_code_block + + +@pytest.fixture +def repo_root(tmp_path): + """Minimal repo layout with the four prompt inputs the builder reads.""" + (tmp_path / "prompts" / "library").mkdir(parents=True) + (tmp_path / "plots" / "scatter-basic").mkdir(parents=True) + (tmp_path / "prompts" / "plot-generator.md").write_text("BASE RULES", encoding="utf-8") + (tmp_path / "prompts" / "default-style-guide.md").write_text("STYLE GUIDE", encoding="utf-8") + (tmp_path / "prompts" / "library" / "matplotlib.md").write_text("LIBRARY RULES", encoding="utf-8") + (tmp_path / "plots" / "scatter-basic" / "specification.md").write_text("THE SPEC", encoding="utf-8") + return tmp_path + + +class TestBuildGenerationPrompt: + def test_includes_all_inputs_and_contract(self, repo_root): + system, user = build_generation_prompt(repo_root, "scatter-basic", "matplotlib") + + assert "data-visualization engineer" in system + for fragment in ("BASE RULES", "STYLE GUIDE", "LIBRARY RULES", "THE SPEC"): + assert fragment in user + # Single-shot contract essentials + assert "ANYPLOT_THEME" in user + assert "3200x1800" in user + assert "ONE fenced code block" in user + + def test_missing_spec_raises(self, repo_root): + with pytest.raises(FileNotFoundError, match="specification.md"): + build_generation_prompt(repo_root, "does-not-exist", "matplotlib") + + def test_missing_library_prompt_raises(self, repo_root): + with pytest.raises(FileNotFoundError, match="nosuchlib.md"): + build_generation_prompt(repo_root, "scatter-basic", "nosuchlib") + + +class TestBuildRepairPrompt: + def test_appends_code_and_error(self): + repaired = build_repair_prompt("BASE", "import x", "Traceback: boom") + assert repaired.startswith("BASE") + assert "import x" in repaired + assert "Traceback: boom" in repaired + assert "Previous attempt failed" in repaired + + def test_error_is_truncated_to_tail(self): + long_error = "x" * 10000 + "THE END" + repaired = build_repair_prompt("BASE", "code", long_error) + assert "THE END" in repaired + assert len(repaired) < 10000 + + +class TestExtractCodeBlock: + def test_python_fence(self): + text = "Here you go:\n```python\nimport matplotlib\n```\nEnjoy!" + assert extract_code_block(text) == "import matplotlib" + + def test_bare_fence(self): + assert extract_code_block("```\nimport numpy\n```") == "import numpy" + + def test_longest_block_wins(self): + text = "```python\nprint('short')\n```\n\n```python\nimport pandas\nimport numpy\nprint('real one')\n```" + assert "real one" in extract_code_block(text) + assert "short" not in extract_code_block(text) + + def test_unfenced_code_falls_back_to_raw_text(self): + text = "import matplotlib.pyplot as plt\nplt.plot([1, 2])" + assert extract_code_block(text) == text + + def test_prose_without_code_raises(self): + with pytest.raises(ValueError, match="No fenced code block"): + extract_code_block("I am sorry, I cannot help with that.") diff --git a/tests/unit/automation/benchmark/test_runner.py b/tests/unit/automation/benchmark/test_runner.py new file mode 100644 index 0000000000..6c95f90da1 --- /dev/null +++ b/tests/unit/automation/benchmark/test_runner.py @@ -0,0 +1,131 @@ +"""Tests for automation.benchmark.runner.""" + +import os +import textwrap + +from PIL import Image + +from automation.benchmark.runner import check_canvas, render_env, run_python_implementation + + +def _write_png(path, width, height): + Image.new("RGB", (width, height), color=(0, 158, 115)).save(path) + + +class TestCheckCanvas: + def test_exact_landscape_target(self, tmp_path): + png = tmp_path / "plot.png" + _write_png(png, 3200, 1800) + assert check_canvas(png) is True + + def test_exact_square_target(self, tmp_path): + png = tmp_path / "plot.png" + _write_png(png, 2400, 2400) + assert check_canvas(png) is True + + def test_within_tolerance(self, tmp_path): + png = tmp_path / "plot.png" + _write_png(png, 3200 - 16, 1800 + 16) + assert check_canvas(png) is True + + def test_drifted_canvas_fails(self, tmp_path): + png = tmp_path / "plot.png" + _write_png(png, 3000, 1800) + assert check_canvas(png) is False + + +class TestRenderEnv: + def test_secrets_are_not_inherited(self, monkeypatch, tmp_path): + monkeypatch.setenv("FAKE_GCP_CREDENTIAL", "super-secret") + env = render_env("light", home=tmp_path) + assert "FAKE_GCP_CREDENTIAL" not in env + + def test_theme_and_backend_are_set(self, tmp_path): + env = render_env("dark", home=tmp_path) + assert env["ANYPLOT_THEME"] == "dark" + assert env["MPLBACKEND"] == "Agg" + + def test_path_passes_through(self, tmp_path): + assert "PATH" in render_env("light", home=tmp_path) + + def test_home_is_remapped_to_scratch_dir(self, tmp_path): + # The caller's real HOME (~/.config with gcloud credentials, …) must + # not be visible to LLM-generated code. + env = render_env("light", home=tmp_path) + assert env["HOME"] == str(tmp_path) + assert env["HOME"] != os.environ.get("HOME") + + +class TestRunPythonImplementation: + def _write_script(self, tmp_path, body): + script = tmp_path / "impl.py" + script.write_text(textwrap.dedent(body), encoding="utf-8") + return script + + def test_successful_render(self, tmp_path): + script = self._write_script( + tmp_path, + """ + import os + from PIL import Image + theme = os.getenv("ANYPLOT_THEME", "light") + Image.new("RGB", (100, 100)).save(f"plot-{theme}.png") + """, + ) + result = run_python_implementation(script, workdir=tmp_path, timeout=60) + assert result.success is True + assert result.canvas_ok is False # 100x100 is far off both canonical targets + assert set(result.images) == {"light", "dark"} + + def test_failing_script_reports_stderr(self, tmp_path): + script = self._write_script(tmp_path, "raise SystemExit('boom')") + result = run_python_implementation(script, workdir=tmp_path, timeout=60) + assert result.success is False + assert "exited 1" in result.error + assert "boom" in result.error + + def test_missing_png_is_an_error(self, tmp_path): + script = self._write_script(tmp_path, "print('renders nothing')") + result = run_python_implementation(script, workdir=tmp_path, timeout=60) + assert result.success is False + assert "wrote no plot-light.png" in result.error + + def test_canvas_gate_requires_both_themes_on_target(self, tmp_path): + # Light hits 3200x1800 but dark drifts — the gate must fail. + script = self._write_script( + tmp_path, + """ + import os + from PIL import Image + theme = os.getenv("ANYPLOT_THEME", "light") + size = (3200, 1800) if theme == "light" else (3000, 1800) + Image.new("RGB", size).save(f"plot-{theme}.png") + """, + ) + result = run_python_implementation(script, workdir=tmp_path, timeout=60) + assert result.success is True + assert result.canvas_ok is False + + def test_stale_pngs_from_prior_attempt_do_not_count(self, tmp_path): + # A previous attempt left renders behind; this attempt writes nothing. + _write_png(tmp_path / "plot-light.png", 10, 10) + _write_png(tmp_path / "plot-dark.png", 10, 10) + script = self._write_script(tmp_path, "print('renders nothing')") + result = run_python_implementation(script, workdir=tmp_path, timeout=60) + assert result.success is False + assert "wrote no plot-light.png" in result.error + + def test_generated_code_does_not_see_caller_secrets(self, tmp_path, monkeypatch): + monkeypatch.setenv("FAKE_GCP_CREDENTIAL", "super-secret") + script = self._write_script( + tmp_path, + """ + import os + from PIL import Image + assert "FAKE_GCP_CREDENTIAL" not in os.environ, "secret leaked into render env" + theme = os.getenv("ANYPLOT_THEME", "light") + Image.new("RGB", (10, 10)).save(f"plot-{theme}.png") + """, + ) + result = run_python_implementation(script, workdir=tmp_path, timeout=60) + assert result.success is True diff --git a/tests/unit/automation/benchmark/test_vertex_client.py b/tests/unit/automation/benchmark/test_vertex_client.py new file mode 100644 index 0000000000..e5de619d90 --- /dev/null +++ b/tests/unit/automation/benchmark/test_vertex_client.py @@ -0,0 +1,216 @@ +"""Tests for automation.benchmark.vertex_client provider routing.""" + +from types import SimpleNamespace + +import pytest + +from automation.benchmark.vertex_client import ( + PROVIDER_ANTHROPIC, + PROVIDER_OPENAPI, + GenerationResult, + VertexClient, + _extract_openapi_text, + resolve_provider, +) + + +class TestResolveProvider: + """Model ids route to the Vertex surface that actually serves them.""" + + def test_bare_claude_id(self): + assert resolve_provider("claude-sonnet-4-5@20250929") == (PROVIDER_ANTHROPIC, "claude-sonnet-4-5@20250929") + + def test_anthropic_prefixed_id_is_stripped(self): + assert resolve_provider("anthropic/claude-opus-4-1@20250805") == ( + PROVIDER_ANTHROPIC, + "claude-opus-4-1@20250805", + ) + + def test_bare_gemini_id_gets_google_prefix(self): + assert resolve_provider("gemini-2.5-pro") == (PROVIDER_OPENAPI, "google/gemini-2.5-pro") + + def test_google_prefixed_gemini_id_passes_through(self): + assert resolve_provider("google/gemini-2.5-flash") == (PROVIDER_OPENAPI, "google/gemini-2.5-flash") + + def test_model_garden_partner_ids_pass_through(self): + for model in ("meta/llama-3.3-70b-instruct-maas", "mistralai/mistral-large-2411", "qwen/qwen3-coder"): + assert resolve_provider(model) == (PROVIDER_OPENAPI, model) + + def test_whitespace_is_trimmed(self): + assert resolve_provider(" gemini-2.5-pro ") == (PROVIDER_OPENAPI, "google/gemini-2.5-pro") + + def test_unroutable_id_raises(self): + with pytest.raises(ValueError, match="Cannot route model id"): + resolve_provider("gpt-5") + + def test_empty_id_raises(self): + with pytest.raises(ValueError, match="must not be empty"): + resolve_provider(" ") + + +class TestOpenapiUrl: + """The chat-completions URL follows Vertex's regional host scheme.""" + + def test_regional_location(self): + client = VertexClient(project="anyplot", location="us-central1") + assert client._openapi_url() == ( + "https://us-central1-aiplatform.googleapis.com/v1/projects/anyplot" + "/locations/us-central1/endpoints/openapi/chat/completions" + ) + + def test_global_location_drops_region_prefix(self): + client = VertexClient(project="anyplot", location="global") + assert client._openapi_url() == ( + "https://aiplatform.googleapis.com/v1/projects/anyplot/locations/global/endpoints/openapi/chat/completions" + ) + + +class TestGenerateRouting: + """generate() dispatches by provider and stamps wall-clock latency.""" + + def _stub_result(self, model): + return GenerationResult(text="code", model=model, provider="stub", latency_seconds=0.0) + + def test_claude_id_dispatches_to_anthropic(self, monkeypatch): + client = VertexClient(project="anyplot") + called = {} + + def fake_anthropic(model, system, prompt, max_tokens, temperature): + called["args"] = (model, system, prompt, max_tokens, temperature) + return self._stub_result(model) + + monkeypatch.setattr(client, "_generate_anthropic", fake_anthropic) + result = client.generate("claude-sonnet-4-5@20250929", "sys", "user", max_tokens=42, temperature=0.5) + assert called["args"] == ("claude-sonnet-4-5@20250929", "sys", "user", 42, 0.5) + assert result.latency_seconds >= 0.0 + + def test_gemini_id_dispatches_to_openapi_with_publisher_prefix(self, monkeypatch): + client = VertexClient(project="anyplot") + called = {} + + def fake_openapi(model, system, prompt, max_tokens, temperature): + called["model"] = model + return self._stub_result(model) + + monkeypatch.setattr(client, "_generate_openapi", fake_openapi) + client.generate("gemini-2.5-pro", "sys", "user") + assert called["model"] == "google/gemini-2.5-pro" + + +class TestGenerateAnthropic: + """Claude on Vertex goes through the Anthropic SDK transport.""" + + class _FakeAnthropicVertex: + last_init = None + + def __init__(self, **kwargs): + type(self).last_init = kwargs + self.messages = SimpleNamespace(create=self._create) + + def _create(self, **kwargs): + self.last_create = kwargs + return SimpleNamespace( + content=[ + SimpleNamespace(type="thinking", text="ignored"), + SimpleNamespace(type="text", text="the code"), + ], + usage=SimpleNamespace(input_tokens=100, output_tokens=50), + ) + + def test_response_text_and_usage(self, monkeypatch): + import anthropic + + monkeypatch.setattr(anthropic, "AnthropicVertex", self._FakeAnthropicVertex) + client = VertexClient(project="anyplot", anthropic_location="europe-west1") + result = client._generate_anthropic("claude-sonnet-4-5@20250929", "sys", "user", 1000, 0.2) + + assert result.text == "the code" # non-text blocks are skipped + assert result.provider == PROVIDER_ANTHROPIC + assert result.input_tokens == 100 + assert result.output_tokens == 50 + assert self._FakeAnthropicVertex.last_init["project_id"] == "anyplot" + assert self._FakeAnthropicVertex.last_init["region"] == "europe-west1" + + +class TestGenerateOpenapi: + """Gemini / Model Garden go through the OpenAI-compatible endpoint.""" + + def _fake_response(self, status_code=200, body=None, text="err"): + return SimpleNamespace(status_code=status_code, json=lambda: body, text=text) + + def test_success_parses_text_and_usage(self, monkeypatch): + from automation.benchmark import vertex_client + + client = VertexClient(project="anyplot") + monkeypatch.setattr(client, "_access_token", lambda: "tok") + captured = {} + + def fake_post(url, json=None, headers=None, timeout=None): + captured.update(url=url, json=json, headers=headers) + return self._fake_response( + body={ + "choices": [{"message": {"content": "gemini code"}}], + "usage": {"prompt_tokens": 11, "completion_tokens": 7}, + } + ) + + monkeypatch.setattr(vertex_client.httpx, "post", fake_post) + result = client._generate_openapi("google/gemini-2.5-pro", "sys", "user", 1000, 0.2) + + assert result.text == "gemini code" + assert result.provider == PROVIDER_OPENAPI + assert result.input_tokens == 11 + assert result.output_tokens == 7 + assert captured["headers"]["Authorization"] == "Bearer tok" + assert captured["json"]["model"] == "google/gemini-2.5-pro" + assert captured["json"]["messages"][0]["role"] == "system" + + def test_non_200_raises_with_model_and_status(self, monkeypatch): + from automation.benchmark import vertex_client + + client = VertexClient(project="anyplot") + monkeypatch.setattr(client, "_access_token", lambda: "tok") + monkeypatch.setattr( + vertex_client.httpx, "post", lambda *a, **k: self._fake_response(status_code=404, text="model not found") + ) + with pytest.raises(RuntimeError, match="HTTP 404.*model not found"): + client._generate_openapi("meta/nope", "sys", "user", 1000, 0.2) + + +class TestAccessToken: + def test_returns_token_from_adc(self, monkeypatch): + import google.auth + + creds = SimpleNamespace(valid=True, token="adc-token") + monkeypatch.setattr(google.auth, "default", lambda scopes: (creds, "anyplot")) + client = VertexClient(project="anyplot") + assert client._access_token() == "adc-token" + + def test_empty_token_raises(self, monkeypatch): + import google.auth + + creds = SimpleNamespace(valid=True, token=None) + monkeypatch.setattr(google.auth, "default", lambda scopes: (creds, "anyplot")) + client = VertexClient(project="anyplot") + with pytest.raises(RuntimeError, match="access token"): + client._access_token() + + +class TestExtractOpenapiText: + """Both OpenAI content shapes (string and parts list) are accepted.""" + + def test_string_content(self): + body = {"choices": [{"message": {"content": "hello"}}]} + assert _extract_openapi_text(body) == "hello" + + def test_parts_list_content(self): + body = {"choices": [{"message": {"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}}]} + assert _extract_openapi_text(body) == "ab" + + def test_no_choices_raises(self): + with pytest.raises(RuntimeError, match="no choices"): + _extract_openapi_text({"choices": []}) + + def test_unexpected_content_shape_raises(self): + with pytest.raises(RuntimeError, match="Unexpected content shape"): + _extract_openapi_text({"choices": [{"message": {"content": 42}}]})