tenant.openagents/omega

1"""Benchmark registry.
2
3A `Benchmark` describes everything the orchestrator needs to run one benchmark
4end-to-end that is *not* about the build source or the model: which harness runs
5it (Harbor or its Pier fork), how its dataset is provisioned, how a trial is
6scored, whether it needs an LLM judge, and the per-task timeout.
7
8Adding a benchmark should be a data change here, not new control flow. The three
9benchmark families currently supported are all Harbor-family harnesses:
10
11  - SWE-Atlas (qna / rf / tw): Harbor, rubric LLM judge.
12  - Terminal-Bench 2.1: Harbor, test-script scoring, no judge.
13  - DeepSWE: Pier (a Harbor fork with per-agent network allowlists, required
14    because DeepSWE tasks run with `allow_internet = false`), test scoring.
15"""
16
17from __future__ import annotations
18
19from dataclasses import dataclass, field
20
21# How a benchmark's dataset is made available to the harness.
22#   registry -> `-d <name>` pulled from the harness hub
23#   path     -> `-p <dir>` from a git repo cloned by the controller
24#   pier_path -> path dataset run under Pier instead of Harbor
25DATASET_REGISTRY = "registry"
26DATASET_PATH = "path"
27DATASET_PIER_PATH = "pier_path"
28
29HARNESS_HARBOR = "harbor"
30HARNESS_PIER = "pier"
31
32SCORING_RUBRIC_JUDGE = "rubric-judge"
33SCORING_TESTS = "tests"
34
35
36@dataclass(frozen=True)
37class DatasetRef:
38    kind: str
39    # For registry datasets: the hub dataset name. For path datasets: the repo
40    # URL plus a ref and the sub-directory the tasks live in.
41    name: str | None = None
42    repo_url: str | None = None
43    repo_ref: str | None = None
44    data_dir: str | None = None
45
46
47@dataclass(frozen=True)
48class Benchmark:
49    id: str
50    label: str
51    harness: str
52    dataset: DatasetRef
53    default_timeout_secs: int
54    scoring: str
55    needs_judge: bool
56    default_judge: str | None = None
57    # Hosts the in-sandbox agent must reach even on air-gapped tasks. Required by
58    # Pier benchmarks (DeepSWE) so eval-cli can still call the model API.
59    network_allowlist: tuple[str, ...] = ()
60    env: dict[str, str] = field(default_factory=dict)
61
62
63SWE_ATLAS_REPO_URL = "https://github.com/scaleapi/SWE-Atlas.git"
64SWE_ATLAS_REPO_REF = "main"
65DEEPSWE_REPO_URL = "https://github.com/datacurve-ai/deep-swe.git"
66DEEPSWE_REPO_REF = "main"
67
68# Model/judge API hosts the agent needs even under `allow_internet = false`.
69# Pier grants these to the agent via its per-agent allowlist.
70AGENT_API_HOSTS: tuple[str, ...] = (
71    "api.anthropic.com",
72    "api.openai.com",
73    "inference.baseten.co",
74)
75
76
77BENCHMARKS: dict[str, Benchmark] = {
78    "swe-atlas-qna": Benchmark(
79        id="swe-atlas-qna",
80        label="SWE-Atlas Codebase Q&A",
81        harness=HARNESS_HARBOR,
82        dataset=DatasetRef(kind=DATASET_REGISTRY, name="scale-ai/swe-atlas-qna"),
83        default_timeout_secs=7200,
84        scoring=SCORING_RUBRIC_JUDGE,
85        needs_judge=True,
86        default_judge="deepseek-v4-pro",
87    ),
88    "swe-atlas-rf": Benchmark(
89        id="swe-atlas-rf",
90        label="SWE-Atlas Refactoring",
91        harness=HARNESS_HARBOR,
92        dataset=DatasetRef(kind=DATASET_REGISTRY, name="scale-ai/swe-atlas-rf"),
93        default_timeout_secs=3300,
94        scoring=SCORING_RUBRIC_JUDGE,
95        needs_judge=True,
96        default_judge="kimi-k2.7-code",
97    ),
98    "swe-atlas-tw": Benchmark(
99        id="swe-atlas-tw",
100        label="SWE-Atlas Test Writing",
101        harness=HARNESS_HARBOR,
102        # Test writing is not consistently published in the same registry shape as qna/rf.
103        dataset=DatasetRef(
104            kind=DATASET_PATH,
105            repo_url=SWE_ATLAS_REPO_URL,
106            repo_ref=SWE_ATLAS_REPO_REF,
107            data_dir="data/tw",
108        ),
109        default_timeout_secs=3300,
110        scoring=SCORING_RUBRIC_JUDGE,
111        needs_judge=True,
112        default_judge="kimi-k2.7-code",
113    ),
114    "terminal-bench-2.1": Benchmark(
115        id="terminal-bench-2.1",
116        label="Terminal-Bench 2.1",
117        harness=HARNESS_HARBOR,
118        dataset=DatasetRef(
119            kind=DATASET_REGISTRY, name="terminal-bench/terminal-bench-2-1"
120        ),
121        default_timeout_secs=3300,
122        scoring=SCORING_TESTS,
123        needs_judge=False,
124        # Air-gapped tasks: the fetch/web-search tools can't reach the network,
125        # so disable them (via the agent profile in eval-cli) to stop the agent
126        # wasting its budget on tools that can only fail.
127        env={"ZED_EVAL_DISABLE_TOOLS": "fetch,search_web"},
128    ),
129    "deepswe": Benchmark(
130        id="deepswe",
131        label="DeepSWE",
132        harness=HARNESS_PIER,
133        dataset=DatasetRef(
134            kind=DATASET_PIER_PATH,
135            repo_url=DEEPSWE_REPO_URL,
136            repo_ref=DEEPSWE_REPO_REF,
137            # DeepSWE keeps its Harbor-compatible tasks at the repo root.
138            data_dir="tasks",
139        ),
140        default_timeout_secs=7200,
141        scoring=SCORING_TESTS,
142        needs_judge=False,
143        network_allowlist=AGENT_API_HOSTS,
144        # Air-gapped except the model API allowlist; fetch/web-search are useless.
145        env={"ZED_EVAL_DISABLE_TOOLS": "fetch,search_web"},
146    ),
147}
148
149
150# Groups expand to multiple benchmarks in one launch.
151BENCHMARK_GROUPS: dict[str, tuple[str, ...]] = {
152    "swe-atlas": ("swe-atlas-qna", "swe-atlas-rf", "swe-atlas-tw"),
153}
154
155# One short alias per benchmark whose canonical id is verbose.
156BENCHMARK_ALIASES: dict[str, str] = {
157    "qna": "swe-atlas-qna",
158    "rf": "swe-atlas-rf",
159    "tw": "swe-atlas-tw",
160    "tb21": "terminal-bench-2.1",
161}
162
163SWE_ATLAS_PART_BENCHMARKS: dict[str, str] = {
164    "qna": "swe-atlas-qna",
165    "rf": "swe-atlas-rf",
166    "tw": "swe-atlas-tw",
167}
168
169
170def get_benchmark(benchmark_id: str) -> Benchmark:
171    try:
172        return BENCHMARKS[benchmark_id]
173    except KeyError as error:
174        valid = ", ".join(sorted(BENCHMARKS))
175        raise ValueError(
176            f"unknown benchmark '{benchmark_id}' (valid: {valid})"
177        ) from error
178
179
180def is_benchmark_selector(selector: str) -> bool:
181    """Whether `selector` names a known benchmark id, alias, or group.
182
183    Used by `run` to validate benchmark positionals before preparing builds."""
184    normalized = selector.strip().lower()
185    return (
186        normalized in BENCHMARK_GROUPS
187        or normalized in BENCHMARKS
188        or normalized in BENCHMARK_ALIASES
189    )
190
191
192def resolve_benchmark_selector(selector: str) -> list[str]:
193    """Expands a user-facing selector (benchmark id, alias, or group) into the
194    concrete benchmark ids it refers to, preserving order."""
195    normalized = selector.strip().lower()
196    if normalized in BENCHMARK_GROUPS:
197        return list(BENCHMARK_GROUPS[normalized])
198    if normalized in BENCHMARKS:
199        return [normalized]
200    if normalized in BENCHMARK_ALIASES:
201        return [BENCHMARK_ALIASES[normalized]]
202    valid = ", ".join(sorted({*BENCHMARKS, *BENCHMARK_GROUPS, *BENCHMARK_ALIASES}))
203    raise ValueError(f"unknown benchmark '{selector}' (valid: {valid})")
204
205
206def resolve_benchmarks(selectors: list[str]) -> list[str]:
207    resolved: list[str] = []
208    for selector in selectors:
209        for part in selector.split(","):
210            part = part.strip()
211            if not part:
212                continue
213            for benchmark_id in resolve_benchmark_selector(part):
214                if benchmark_id not in resolved:
215                    resolved.append(benchmark_id)
216    return resolved
217
218
219def benchmark_metadata(benchmark: Benchmark) -> dict[str, object]:
220    """The self-describing block embedded in a run request so the controller and
221    harness-command builder need no separate registry lookup."""
222    return {
223        "id": benchmark.id,
224        "label": benchmark.label,
225        "harness": benchmark.harness,
226        "dataset": {
227            "kind": benchmark.dataset.kind,
228            "name": benchmark.dataset.name,
229            "repo_url": benchmark.dataset.repo_url,
230            "repo_ref": benchmark.dataset.repo_ref,
231            "data_dir": benchmark.dataset.data_dir,
232        },
233        "default_timeout_secs": benchmark.default_timeout_secs,
234        "scoring": benchmark.scoring,
235        "needs_judge": benchmark.needs_judge,
236        "default_judge": benchmark.default_judge,
237        "network_allowlist": list(benchmark.network_allowlist),
238        "env": dict(benchmark.env),
239    }
240