|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Run Argus agent on WebGen-Bench instances.""" |
| 3 | + |
| 4 | +from __future__ import annotations |
| 5 | + |
| 6 | +import argparse |
| 7 | +import datetime |
| 8 | +import json |
| 9 | +import logging |
| 10 | +import random |
| 11 | +import threading |
| 12 | +from concurrent.futures import ThreadPoolExecutor, as_completed |
| 13 | +from pathlib import Path |
| 14 | + |
| 15 | +from datasets import load_dataset |
| 16 | + |
| 17 | +from argus import Agent, AskWebAgentTool, WebAgent |
| 18 | +from argus.config import AgentConfig, Config, LLMConfig, WebAgentConfig |
| 19 | +from argus.data.content import Content |
| 20 | +from argus.llm.anthropic import AnthropicClient |
| 21 | +from argus.llm.base import LLMClient |
| 22 | +from argus.llm.openai import OpenAIClient |
| 23 | +from argus.tools.checklist import ChecklistTool |
| 24 | +from argus.tools.shell import ShellTool |
| 25 | + |
| 26 | +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") |
| 27 | +logger = logging.getLogger(__name__) |
| 28 | + |
| 29 | +DATASET_NAME = "luzimu/WebGen-Bench" |
| 30 | +# General-purpose web dev image with Node 20 + Python 3 |
| 31 | +DEFAULT_DOCKER_IMAGE = "node:20-bookworm" |
| 32 | + |
| 33 | + |
| 34 | +def _build_task(instruction: str, port: int, max_steps: int) -> str: |
| 35 | + """Build the agent prompt from a WebGen-Bench instruction.""" |
| 36 | + return f""" |
| 37 | +You are a skilled web developer. Your task is to build a fully working website from scratch |
| 38 | +based on the specification below. |
| 39 | +
|
| 40 | +The website MUST be accessible at http://localhost:{port} when you are done. |
| 41 | +Your working directory is /workspace — put all generated files there. |
| 42 | +
|
| 43 | +You have {max_steps} steps available. Use them wisely. |
| 44 | +
|
| 45 | +## Specification |
| 46 | +{instruction} |
| 47 | +""" |
| 48 | + |
| 49 | + |
| 50 | +def _get_source_files(shell: ShellTool) -> str: |
| 51 | + """Return a git-style diff of all files created under /workspace.""" |
| 52 | + if shell._container is None: |
| 53 | + return "" |
| 54 | + # git init + add everything + show staged diff to capture all new files |
| 55 | + result = shell._container.exec_run( |
| 56 | + [ |
| 57 | + "bash", |
| 58 | + "-c", |
| 59 | + ( |
| 60 | + "cd /workspace && git init -q && git add -A && " |
| 61 | + "git -c core.fileMode=false diff --staged" |
| 62 | + ), |
| 63 | + ] |
| 64 | + ) |
| 65 | + return result.output.decode("utf-8", errors="replace") |
| 66 | + |
| 67 | + |
| 68 | +def _build_llm(cfg: LLMConfig) -> LLMClient: |
| 69 | + provider = cfg.provider.lower() |
| 70 | + if provider == "anthropic": |
| 71 | + return AnthropicClient( |
| 72 | + model=cfg.model, |
| 73 | + api_key=cfg.api_key, |
| 74 | + base_url=cfg.base_url or None, |
| 75 | + temperature=cfg.temperature, |
| 76 | + ) |
| 77 | + if provider == "openai": |
| 78 | + return OpenAIClient( |
| 79 | + model=cfg.model, |
| 80 | + api_key=cfg.api_key, |
| 81 | + base_url=cfg.base_url or None, |
| 82 | + temperature=cfg.temperature, |
| 83 | + ) |
| 84 | + raise ValueError(f"Unknown LLM provider: {cfg.provider!r}. Supported: openai, anthropic") |
| 85 | + |
| 86 | + |
| 87 | +def _build_web_agent(cfg: WebAgentConfig, instance_id: str) -> WebAgent: |
| 88 | + llm = _build_llm(cfg.LLM) |
| 89 | + log_dir = Path(cfg.log_dir) / instance_id if cfg.log_dir else None |
| 90 | + return WebAgent( |
| 91 | + llm=llm, |
| 92 | + system_prompt=cfg.system_prompt, |
| 93 | + max_steps=cfg.max_steps, |
| 94 | + log_dir=log_dir, |
| 95 | + headless=cfg.headless, |
| 96 | + ) |
| 97 | + |
| 98 | + |
| 99 | +def _build_agent( |
| 100 | + cfg: AgentConfig, instance: dict, port: int, web_agent: WebAgent, docker_image: str |
| 101 | +) -> Agent: |
| 102 | + shell = ShellTool( |
| 103 | + docker_image, |
| 104 | + workdir="/workspace", |
| 105 | + remove_on_cleanup=True, |
| 106 | + port=port, |
| 107 | + ) |
| 108 | + log_dir = Path(cfg.log_dir) / instance["id"] if cfg.log_dir else None |
| 109 | + llm = _build_llm(cfg.LLM) |
| 110 | + |
| 111 | + agent = Agent( |
| 112 | + llm=llm, |
| 113 | + tools=[ChecklistTool(), shell, AskWebAgentTool(web_agent)], |
| 114 | + system_prompt=cfg.system_prompt, |
| 115 | + max_steps=cfg.max_steps, |
| 116 | + log_dir=log_dir, |
| 117 | + summarize_threshold=cfg.summarize_threshold, |
| 118 | + ) |
| 119 | + return agent |
| 120 | + |
| 121 | + |
| 122 | +def _run_instance( |
| 123 | + cfg: Config, |
| 124 | + instance: dict, |
| 125 | + port: int, |
| 126 | + docker_image: str, |
| 127 | + predictions_file, |
| 128 | + file_lock: threading.Lock, |
| 129 | +) -> None: |
| 130 | + instance_id = instance["id"] |
| 131 | + logger.info("=== %s starting (port %d) ===", instance_id, port) |
| 132 | + |
| 133 | + web_agent = _build_web_agent(cfg.web_agent, instance_id) |
| 134 | + agent = _build_agent(cfg.agent, instance, port, web_agent, docker_image) |
| 135 | + |
| 136 | + task = Content(text=_build_task(instance["instruction"], port, cfg.agent.max_steps)) |
| 137 | + agent.run(task) |
| 138 | + |
| 139 | + logger.info("=== %s finished, saving results ===", instance_id) |
| 140 | + shell = next((t for t in agent.tool_list if isinstance(t, ShellTool)), None) |
| 141 | + source_patch = _get_source_files(shell) if shell is not None else "" |
| 142 | + |
| 143 | + record = json.dumps( |
| 144 | + { |
| 145 | + "instance_id": instance_id, |
| 146 | + "model_name_or_path": "Argus", |
| 147 | + "source_patch": source_patch, |
| 148 | + "instruction": instance["instruction"], |
| 149 | + "ui_instruct": instance["ui_instruct"], |
| 150 | + } |
| 151 | + ) |
| 152 | + with file_lock: |
| 153 | + predictions_file.write(record + "\n") |
| 154 | + predictions_file.flush() |
| 155 | + logger.info("Prediction written for %s", instance_id) |
| 156 | + |
| 157 | + |
| 158 | +def main() -> None: |
| 159 | + parser = argparse.ArgumentParser(description="Run Argus on WebGen-Bench") |
| 160 | + parser.add_argument("--config", default="webgen_bench.yaml", help="Path to config YAML") |
| 161 | + parser.add_argument("--split", default="test", help="Dataset split (test / train)") |
| 162 | + parser.add_argument( |
| 163 | + "--instance-ids", |
| 164 | + nargs="*", |
| 165 | + metavar="ID", |
| 166 | + help="Run only the specified instance IDs (runs all if omitted)", |
| 167 | + ) |
| 168 | + parser.add_argument( |
| 169 | + "--workers", |
| 170 | + type=int, |
| 171 | + default=1, |
| 172 | + help="Number of parallel workers (default: 1)", |
| 173 | + ) |
| 174 | + parser.add_argument( |
| 175 | + "--docker-image", |
| 176 | + default=DEFAULT_DOCKER_IMAGE, |
| 177 | + help=f"Docker image for agent containers (default: {DEFAULT_DOCKER_IMAGE})", |
| 178 | + ) |
| 179 | + args = parser.parse_args() |
| 180 | + |
| 181 | + cfg = Config.from_yaml(Path(__file__).parent / args.config) |
| 182 | + |
| 183 | + dataset = load_dataset(DATASET_NAME, split=args.split) |
| 184 | + if args.instance_ids: |
| 185 | + keep = set(args.instance_ids) |
| 186 | + dataset = dataset.filter(lambda x: x["id"] in keep) |
| 187 | + |
| 188 | + logger.info("Running %d instances with %d worker(s)", len(dataset), args.workers) |
| 189 | + |
| 190 | + predictions_path = Path( |
| 191 | + f"webgen_predictions_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.jsonl" |
| 192 | + ) |
| 193 | + file_lock = threading.Lock() |
| 194 | + |
| 195 | + used_ports: set[int] = set() |
| 196 | + instance_ports: list[int] = [] |
| 197 | + for _ in dataset: |
| 198 | + port = random.randint(10000, 65535) |
| 199 | + while port in used_ports: |
| 200 | + port = random.randint(10000, 65535) |
| 201 | + used_ports.add(port) |
| 202 | + instance_ports.append(port) |
| 203 | + |
| 204 | + logger.info("Predictions will be written to %s", predictions_path) |
| 205 | + with predictions_path.open("a", encoding="utf-8") as predictions_file: |
| 206 | + with ThreadPoolExecutor(max_workers=args.workers) as executor: |
| 207 | + futures = { |
| 208 | + executor.submit( |
| 209 | + _run_instance, |
| 210 | + cfg, |
| 211 | + instance, |
| 212 | + port, |
| 213 | + args.docker_image, |
| 214 | + predictions_file, |
| 215 | + file_lock, |
| 216 | + ): instance["id"] |
| 217 | + for instance, port in zip(dataset, instance_ports) |
| 218 | + } |
| 219 | + for future in as_completed(futures): |
| 220 | + instance_id = futures[future] |
| 221 | + try: |
| 222 | + future.result() |
| 223 | + except Exception: |
| 224 | + logger.exception("Instance %s failed", instance_id) |
| 225 | + |
| 226 | + |
| 227 | +if __name__ == "__main__": |
| 228 | + main() |
0 commit comments