mindxtrain/research/contract.py

"""Declarative autoresearch contract (discipline 2).

One `AttemptContract` decouples the keep/reset search loop from the benchmark: it
names the command that runs a single measured attempt, the one file the agent may
edit, and the scalar metric to optimise plus where to read it. The same loop drives
any recipe — CPU imprint, a real checkpoint's MEI, BFCL tool-call accuracy — by
swapping the contract, not the harness.

Contracts are hand-authored TOML under `mindxtrain/research/contracts/`; loading uses
stdlib `tomllib` (no new dependency). Pure stdlib + pydantic; base-install importable.
"""

from __future__ import annotations

import tomllib
from pathlib import Path
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field

# Where the optimised scalar is read from. Each value resolves to an existing
# in-tree eval; `metrics_jsonl` is the universal stdlib path (open-rl convention).
MetricSource = Literal[
    "mei",
    "classroom_delta",
    "bfcl",
    "tau_bench",
    "held_out_loss",
    "metrics_jsonl",
]


class AttemptContract(BaseModel):
    """Declarative contract for one autoresearch attempt.

    `command` runs one measured attempt and must exit nonzero on failure; the runner
    substitutes ``{run_dir}``, ``{run_name}`` and ``{attempt_timeout_minutes}``.
    `editable` is the *only* path the agent may change — the guard enforces it.
    `metric` names the scalar to optimise and `source` says where to read it.
    """

    model_config = ConfigDict(extra="forbid", frozen=True)

    task: str = Field(min_length=1, description="Human label for the search.")
    command: str = Field(min_length=1, description="Shell command running one attempt.")
    editable: list[Path] = Field(min_length=1, description="The only paths the agent may edit.")
    metric: str = Field(min_length=1, description="Name of the scalar to optimise.")
    source: MetricSource = "metrics_jsonl"
    mode: Literal["max", "min"] = "max"
    attempt_timeout_minutes: float = Field(default=10.0, gt=0)
    # Keep iff strictly better by this margin (mirrors the §8 MEI strict-higher gate).
    min_improvement: float = Field(default=0.0, ge=0)

    @classmethod
    def from_toml(cls, path: Path) -> AttemptContract:
        """Load and validate a contract from a TOML file.

        `tomllib.load` requires a binary handle. Unknown keys raise (extra=forbid),
        so a typo'd contract field fails loudly rather than being silently ignored.
        """
        with Path(path).open("rb") as fh:
            data = tomllib.load(fh)
        return cls.model_validate(data)


__all__ = ["AttemptContract", "MetricSource"]