hack-house/hh/scripts/bench/olympics/challenge.py

"""Challenge adapter — MBPP bootstrap with a public/hidden test split.

The arena consumes a ``Challenge``: a brief the referee posts to the room, a
*public* test the team may run during TEST, and a *hidden* test held out for
SUBMIT grading (so teams can't teach to the test — SPEC §8/§10).

M1 implements the MBPP-Python adapter only (the spine; multi-language MultiPL-E
splits arrive in M3). It reuses ``bench/datasets.py`` for rows and
``bench/langs.py`` for the execution recipe; the per-assert public/hidden split
and prompt synthesis live here.

Split policy (SPEC §16 decision): reveal exactly ONE assert from ``test_list``
as the public example; hold the remainder as hidden. If a problem ships only one
assert, that single assert is both the public guide and the hidden gate.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable

from .. import datasets
from ..langs import Lang, resolve

_MBPP_PY_STOP = ["\nassert", "\nprint(", "\n# Test", "\nif __name__", '\n"""']


@dataclass
class Challenge:
    id: str
    language: str
    lang: Lang
    text: str                       # natural-language task
    public_tests: list[str]         # asserts revealed to the team
    hidden_tests: list[str]         # asserts held out for grading
    setup: str = ""                 # test_setup_code, prepended to programs
    reference: str = ""             # canonical solution (audit / novelty later)
    meta: dict = field(default_factory=dict)
    # ── MultiPL-E continuation mode (M3) ──────────────────────────────────
    # When ``prompt_prefix`` is set the challenge is *continuation*-style (a
    # runnable function prefix the model completes), not the MBPP-Python
    # description+asserts style. MultiPL-E ships one ``tests`` block per row with
    # no per-assert structure, so there is no public/hidden split: the program is
    # ``prefix + completion + tests`` (via ``_assemble_fn``) for both public and
    # hidden grading. ``public_tests``/``hidden_tests`` stay empty in this mode.
    prompt_prefix: str = ""         # runnable function prefix (MultiPL-E `prompt`)
    tests_block: str = ""           # the single MultiPL-E `tests` block (audit)
    stop: list[str] | None = None   # per-row stop tokens (MultiPL-E)
    _assemble_fn: Callable[[str, str, dict], str] | None = None
    _row: dict = field(default_factory=dict)

    @property
    def _continuation(self) -> bool:
        return bool(self.prompt_prefix)

    # ── what the room sees ────────────────────────────────────────────────
    def brief(self) -> str:
        if self._continuation:
            return (
                f"CHALLENGE {self.id} ({self.language}). Complete this "
                f"{self.language} function so the hidden tests pass:\n"
                f"{self.prompt_prefix.rstrip()}\n"
                f"Deliver one self-contained {self.language} solution that "
                f"continues the prefix above. Hidden tests grade it at SUBMIT."
            )
        example = self.public_tests[0] if self.public_tests else ""
        return (
            f"CHALLENGE {self.id} ({self.language}). Implement a solution to:\n"
            f"  {self.text.strip()}\n"
            f"Example test that must pass:\n  {example}\n"
            f"Deliver one self-contained {self.language} solution. "
            f"Hidden tests will grade it at SUBMIT."
        )

    # ── what the implementer model is asked to continue (raw, no hidden) ───
    def gen_prompt(self) -> str:
        if self._continuation:
            return self.prompt_prefix
        asserts = "\n".join(self.public_tests)
        return f'"""\n{self.text.strip()}\n\n{asserts}\n"""\n'

    def stop_tokens(self) -> list[str]:
        if self.stop is not None:
            return self.stop
        return _MBPP_PY_STOP

    # ── assemble a runnable program for a given test set ──────────────────
    def _assemble(self, completion: str, asserts: list[str]) -> str:
        body = "\n".join(asserts)
        return f"{self.setup}\n{completion}\n{body}\n"

    def public_program(self, completion: str) -> str:
        if self._continuation:
            return self._assemble_fn(self.prompt_prefix, completion, self._row)
        return self._assemble(completion, self.public_tests)

    def hidden_program(self, completion: str) -> str:
        if self._continuation:
            # single tests block — public == hidden for MultiPL-E.
            return self._assemble_fn(self.prompt_prefix, completion, self._row)
        # grade on the full spec (public example + held-out asserts) so a
        # solution that only satisfies the revealed example still fails.
        return self._assemble(completion, self.public_tests + self.hidden_tests)


