arena/environments/schema_migration/schema_migration/taskset.py

"""schema-migration: a migration is judged on the rows you did not look at.

`readings.value_text` holds numbers and units as one string. The migration has to become
`value_num REAL` plus `unit TEXT`, drop the original column, and lose nothing. The agent
sees five tidy rows; it is graded on forty it never sees, which carry thousands separators,
negatives, decimals, a unit containing a digit and a slash, a value with no unit, and
doubled internal whitespace.

The reward is the migration actually executing against a real database, so there is no
judge and no rubric — SQLite either recomposes the value or it does not. Row count is
scored separately from fidelity because otherwise `DELETE FROM readings` is a perfect
migration: nothing left to be wrong about.
"""

from __future__ import annotations

import re

from pydantic import Field

import verifiers.v1 as vf

from schema_migration.fixture import connect
from schema_migration.run import measure

SYSTEM = """You are writing a SQLite migration.

The table is:
  CREATE TABLE readings (id INTEGER PRIMARY KEY, probe TEXT NOT NULL, value_text TEXT NOT NULL);

Migrate it so that `readings` has columns id, probe, value_num (REAL) and unit (TEXT), and
no longer has value_text. `value_num` is the number in value_text; `unit` is what follows
it, trimmed, or an empty string when there is no unit.

Return ONLY SQL in one ```sql code block. It is run with executescript(), so several
statements are fine. Preserve every row and every id.

You are graded on rows you have not seen, from the same column. They are messier than the
ones below."""

_BLOCK = re.compile(r"```(?:sql)?\s*\n(.*?)```", re.DOTALL)


def parse_sql(reply: str) -> str:
    blocks = _BLOCK.findall(reply or "")
    return (blocks[-1] if blocks else (reply or "")).strip()


class MigrationData(vf.TaskData):
    seed: int
    held_out: int


class MigrationTask(vf.Task[MigrationData]):
    @vf.stop
    async def single_turn(self, trace: vf.Trace) -> bool:
        return trace.num_turns >= 1

    @vf.metric
    async def scan(self, trace: vf.Trace) -> dict[str, float]:
        outcome = measure(self.data.seed, self.data.held_out, parse_sql(trace.last_reply))
        return {
            "schema": float(outcome.schema_ok),
            "fidelity": outcome.fidelity,
            "rows_kept": outcome.rows_kept,
            # Rows are a MULTIPLIER on fidelity, not a reward beside it. Scored separately
            # they paid 0.15 for adding two empty columns and touching nothing — a floor
            # above zero for near-inaction, which the house rules forbid. As a multiplier,
            # losing rows still zeroes an otherwise perfect migration and preserving them
            # on its own earns nothing.
            "integrity": outcome.fidelity * outcome.rows_kept,
            "clean": float(outcome.clean),
            "errored": float(bool(outcome.error)),
        }

    @vf.reward(weight=0.30)
    async def schema(self, trace: vf.Trace) -> float:
        return trace.metrics.get("schema", 0.0)

    @vf.reward(weight=0.45)
    async def integrity(self, trace: vf.Trace) -> float:
        return trace.metrics.get("integrity", 0.0)

    @vf.reward(weight=0.25)
    async def gate(self, trace: vf.Trace) -> float:
        return trace.metrics.get("clean", 0.0)


class MigrationConfig(vf.TasksetConfig):
    num_tasks: int = Field(48, ge=1)
    visible: int = Field(5, ge=1)
    held_out: int = Field(40, ge=1)


class MigrationTaskset(vf.Taskset[MigrationTask, MigrationConfig]):
    SEED_BASE = 30_000

    def load(self) -> list[MigrationTask]:
        tasks = []
        for i in range(self.config.num_tasks):
            seed = self.SEED_BASE + i
            # The visible slice is tidy on purpose; the graded one is not.
            con, data = connect(seed, self.config.visible, awkward=False)
            sample = "\n".join(
                f"  {i + 1} | probe_{i:03d} | {text}" for i, (text, _, _) in enumerate(data)
            )
            con.close()
            tasks.append(
                MigrationTask(
                    MigrationData(
                        idx=i,
                        name=f"readings-{seed}",
                        prompt=f"Rows in the table:\n\n id | probe | value_text\n{sample}\n\nWrite the migration.",
                        system_prompt=SYSTEM,
                        seed=seed,
                        held_out=self.config.held_out,
                    ),
                    self.config.task,
                )
            )
        return tasks