Initial public upload from Projekte/AgentArena

3b0d6b3 verified about 1 month ago

7.28 kB

import json
import logging
import os
import re
import shlex
import shutil
import subprocess
import tempfile
import time

from config import CODEX_CLI_PATH, get_pricing
from models import RunStatus

from adapters.base import AgentAdapter, RunContext, RunResult

logger = logging.getLogger(__name__)

# Codex prints token usage like: "tokens used\n1.019"
_TOKENS_USED_RE = re.compile(r"tokens\s+used\s*\n\s*([0-9.]+)", re.IGNORECASE)
# Broader alternatives
_TOKEN_ALT_RE = re.compile(r"(?:total\s+)?tokens[:\s]+(\d[\d,]*)", re.I)
_COST_ALT_RE = re.compile(r"(?:cost|price|total)[:\s]*\$([0-9.]+)", re.I)

# Default model assumed when Codex doesn't report one
_DEFAULT_MODEL = "gpt-4o"


class CodexAdapter(AgentAdapter):
    def name(self) -> str:
        return "Codex"

    def is_available(self) -> bool:
        return shutil.which(CODEX_CLI_PATH) is not None

    async def run(self, ctx: RunContext, sandbox) -> RunResult:
        """Run Codex CLI on the host against a local clone of the task repo.
        Codex uses ChatGPT Plus OAuth (no API key needed).
        Runs non-interactively via `codex exec -s full-auto`.
        """
        start = time.time()
        tmpdir = tempfile.mkdtemp(prefix="arena-codex-")

        try:
            # Clone repo to host temp dir.
            repo_url = ctx.repo_url or ""
            repo_ref = ctx.repo_ref or "main"

            if repo_url:
                clone_result = subprocess.run(
                    ["git", "clone", "--depth=1", "-b", repo_ref, repo_url, tmpdir],
                    capture_output=True, text=True, timeout=120,
                )
                if clone_result.returncode != 0:
                    return RunResult(
                        status=RunStatus.error,
                        duration_s=round(time.time() - start, 2),
                        raw_output=clone_result.stderr[:5000],
                    )
                work_dir = tmpdir
            else:
                work_dir = ctx.work_dir

            # Git checkpoint for diff.
            subprocess.run(
                ["git", "tag", "arena-checkpoint"],
                cwd=work_dir, capture_output=True, timeout=10,
            )

            # Use model from agent config, fall back to default.
            model = self.config.get("model", _DEFAULT_MODEL)

            prompt = (
                f"Fix the following issue in the codebase:\n\n"
                f"{ctx.task_prompt}\n\n"
                f"After fixing, verify your changes pass: {ctx.success_cmd}"
            )

            cmd = [
                CODEX_CLI_PATH,
                "exec",
                "-s", "danger-full-access",
                "--ephemeral",
                "--model", model,
                prompt,
            ]

            try:
                proc = subprocess.run(
                    cmd,
                    cwd=work_dir,
                    capture_output=True,
                    text=True,
                    timeout=ctx.time_limit_s,
                )
                returncode = proc.returncode
                stdout = proc.stdout
                stderr = proc.stderr
            except subprocess.TimeoutExpired:
                returncode = 124
                stdout = ""
                stderr = f"Command timed out after {ctx.time_limit_s}s"

            duration = time.time() - start

            # Collect diff, exclude __pycache__ and .pyc.
            def _get_diff(rev):
                r = subprocess.run(
                    ["git", "diff", rev, "--", ".", ":!**/__pycache__/", ":!**/*.pyc"],
                    cwd=work_dir, capture_output=True, text=True, timeout=30,
                )
                return r.stdout if r.returncode == 0 else ""

            diff_text = _get_diff("arena-checkpoint")
            if not diff_text.strip():
                diff_text = _get_diff("arena-checkpoint..HEAD")

            # Parse token usage from output.
            tokens_input = 0
            tokens_output = 0
            cost_usd = 0.0
            combined = stdout + "\n" + stderr

            # Log raw output for diagnostics
            logger.info("CODEX RAW OUTPUT (first 2000 chars):\n%s", combined[:2000])

            # Codex prints "tokens used\n1.019" (in thousands)
            m = _TOKENS_USED_RE.search(combined)
            if m:
                total_k = float(m.group(1))
                total_tokens = int(total_k * 1000)
                # Rough split: 70% input, 30% output
                tokens_input = int(total_tokens * 0.7)
                tokens_output = total_tokens - tokens_input

            # Fallback: try alternative token/cost patterns
            if tokens_input == 0 and tokens_output == 0:
                # Filter out metadata lines (e.g., "Repo-map: using 1024 tokens")
                content_lines = [
                    line for line in combined.split("\n")
                    if not re.search(r"repo.?map|using \d+ tokens.*refresh", line, re.I)
                ]
                filtered = "\n".join(content_lines)
                alt_tokens = _TOKEN_ALT_RE.findall(filtered)
                if alt_tokens:
                    total = sum(int(t.replace(",", "")) for t in alt_tokens)
                    if total >= 100:  # Sanity check: real LLM usage > 100 tokens
                        tokens_input = int(total * 0.7)
                        tokens_output = total - tokens_input
            if cost_usd == 0.0:
                alt_costs = _COST_ALT_RE.findall(combined)
                if alt_costs:
                    cost_usd = sum(float(c) for c in alt_costs)

            # Estimate cost from tokens using dynamic pricing lookup.
            if cost_usd == 0.0 and (tokens_input or tokens_output):
                rate_in, _, _, rate_out = get_pricing(model)
                cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000

            # Last resort: estimate from raw output length (better than $0)
            if cost_usd == 0.0 and (tokens_input + tokens_output) < 50 and len(combined) > 100:
                estimated_output_tokens = len(combined) // 4
                rate_in, _, _, rate_out = get_pricing(model)
                tokens_input = estimated_output_tokens * 3
                tokens_output = estimated_output_tokens
                cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000
                logger.warning("Cost estimated from output length (%d chars) — no token data found", len(combined))

            trajectory = [{"tool_name": "Bash", "input_preview": f"codex exec --model {model}"}]

            if duration >= ctx.time_limit_s:
                status = RunStatus.timeout
            elif returncode == 0:
                status = RunStatus.success
            else:
                status = RunStatus.failure

            return RunResult(
                status=status,
                duration_s=round(duration, 2),
                tokens_input=tokens_input,
                tokens_output=tokens_output,
                cost_usd=round(cost_usd, 6),
                diff_text=diff_text,
                trajectory=trajectory,
                raw_output=(stdout + stderr)[:50_000],
            )
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)

	import json
	import logging
	import os
	import re
	import shlex
	import shutil
	import subprocess
	import tempfile
	import time

	from config import CODEX_CLI_PATH, get_pricing
	from models import RunStatus

	from adapters.base import AgentAdapter, RunContext, RunResult

	logger = logging.getLogger(__name__)

	# Codex prints token usage like: "tokens used\n1.019"
	_TOKENS_USED_RE = re.compile(r"tokens\s+used\s\n\s([0-9.]+)", re.IGNORECASE)
	# Broader alternatives
	_TOKEN_ALT_RE = re.compile(r"(?:total\s+)?tokens[:\s]+(\d[\d,]*)", re.I)
	_COST_ALT_RE = re.compile(r"(?:cost\|price\|total)[:\s]*\$([0-9.]+)", re.I)

	# Default model assumed when Codex doesn't report one
	_DEFAULT_MODEL = "gpt-4o"


	class CodexAdapter(AgentAdapter):
	def name(self) -> str:
	return "Codex"

	def is_available(self) -> bool:
	return shutil.which(CODEX_CLI_PATH) is not None

	async def run(self, ctx: RunContext, sandbox) -> RunResult:
	"""Run Codex CLI on the host against a local clone of the task repo.

	Codex uses ChatGPT Plus OAuth (no API key needed).
	Runs non-interactively via `codex exec -s full-auto`.
	"""
	start = time.time()
	tmpdir = tempfile.mkdtemp(prefix="arena-codex-")

	try:
	# Clone repo to host temp dir.
	repo_url = ctx.repo_url or ""
	repo_ref = ctx.repo_ref or "main"

	if repo_url:
	clone_result = subprocess.run(
	["git", "clone", "--depth=1", "-b", repo_ref, repo_url, tmpdir],
	capture_output=True, text=True, timeout=120,
	)
	if clone_result.returncode != 0:
	return RunResult(
	status=RunStatus.error,
	duration_s=round(time.time() - start, 2),
	raw_output=clone_result.stderr[:5000],
	)
	work_dir = tmpdir
	else:
	work_dir = ctx.work_dir

	# Git checkpoint for diff.
	subprocess.run(
	["git", "tag", "arena-checkpoint"],
	cwd=work_dir, capture_output=True, timeout=10,
	)

	# Use model from agent config, fall back to default.
	model = self.config.get("model", _DEFAULT_MODEL)

	prompt = (
	f"Fix the following issue in the codebase:\n\n"
	f"{ctx.task_prompt}\n\n"
	f"After fixing, verify your changes pass: {ctx.success_cmd}"
	)

	cmd = [
	CODEX_CLI_PATH,
	"exec",
	"-s", "danger-full-access",
	"--ephemeral",
	"--model", model,
	prompt,
	]

	try:
	proc = subprocess.run(
	cmd,
	cwd=work_dir,
	capture_output=True,
	text=True,
	timeout=ctx.time_limit_s,
	)
	returncode = proc.returncode
	stdout = proc.stdout
	stderr = proc.stderr
	except subprocess.TimeoutExpired:
	returncode = 124
	stdout = ""
	stderr = f"Command timed out after {ctx.time_limit_s}s"

	duration = time.time() - start

	# Collect diff, exclude __pycache__ and .pyc.
	def _get_diff(rev):
	r = subprocess.run(
	["git", "diff", rev, "--", ".", ":!/__pycache__/", ":!/*.pyc"],
	cwd=work_dir, capture_output=True, text=True, timeout=30,
	)
	return r.stdout if r.returncode == 0 else ""

	diff_text = _get_diff("arena-checkpoint")
	if not diff_text.strip():
	diff_text = _get_diff("arena-checkpoint..HEAD")

	# Parse token usage from output.
	tokens_input = 0
	tokens_output = 0
	cost_usd = 0.0
	combined = stdout + "\n" + stderr

	# Log raw output for diagnostics
	logger.info("CODEX RAW OUTPUT (first 2000 chars):\n%s", combined[:2000])

	# Codex prints "tokens used\n1.019" (in thousands)
	m = _TOKENS_USED_RE.search(combined)
	if m:
	total_k = float(m.group(1))
	total_tokens = int(total_k * 1000)
	# Rough split: 70% input, 30% output
	tokens_input = int(total_tokens * 0.7)
	tokens_output = total_tokens - tokens_input

	# Fallback: try alternative token/cost patterns
	if tokens_input == 0 and tokens_output == 0:
	# Filter out metadata lines (e.g., "Repo-map: using 1024 tokens")
	content_lines = [
	line for line in combined.split("\n")
	if not re.search(r"repo.?map\|using \d+ tokens.*refresh", line, re.I)
	]
	filtered = "\n".join(content_lines)
	alt_tokens = _TOKEN_ALT_RE.findall(filtered)
	if alt_tokens:
	total = sum(int(t.replace(",", "")) for t in alt_tokens)
	if total >= 100: # Sanity check: real LLM usage > 100 tokens
	tokens_input = int(total * 0.7)
	tokens_output = total - tokens_input
	if cost_usd == 0.0:
	alt_costs = _COST_ALT_RE.findall(combined)
	if alt_costs:
	cost_usd = sum(float(c) for c in alt_costs)

	# Estimate cost from tokens using dynamic pricing lookup.
	if cost_usd == 0.0 and (tokens_input or tokens_output):
	rate_in, _, _, rate_out = get_pricing(model)
	cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000

	# Last resort: estimate from raw output length (better than $0)
	if cost_usd == 0.0 and (tokens_input + tokens_output) < 50 and len(combined) > 100:
	estimated_output_tokens = len(combined) // 4
	rate_in, _, _, rate_out = get_pricing(model)
	tokens_input = estimated_output_tokens * 3
	tokens_output = estimated_output_tokens
	cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000
	logger.warning("Cost estimated from output length (%d chars) — no token data found", len(combined))

	trajectory = [{"tool_name": "Bash", "input_preview": f"codex exec --model {model}"}]

	if duration >= ctx.time_limit_s:
	status = RunStatus.timeout
	elif returncode == 0:
	status = RunStatus.success
	else:
	status = RunStatus.failure

	return RunResult(
	status=status,
	duration_s=round(duration, 2),
	tokens_input=tokens_input,
	tokens_output=tokens_output,
	cost_usd=round(cost_usd, 6),
	diff_text=diff_text,
	trajectory=trajectory,
	raw_output=(stdout + stderr)[:50_000],
	)
	finally:
	shutil.rmtree(tmpdir, ignore_errors=True)