def _split(test_list: list[str]) -> tuple[list[str], list[str]]:
    if not test_list:
        return [], []
    if len(test_list) == 1:
        return [test_list[0]], [test_list[0]]
    return [test_list[0]], test_list[1:]


def _doc_line(prompt: str) -> str:
    """Best-effort one-liner from a MultiPL-E prompt's leading comment block."""
    for raw in prompt.splitlines():
        s = raw.lstrip("#/ \t").strip()
        if s and not s.startswith("!") and "(" not in s:
            return s
    return prompt.strip().splitlines()[0] if prompt.strip() else ""


def _mbpp_task_id(name: str) -> int | None:
    # MultiPL-E mbpp names look like "mbpp_3_is_not_prime".
    parts = name.split("_")
    if len(parts) >= 2 and parts[0] == "mbpp" and parts[1].isdigit():
        return int(parts[1])
    return None


def load_multipl_e(task_id: int | None = None, *, index: int = 0,
                   language: str = "javascript",
                   suite: str = "mbpp") -> Challenge:
    """Build one MultiPL-E continuation challenge for a non-Python language.

    Reuses ``suites.binding`` for the (suite, language) format recipe. Pick by
    embedded MBPP ``task_id`` (e.g. ``mbpp_11_...``) if given, else by ``index``
    into the cached split. The row's single ``tests`` block grades both public
    and hidden (MultiPL-E has no per-assert split — see ``Challenge``)."""
    from ..suites import binding as _binding
    lang = resolve(language)
    b = _binding(suite, language)
    rows = datasets.load(b.dataset, b.config, split="test")
    if task_id is not None:
        row = next((r for r in rows
                    if _mbpp_task_id(r.get("name", "")) == task_id), None)
        if row is None:
            raise KeyError(f"{suite} task_id {task_id} not in {b.config} split")
    else:
        row = rows[index]
    name = row.get("name", f"{suite}-{lang.id}-{index}")
    tid = _mbpp_task_id(name)
    prompt = b.build_prompt(row)
    return Challenge(
        id=f"{suite}-{lang.id}-{tid if tid is not None else index}",
        language=lang.id, lang=lang, text=_doc_line(prompt),
        public_tests=[], hidden_tests=[],
        reference=row.get("original", "") or "",
        prompt_prefix=prompt, tests_block=row.get("tests", "") or "",
        stop=b.stop_tokens(row), _assemble_fn=b.assemble, _row=row,
        meta={"name": name, "suite": suite, "config": b.config,
              "task_id": tid})


def load_mbpp(task_id: int | None = None, *, index: int = 0,
              language: str = "python") -> Challenge:
    """Build one MBPP challenge. Pick by ``task_id`` if given, else by ``index``
    into the cached test split. Non-Python languages route to the MultiPL-E
    continuation adapter (``load_multipl_e``)."""
    if resolve(language).id != "python":
        return load_multipl_e(task_id, index=index, language=language,
                              suite="mbpp")
    lang = resolve(language)
    rows = datasets.load("google-research-datasets/mbpp", "full")
    if task_id is not None:
        row = next((r for r in rows if r.get("task_id") == task_id), None)
        if row is None:
            raise KeyError(f"MBPP task_id {task_id} not in cached split")
    else:
        row = rows[index]
    public, hidden = _split(row.get("test_list", []))
    tid = row.get("task_id", index)
    return Challenge(
        id=f"mbpp-{tid}", language=lang.id, lang=lang,
        text=row.get("text", ""), public_tests=public, hidden_tests=hidden,
        setup=row.get("test_setup_code", "") or "",
        reference=row.get("code", ""),
        meta={"task_id": tid, "n_tests": len(row.get("test_list", []))})