BenchGen Router Lite

PROBE_CALLS = 20
PROBE_PROMPT = "Reply with exactly: OK"
# Deliberately unsatisfiable within max_tokens. A compliant agent stops at the cap with its
# partial text intact; silent truncation looks identical to a short answer once it reaches the
# reward matrix, so it has to be caught here.
LENGTH_PROBE_PROMPT = "Count from 1 to 5000. One number per line. Do not stop early."
MAX_EMPTY_RATE = 0.05


async def check(agent, protocol) -> tuple[str, str]:
    empties = errors = 0
    usage_seen = False

    async with OpenRouterClient(timeout_s=protocol.timeout_s) as client:
        for _ in range(PROBE_CALLS):
            result = await client.chat(
                agent, [{"role": "user", "content": PROBE_PROMPT}], protocol
            )
            errors += bool(result.error)
            empties += bool(result.empty)
            usage_seen = usage_seen or result.completion_tokens is not None
        long = await client.chat(
            agent, [{"role": "user", "content": LENGTH_PROBE_PROMPT}], protocol
        )

    detail = f"empty={empties / PROBE_CALLS:.0%} err={errors / PROBE_CALLS:.0%}"
    over_cap = (
        long.completion_tokens is not None and long.completion_tokens > protocol.max_tokens
    )
    if errors:
        return "FAIL", f"{detail} (call errors)"
    if empties / PROBE_CALLS > MAX_EMPTY_RATE:
        return "FAIL", f"{detail} (empty rate above {MAX_EMPTY_RATE:.0%})"
    if not usage_seen:
        return "FAIL", f"{detail} (no token counts, so cost cannot be tracked)"
    # A probe that errored or came back empty measured nothing, so it cannot pass: the point is
    # positive evidence that the cap is honoured.
    if long.error or over_cap or not long.text.strip():
        return "FAIL", f"{detail} (ignored the {protocol.max_tokens} token cap)"
    return "PASS", detail

config = load_collection(VERSION)
request = BuildRequest(
    per_source_limit=dict(config["per_source_limit"]),
    max_share=float(config["max_domain_share"]),
    seed=int(config["seed"]),
    cache_dir=CACHE_DIR,
    ratios=config["split_ratios"],
)
result = build_pool(request)

out = tasks_path(VERSION)
print(f"wrote {write_jsonl(out, result.tasks)} tasks to {out}")

problems = gate_2(result, request.max_share)
if problems:
    print("\nGate 2 FAILED:")
    for problem in problems[:20]:
        print(f"  - {problem}")
    return 1
print("\nGate 2 passed: balanced, answered, graded, split.")

def effective_max_share(n_domains: int, max_share: float = DEFAULT_MAX_SHARE) -> float:
    """With `n` domains no allocation can put every domain below `1/n`.

    Asking for 30% across 3 domains is arithmetically impossible; the cap becomes 33.3% and the
    report says so rather than the build silently failing its own gate.
    """
    if n_domains <= 0:
        return max_share
    return max(max_share, 1.0 / n_domains)

def sample_per_domain(tasks: list[Task], per_domain: dict[str, int], seed: int) -> list[Task]:
    """Seeded and sorted, so the same seed always draws the same tasks."""
    buckets: dict[str, list[Task]] = defaultdict(list)
    for task in tasks:
        buckets[task.domain].append(task)

    chosen: list[Task] = []
    for domain, want in sorted(per_domain.items()):
        bucket = sorted(buckets.get(domain, []), key=lambda t: t.task_id)
        rng = random.Random(f"{seed}/{domain}")
        chosen.extend(rng.sample(bucket, min(want, len(bucket))))
    return sorted(chosen, key=lambda t: t.task_id)

require_verified(pool)  # refuses to start against a slug Stage 1 never verified

