Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
7.3 kB · 177 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178from __future__ import annotations
import argparseimport jsonimport osimport subprocessfrom datetime import UTC, datetimefrom pathlib import Path
from prime_sandboxes import APIClient, SandboxClient
from controller import Controller, cleanupfrom github import WORKFLOW, GitHubfrom report import renderfrom schema import Config, Report, Side, load_report, report_complete, write_json
def workflow_source(event: dict) -> dict: run = event["workflow_run"] if ( run["event"] not in ("pull_request_target", "workflow_dispatch") or run["path"].split("@")[0] != f".github/workflows/{WORKFLOW}" or run["repository"]["full_name"] != os.environ["GITHUB_REPOSITORY"] or (run["event"] == "workflow_dispatch" and run["head_branch"] != "main") ): raise ValueError("Unexpected source workflow") return run
def validate_completion(report: Report, request: Report, run: dict) -> None: identity = { "repository", "head_repository", "pr", "run_id", "attempt", "harness_sha", "base_sha", "head_sha", "started_at", "config", } if report.model_dump(include=identity) != request.model_dump(include=identity): raise ValueError("Result identity does not match the trusted benchmark request") if (request.run_id, request.attempt) != (run["id"], run["run_attempt"]) or run[ "display_title" ] != f"Prime Agent benchmarks · PR #{request.pr}": raise ValueError("Benchmark request does not match the completed workflow")
def completed_report(result_path: Path, request_path: Path, run: dict) -> Report: request = load_report(request_path) validate_completion(request, request, run) try: report = load_report(result_path) if result_path.exists() else request validate_completion(report, request, run) except (ValueError, OSError): report = request report.status = "failed" report.errors.append("Benchmark report failed validation; see the workflow artifacts") if report.status == "completed" and not report_complete(report): report.status = "partial" report.errors.append("Benchmark claimed completion with failed or missing measurements") if report.status == "running": report.status = "canceled" if run["conclusion"] == "cancelled" else "failed" report.errors.append("Benchmark did not produce a final report; see the workflow logs") if report.status != "pending-trust" and report.finished_at is None: report.finished_at = datetime.now(UTC) return report
def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("command", choices=("resolve", "run", "publish", "finalize", "cleanup", "local")) parser.add_argument("--results", type=Path, default=Path("results")) parser.add_argument("--request", type=Path, default=Path("request/report.json")) parser.add_argument("--base") parser.add_argument("--head") parser.add_argument("--config", type=Path) parser.add_argument("--pending-trust", action="store_true") args = parser.parse_args() args.results.mkdir(parents=True, exist_ok=True) repository = os.environ.get("GITHUB_REPOSITORY", "PrimeIntellect-ai/prime-agent") config = Config.load(args.config) if args.config else Config.load() if args.command == "local": if not args.base or not args.head: parser.error("local requires exact --base and --head commit SHAs") report = Report( repository=repository, head_repository=repository, pr=0, run_id=int(datetime.now(UTC).timestamp()), attempt=1, harness_sha=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), base_sha=args.base, head_sha=args.head, started_at=datetime.now(UTC), config=config, main=Side(sha=args.base), pr_head=Side(sha=args.head), ) controller = Controller(report, args.results, live_github=False) controller.run() require_success(controller.report) return event = json.loads(Path(os.environ["GITHUB_EVENT_PATH"]).read_text()) if args.command == "cleanup": run = workflow_source(event) client = SandboxClient(APIClient(api_key=os.environ["PRIME_SANDBOX_API_KEY"])) deleted = cleanup(client, repository, run["id"], run["run_attempt"]) print(f"Cleaned up {len(deleted)} remaining benchmark sandboxes") return github = GitHub(repository) if args.command == "resolve": pr = int(event.get("pull_request", {}).get("number") or event["inputs"]["pr"]) if pr <= 0: raise ValueError("Invalid PR number") report, author = github.resolve( pr, os.environ["GITHUB_SHA"], int(os.environ["GITHUB_RUN_ID"]), int(os.environ["GITHUB_RUN_ATTEMPT"]), config, ) write_json(args.results / "report.json", report) manual = os.environ.get("GITHUB_EVENT_NAME") == "workflow_dispatch" or report.attempt > 1 needed = manual or not github.duplicate(report) with Path(os.environ["GITHUB_OUTPUT"]).open("a") as output: output.write( f"pr={pr}\nauthor={author}\nharness={report.harness_sha}\nneeded={str(needed).lower()}\n" ) return report_path = args.results / "report.json" report = ( completed_report(report_path, args.request, workflow_source(event)) if args.command == "finalize" else load_report(report_path if report_path.exists() else args.request) ) if report.repository != repository: raise ValueError("Result repository does not match the workflow") if args.command == "finalize": github.publish(report) else: if (report.run_id, report.attempt) != ( int(os.environ["GITHUB_RUN_ID"]), int(os.environ["GITHUB_RUN_ATTEMPT"]), ): raise ValueError("Result identity does not match the active workflow") if args.command == "run": if not os.environ.get("PRIME_SANDBOX_API_KEY"): report.status = "failed" report.errors.append("Repository secret PRIME_SANDBOX_API_KEY is not configured") if report.status == "running": controller = Controller(report, args.results) controller.run() report = controller.report else: write_json(report_path, report) else: if args.pending_trust: report.status = "pending-trust" write_json(report_path, report) github.publish(report) summary = os.environ.get("GITHUB_STEP_SUMMARY") if summary: Path(summary).write_text(render(report)) if args.command == "run": require_success(report)
def require_success(report: Report) -> None: if report.status != "completed" or not report_complete(report): raise SystemExit(f"Benchmark {report.status}: failed or incomplete; see saved report and transcripts")
if __name__ == "__main__": main()