run_experiments.py

#!/usr/bin/env python3
"""
Run repeated Gemini Agent Laboratory experiments from YAML files.

The Gemini preflight and all Gemini model calls use Google's native
google-genai SDK. This runner does not import OpenAI.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import yaml
from google import genai


REPO_DIR = Path(__file__).resolve().parent
RESULTS_DIR = REPO_DIR / "results"
MODEL_REGISTRY_FILE = REPO_DIR / "configs" / "models_pricing.json"

# Start with one run per model. Change to (1, 2, 3) after a successful test.
RUN_INDICES = (1,)
STOP_MODEL_AFTER_FAILURE = True

EXPERIMENTS = (
    {
        "yaml": REPO_DIR
        / "experiment_configs"
        / "sgc_gemini_3_5_flash_lite.yaml",
        "name": "gemini_3_5_flash",
    },
    {
        "yaml": REPO_DIR
        / "experiment_configs"
        / "sgc_gemini_3_1_flash_lite.yaml",
        "name": "gemini_3_1_flash",
    },
)

CLI_OPTIONS = {
    "copilot-mode": "--copilot-mode",
    "load-existing": "--load-existing",
    "load-existing-path": "--load-existing-path",
    "research-topic": "--research-topic",
    "compile-latex": "--compile-latex",
    "llm-backend": "--llm-backend",
    "language": "--language",
    "num-papers-lit-review": "--num-papers-lit-review",
    "mlesolver-max-steps": "--mlesolver-max-steps",
    "papersolver-max-steps": "--papersolver-max-steps",
    "ollama-max-tokens": "--ollama-max-tokens",
    "task-note-llm-config-file": "--task-note-llm-config-file",
}


def cli_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    return str(value)


def load_yaml(path: Path) -> dict[str, Any]:
    if not path.is_file():
        raise FileNotFoundError(f"YAML file not found: {path}")

    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"YAML root must be a mapping: {path}")

    for key in ("research-topic", "llm-backend"):
        if not data.get(key):
            raise ValueError(f"Missing required YAML key '{key}': {path}")

    return data


def list_available_models(api_key: str) -> set[str]:
    client = genai.Client(api_key=api_key)
    names: set[str] = set()

    for model in client.models.list():
        name = getattr(model, "name", None)
        if name:
            names.add(str(name).removeprefix("models/"))

    return names


def register_model(model_name: str) -> None:
    """
    Add a Gemini model to the fork's local registry if it is absent.

    Prices are null to avoid displaying invented cost estimates.
    """
    if not MODEL_REGISTRY_FILE.is_file():
        raise FileNotFoundError(
            f"Missing model registry: {MODEL_REGISTRY_FILE}"
        )

    data = json.loads(MODEL_REGISTRY_FILE.read_text(encoding="utf-8"))
    models = data.setdefault("models", {})

    desired = {
        "provider": "google",
        "api_model_name": model_name,
        "aliases": [],
        "cost_per_million_input": None,
        "cost_per_million_output": None,
    }

    if models.get(model_name) == desired:
        return

    backup = MODEL_REGISTRY_FILE.with_suffix(".json.backup")
    if not backup.exists():
        shutil.copy2(MODEL_REGISTRY_FILE, backup)

    models[model_name] = desired
    data["last_updated"] = datetime.now(timezone.utc).isoformat()

    MODEL_REGISTRY_FILE.write_text(
        json.dumps(data, indent=2) + "\n",
        encoding="utf-8",
    )
    print(f"Registered model locally: {model_name}")


def build_command(
    yaml_data: dict[str, Any],
    gemini_api_key: str,
) -> list[str]:
    command = [
        sys.executable,
        "-u",
        "ai_lab_repo.py",
        # The fork requires this argument even when Google is the provider.
        "--api-key",
        str(yaml_data.get("api-key") or "NOT-USED"),
        "--google-api-key",
        gemini_api_key,
    ]

    for yaml_key, cli_flag in CLI_OPTIONS.items():
        value = yaml_data.get(yaml_key)
        if value is not None:
            command.extend([cli_flag, cli_value(value)])

    return command


def output_targets() -> list[Path]:
    candidates: list[Path] = [
        REPO_DIR / "research_dir",
        REPO_DIR / "MATH_research_dir",
        REPO_DIR / "state_saves",
    ]
    candidates.extend(REPO_DIR.glob("*_research_dir"))
    candidates.extend(REPO_DIR.glob("agent_times_*.txt"))

    unique: list[Path] = []
    seen: set[Path] = set()

    for path in candidates:
        resolved = path.resolve()
        if path.exists() and resolved not in seen:
            seen.add(resolved)
            unique.append(path)

    return unique


def remove_stale_outputs() -> None:
    for target in output_targets():
        if target.is_dir():
            shutil.rmtree(target)
        else:
            target.unlink()


def archive_outputs(result_dir: Path) -> None:
    targets = output_targets()

    if not targets:
        print("Warning: no Agent Laboratory output targets were found.")
        return

    for source in targets:
        destination = result_dir / source.name

        if destination.exists():
            if destination.is_dir():
                shutil.rmtree(destination)
            else:
                destination.unlink()

        shutil.move(str(source), str(destination))
        print(f"Archived: {source.name} -> {destination}")


def run_logged(
    command: list[str],
    environment: dict[str, str],
    log_path: Path,
) -> int:
    with log_path.open("w", encoding="utf-8") as log_file:
        process = subprocess.Popen(
            command,
            cwd=REPO_DIR,
            env=environment,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )

        # ai_lab_repo.py calls input() at exit which blocks when stdout
        # is piped.  Pre-load stdin with newlines so those calls return
        # immediately instead of deadlocking.
        process.stdin.write("\n\n\n")
        process.stdin.flush()
        process.stdin.close()

        if process.stdout is None:
            raise RuntimeError("Could not capture process output.")

        for line in process.stdout:
            print(line, end="", flush=True)
            log_file.write(line)
            log_file.flush()

        return process.wait()


def run_experiment(
    yaml_path: Path,
    result_name: str,
    run_index: int,
    api_key: str,
    available_models: set[str],
) -> bool:
    yaml_data = load_yaml(yaml_path)
    model_name = str(yaml_data["llm-backend"])

    print("\n" + "=" * 72)
    print(f"STARTING {result_name} - RUN {run_index}")
    print(f"MODEL: {model_name}")
    print("=" * 72 + "\n")

    # The YAML may use a short alias (e.g. "gemini-3.1-flash-lite") while
    # the API lists the full name (e.g. "gemini-3.1-flash-lite-preview").
    # Try the raw name first, then resolve through the model registry.
    api_name = model_name
    if model_name not in available_models:
        try:
            from model_registry import ModelRegistry
            reg = ModelRegistry(auto_refresh=False)
            api_name = reg.get_api_model_name(model_name)
        except Exception:
            pass

    if api_name not in available_models and model_name not in available_models:
        print(f"ERROR: Gemini did not list model '{model_name}' for this key.")
        if api_name != model_name:
            print(f"  (also tried registry api_model_name '{api_name}')")
        print("No research calls were started.")
        return False

    register_model(model_name)
    remove_stale_outputs()

    result_dir = RESULTS_DIR / f"{result_name}_run_{run_index}"
    result_dir.mkdir(parents=True, exist_ok=True)
    log_path = result_dir / f"{result_name}_run_{run_index}.log"

    environment = os.environ.copy()
    environment["GEMINI_API_KEY"] = api_key
    environment["GOOGLE_API_KEY"] = api_key
    environment["PYTHONUNBUFFERED"] = "1"

    command = build_command(yaml_data, api_key)

    print(f"Configuration: {yaml_path}")
    print(f"Log: {log_path}\n")

    started = time.monotonic()
    return_code = run_logged(command, environment, log_path)
    duration = time.monotonic() - started

    print(
        f"\nFinished run {run_index} for {result_name} "
        f"in {duration:.2f} seconds with exit code {return_code}."
    )

    archive_outputs(result_dir)

    (result_dir / "run_summary.txt").write_text(
        "\n".join(
            (
                f"name={result_name}",
                f"model={model_name}",
                f"yaml={yaml_path}",
                f"run={run_index}",
                f"duration_seconds={duration:.2f}",
                f"exit_code={return_code}",
            )
        )
        + "\n",
        encoding="utf-8",
    )

    if return_code != 0:
        print(f"Run failed. Review: {log_path}")
        return False

    return True


def main() -> int:
    os.chdir(REPO_DIR)
    RESULTS_DIR.mkdir(exist_ok=True)

    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        print(
            "ERROR: GEMINI_API_KEY is not set.\n"
            "In Zsh run:\n"
            '  read -rs "GEMINI_API_KEY?Paste Gemini API key: "\n'
            "  echo\n"
            "  export GEMINI_API_KEY",
            file=sys.stderr,
        )
        return 2

    try:
        available_models = list_available_models(api_key)
    except Exception as error:
        print(f"ERROR: Gemini API preflight failed: {error}", file=sys.stderr)
        return 2

    print(
        f"Native Gemini API preflight succeeded; "
        f"{len(available_models)} models listed."
    )

    failures: list[str] = []

    for experiment in EXPERIMENTS:
        yaml_path = Path(experiment["yaml"])
        result_name = str(experiment["name"])

        for run_index in RUN_INDICES:
            try:
                succeeded = run_experiment(
                    yaml_path,
                    result_name,
                    run_index,
                    api_key,
                    available_models,
                )
            except Exception as error:
                print(
                    f"\nERROR during {result_name} run {run_index}: {error}",
                    file=sys.stderr,
                )
                succeeded = False

            if not succeeded:
                failures.append(f"{result_name} run {run_index}")
                if STOP_MODEL_AFTER_FAILURE:
                    print(
                        f"Skipping remaining runs for {result_name} "
                        "because this run failed."
                    )
                    break

    print("\n" + "=" * 72)
    print("ALL REQUESTED RUNS FINISHED")
    print("=" * 72)

    if failures:
        print("Failed runs:")
        for failure in failures:
            print(f"  - {failure}")
        return 1

    print(f"Results saved under: {RESULTS_DIR}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())