|
30 | 30 | from eval_protocol.mcp_servers.tau2 import get_server_script_path, get_system_prompt |
31 | 31 |
|
32 | 32 |
|
| 33 | +def _ensure_tau2_databases(): |
| 34 | + """Ensure tau2 database files exist, downloading if necessary.""" |
| 35 | + import os |
| 36 | + import urllib.request |
| 37 | + from pathlib import Path |
| 38 | + |
| 39 | + # Get the vendor/tau2/data directory path |
| 40 | + try: |
| 41 | + from vendor.tau2.utils.utils import DATA_DIR |
| 42 | + |
| 43 | + domains_dir = DATA_DIR / "domains" |
| 44 | + except ImportError: |
| 45 | + # Fallback: find vendor/tau2 relative to this file |
| 46 | + vendor_tau2 = Path(__file__).parent.parent.parent / "vendor" / "tau2" |
| 47 | + domains_dir = vendor_tau2 / "data" / "domains" |
| 48 | + |
| 49 | + # Database files to download |
| 50 | + databases = { |
| 51 | + "retail/db.json": "https://raw.githubusercontent.com/sierra-research/tau2-bench/40f46d3540dc95aca145ddecb0464fdd9a1e8c15/data/tau2/domains/retail/db.json", |
| 52 | + "airline/db.json": "https://raw.githubusercontent.com/sierra-research/tau2-bench/main/data/tau2/domains/airline/db.json", |
| 53 | + "mock/db.json": "https://raw.githubusercontent.com/sierra-research/tau2-bench/main/data/tau2/domains/mock/db.json", |
| 54 | + } |
| 55 | + |
| 56 | + for rel_path, url in databases.items(): |
| 57 | + file_path = domains_dir / rel_path |
| 58 | + if not file_path.exists(): |
| 59 | + print(f"📥 Downloading {rel_path} to {file_path}...") |
| 60 | + file_path.parent.mkdir(parents=True, exist_ok=True) |
| 61 | + try: |
| 62 | + urllib.request.urlretrieve(url, file_path) |
| 63 | + print(f"✅ Downloaded {rel_path} ({file_path.stat().st_size:,} bytes)") |
| 64 | + except Exception as e: |
| 65 | + print(f"❌ Failed to download {rel_path}: {e}") |
| 66 | + raise |
| 67 | + |
| 68 | + |
| 69 | +# Ensure databases are available before test runs |
| 70 | +_ensure_tau2_databases() |
| 71 | + |
| 72 | + |
33 | 73 | def _get_retail_dataset_path() -> str: |
34 | 74 | """Get the retail dataset file path.""" |
35 | 75 | return str(Path(__file__).parent / "data" / "retail_dataset.jsonl") |
|
0 commit comments