async with OpenRouterClient(timeout_s=pool.protocol.timeout_s) as client:
    try:
        progress = await collect_rewards(
            tasks=tasks,
            pool=pool,
            chat=lambda agent, messages: client.chat(agent, messages, pool.protocol),
            out_path=OUT_PATH,
            repetitions=repetitions,
            concurrency=int(config["concurrency"]),
            store_raw=bool(config["store_raw_response"]),
            resume=True,
            on_progress=report,
        )
    except ProtocolViolation as exc:
        # The agent's numbers are no longer comparable and its cost is unbounded.
        print(f"\nSTOPPED, protocol violation: {exc}")
        return 2
    except CollectionAborted as exc:
        # A row written now would record a billing failure as a wrong answer, and resume
        # would then skip that task forever.
        print(f"\nSTOPPED, account-level failure: {exc}")
        return 2

@model_validator(mode="after")
def _recompute_bests(self) -> RewardRecord:
    winners = [self.agent_order[i] for i in best_indices(self.mean_reward)]
    # Derived, never trusted from the file: a stale best_agents silently corrupts the gate.
    self.best_agents = winners
    self.is_tie = len(winners) > 1
    return self

metrics = compute_health(records)
outcome = evaluate(load_gates(VERSION), metrics)
write_report("dataset_health.md", render_health(metrics, outcome, {REWARDS.name: REWARDS}))

held_out = {t.task_id for t in tasks if t.split == HELD_OUT_SPLIT}
baselines = compute_baselines([r for r in records if r.task_id in held_out])
write_report(
    "baselines.md", render_baselines(baselines, HELD_OUT_SPLIT, {REWARDS.name: REWARDS})
)

return 1 if outcome.decision == "STOP" else 0

LEAK_MARKERS = ('"prompt"', '"raw_response"', '"answer"')


def assert_no_prompt_text(out_dir: Path) -> None:
    """Cheap guard against a future change re-introducing prompts into the export."""
    for path in out_dir.iterdir():
        if path.suffix not in {".json", ".jsonl"}:
            continue
        body = path.read_text(encoding="utf-8")
        if any(marker in body for marker in LEAK_MARKERS):
            raise SystemExit(f"refusing to publish: {path.name} appears to contain question text")

def head_param_count(in_dim: int, hidden_dim: int, n_agents: int) -> int:
    if hidden_dim <= 0:
        return in_dim * n_agents + n_agents
    return in_dim * hidden_dim + hidden_dim + hidden_dim * n_agents + n_agents


def build_head(theta, in_dim: int, hidden_dim: int, n_agents: int):
    """Unflatten a CMA-ES candidate vector into a callable head(x) -> logits."""
    if hidden_dim <= 0:
        w = theta[: in_dim * n_agents].reshape(in_dim, n_agents)
        b = theta[in_dim * n_agents :]

        def head(x):
            return x @ w + b

        return head
    # ... a one-hidden-layer MLP when head_hidden_dim > 0

def train_cma(E_train, R_train, in_dim: int, n_agents: int, cfg: dict):
    import cma

    hidden_dim = int(cfg.get("head_hidden_dim") or 0)
    n_params = head_param_count(in_dim, hidden_dim, n_agents)
    es = cma.CMAEvolutionStrategy(np.zeros(n_params), float(cfg.get("cma_sigma0", 0.5)), {
        "seed": int(cfg.get("cma_seed", 42)), "verbose": -9,
    })

    best_theta, best_score = np.zeros(n_params), -1.0
    while not es.stop() and gen < int(cfg.get("cma_generations", 60)):
        population = es.ask()
        neg_scores = []
        for theta in population:
            head = build_head(np.asarray(theta), in_dim, hidden_dim, n_agents)
            picks = np.argmax(head(E_train), axis=1)
            reward = float(R_train[np.arange(len(picks)), picks].mean())
            neg_scores.append(-reward)  # CMA-ES minimizes
        es.tell(population, neg_scores)
        # ... track the best theta seen so far, log it, gen += 1