code/local_train/dpo_pack.py

"""Chosen/rejected pairs for the live pre-eval policy, not the local 3-file ritual.

Chosen is always a unique sed -i. Rejected is the live ban: another marker echo,
sed && submit, recat, or grep-as-work.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from albedo_eval_service.remote.dataset import format_messages
from sanity_service.chain import followup_instruction, micro_instruction

from .chain_gold import reject_edit_thought
from .constants import DEFAULT_PACK_DIR, TOKENIZER_DIR
from .think import wrap_completion

CLAUSES = (
    ("echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT", "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT"),
    (
        "echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT && cat patch.txt",
        "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT",
    ),
    ("echo FINALIZE_AND_SUBMIT_TASK_OUTPUT", "FINALIZE_AND_SUBMIT_TASK_OUTPUT"),
    ("echo ALBEDO_TASK_DONE_SUBMIT_NOW", "ALBEDO_TASK_DONE_SUBMIT_NOW"),
    ("echo SUBMIT_TASK_AAAA && cat patch.txt", "SUBMIT_TASK_AAAA"),
)

# Live-ish and off-distribution targets. Not the three local gate files.
CASES = (
    {
        "id": "nosa-scikit",
        "kind": "nosa",
        "file": "tests/test_builder.py",
        "function": "test_wheel_tag_expand",
        "request": "fix test_wheel_tag_expand on mac_ver",
    },
    {
        "id": "mixed-lineprof",
        "kind": "mixed",
        "file": "line_profiler/line_profiler.py",
        "function": "add_function",
        "request": "edit add_function so the missed call site is covered",
    },
    {
        "id": "messy-cmake-url",
        "kind": "messy",
        "file": "cmake.org/cmake/help/git-master/manual/cmake-language.7.h",
        "function": "THOUGHT",
        "request": "make a small concrete edit in cmake-language.7.h",
    },
    {
        "id": "messy-testbed",
        "kind": "messy",
        "file": "/testbed/src/scikit_build_core/builder/wheel_tag.py",
        "function": "WheelTag",
        "request": "touch WheelTag in wheel_tag.py",
    },
    {
        "id": "recat-cmakelists",
        "kind": "recat",
        "file": "/testbed/CMakeLists.txt",
        "function": "GCC",
        "request": "edit CMakeLists.txt so GCC is not forced to libc++",
    },
    {
        "id": "after-submit-urls",
        "kind": "after_submit",
        "file": "testbed/example/api/urls.py",
        "function": "urls",
        "request": "edit urls.py at the missed call site",
    },
)

# Train-only extras. Policy eval stays on CASES so v16/v17 stay comparable.
TRAIN_EXTRA = (
    {
        "id": "mixed-builder",
        "kind": "mixed",
        "file": "tests/test_builder.py",
        "function": "test_wheel_tag_expand",
        "request": "fix test_wheel_tag_expand on mac_ver",
    },
    {
        "id": "recat-follow-cmake",
        "kind": "recat_follow",
        "file": "/testbed/CMakeLists.txt",
        "function": "GCC",
        "request": "edit CMakeLists.txt so GCC is not forced to libc++",
        "leftover": "264: -stdlib=libc++\n265: if(APPLE)\n",
    },
)

# Exact random-6 fails: after a legal submit, another marker with no sed.
NARROW_CASES = (
    {
        "id": "after-aiohttp-pyproject",
        "kind": "after_submit",
        "file": "workspace/aio-libs__aiohttp__1.0/pyproject.toml",
        "function": "Important",
        "request": "make a small concrete edit in pyproject.toml",
    },
    {
        "id": "after-defaults-toml",
        "kind": "after_submit",
        "file": "testbed/cpg_workflows/defaults.toml",
        "function": "defaults",
        "request": "make a small concrete edit in defaults.toml",
    },
    {
        "id": "late-defaults-toml",
        "kind": "late_after_submit",
        "file": "testbed/cpg_workflows/defaults.toml",
        "function": "defaults",
        "request": "make a small concrete edit in defaults.toml",
    },
    {
        "id": "after-pyproject",
        "kind": "after_submit",
        "file": "pyproject.toml",
        "function": "Important",
        "request": "make a small concrete edit in pyproject.toml",
    },
    {
        "id": "after-cargo-toml",
        "kind": "after_submit",
        "file": "testbed/Cargo.toml",
        "function": "package",
        "request": "make a small concrete edit in Cargo.toml",
    },
)

_KIND_COPIES = {
    "nosa": 1,
    "mixed": 1,
    "messy": 1,
    "recat": 1,
    "recat_follow": 1,
    "after_submit": 1,
}


@dataclass(frozen=True)
class PolicyPrefix:
    sample_id: str
    kind: str
    messages: list[dict[str, str]]
    chosen_bash: str
    rejected_bashes: tuple[str, ...]
    submit_command: str
    submit_marker: str
    file: str
    function: str


def strip_open_think(prompt: str, completion: str) -> str:
    """Chat template already opened <think>; do not emit a second one."""
    if prompt.endswith("<think>\n") and completion.startswith("<think>\n"):
        return completion[len("<think>\n") :]
    return completion


def _wrap(thought: str, bash: str) -> str:
    wrapped = wrap_completion(f"THOUGHT: {thought}\n```bash\n{bash}\n```", bash)
    if not wrapped:
        raise ValueError(f"failed to wrap {bash!r}")
    return wrapped


def _asst(thought: str, bash: str) -> dict[str, str]:
    return {"role": "assistant", "content": _wrap(thought, bash)}


def _user(text: str) -> dict[str, str]:
    return {"role": "user", "content": text}


def _comment(path: str) -> str:
    base = path.rsplit("/", 1)[-1]
    if base == "CMakeLists.txt" or base.endswith((".py", ".cmake", ".toml", ".rb", ".h")):
        return "#"
    return "//"


def _sed(path: str, func: str, line: int, tag: str) -> str:
    mark = _comment(path)
    return f"sed -i '{line}a\\{mark} chain-{tag} {func} n{line}' {path}"


def _clause_tag(clause: str) -> str:
    tail = "patch" if "patch.txt" in clause else "bare"
    marker = clause.split()[1] if clause.startswith("echo ") else clause.split()[0]
    return f"{marker}-{tail}"


def _rejected_bashes(clause: str, path: str, *, echo_only: bool = False) -> tuple[str, ...]:
    token = clause.split()[1] if clause.startswith("echo ") else clause.split()[0]
    bare = f"echo {token}"
    if echo_only:
        seen = []
        for bash in (clause, bare):
            if bash and bash not in seen:
                seen.append(bash)
        return tuple(seen)
    mixed = f"{_sed(path, 'helper', 3, 'mix')} && {clause}"
    recat = f"cat {path}"
    grep = f"grep -n chain {path}"
    seen: list[str] = []
    for bash in (clause, bare, mixed, recat, grep):
        if bash and bash not in seen:
            seen.append(bash)
    return tuple(seen)


def _wrong_thought(clause: str, kind: str) -> str:
    if kind == "recat":
        return (
            "The leftover dump is the original issue. Work is done. "
            f"Submit now with the exact command: {clause}"
        )
    if kind == "mixed":
        return (
            "Already edited and the user named the clause. "
            f"Submit again: {clause}"
        )
    return f"The user named the submit clause. Submit now with the exact command: {clause}"


def _chosen_thought(path: str, func: str, clause: str, kind: str) -> str:
    base = reject_edit_thought(path, func, clause)
    if kind == "recat":
        return (
            "The leftover dump looks like the original PR. That is not submit-now "
            "and not a reason to cat again. grep/cat/ls are not work. "
            f"Unique sed -i in {path} ({func}) this turn. Never echo the marker."
        )
    if kind == "mixed":
        return (
            "sed && echo already used the submit. Another marker is "
            "'repeated submissions without doing any work'. "
            + base
        )
    return base


def build_prefix(case: dict[str, str], clause: str, marker: str) -> PolicyPrefix:
    path = case["file"]
    func = case["function"]
    micro = {
        "file": path,
        "function": func,
        "request": case["request"],
        "message": "",
    }
    instruction = micro_instruction(micro, clause)
    follow = followup_instruction(
        "Please handle the missed call site. I do not see a new edit since your last submit.",
        clause,
        first=False,
    )
    thanks = followup_instruction(
        "Thanks, received. Continue with the original issue; do not submit again until you make a new edit.",
        clause,
        first=True,
    )
    leftover = case.get(
        "leftover",
        "264: -stdlib=libc++\n265: if(APPLE)\n",
    )
    grep_obs = "12: maxsize = 2**32\n"
    cat1 = f"cat {path}"
    sed1 = _sed(path, func, 1, "a")
    chosen = _sed(path, func, 9, "late")
    echo = clause
    grep = f"grep -n maxsize {path}"
    recat = f"cat -n {path} | grep {func}"
    mixed = f"{sed1} && {clause}"

    kind = case["kind"]
    if kind == "nosa":
        messages = [
            _user(instruction),
            _asst(f"Unique sed -i in {path} ({func}).", sed1),
            _user(leftover),
            _asst("Work is saved. This turn is ONLY the exact submit command.", echo),
            _user(follow),
            _asst("Search for the missed call site.", grep),
            _user(grep_obs),
            _asst("Read the function body once.", recat),
            _user(follow),
        ]
    elif kind == "mixed":
        messages = [
            _user(instruction),
            _asst(f"Edit {func} and submit in one line.", mixed),
            _user(follow),
        ]
    elif kind == "recat":
        messages = [
            _user(instruction),
            _asst(f"Read {path} once before editing. After this, never cat again.", cat1),
            _user(leftover),
        ]
    elif kind == "recat_follow":
        messages = [
            _user(instruction),
            _asst(f"Read {path} once before editing. After this, never cat again.", cat1),
            _user(leftover),
            _asst(f"Unique sed -i in {path} ({func}). Leftover dump is not submit-now.", sed1),
            _user(follow),
        ]
        kind = "recat"
    elif kind == "after_submit":
        messages = [
            _user(instruction),
            _asst(f"Unique sed -i in {path} ({func}).", sed1),
            _user("command completed with no captured output"),
            _asst("Work is saved. This turn is ONLY the exact submit command.", echo),
            _user(thanks),
        ]
    elif kind == "late_after_submit":
        sed2 = _sed(path, func, 2, "b")
        sed3 = _sed(path, func, 3, "c")
        messages = [
            _user(instruction),
            _asst(f"Read {path} once.", cat1),
            _user(leftover),
            _asst(f"Unique sed -i in {path} ({func}).", sed1),
            _user("command completed with no captured output"),
            _asst("Work is saved. This turn is ONLY the exact submit command.", echo),
            _user(thanks),
            _asst(f"Unique sed -i in {path} ({func}) line 2.", sed2),
            _user("command completed with no captured output"),
            _asst(f"Unique sed -i in {path} ({func}) line 3.", sed3),
            _user("command completed with no captured output"),
            _asst("Work is saved. This turn is ONLY the exact submit command.", echo),
            _user(follow),
        ]
        kind = "after_submit"
    else:
        messages = [
            _user(instruction),
            _asst(f"Read {path} once.", cat1),
            _user(leftover),
            _asst(f"Unique sed -i in {path} ({func}).", sed1),
            _user("command completed with no captured output"),
            _asst("Work is saved. This turn is ONLY the exact submit command.", echo),
            _user(follow),
        ]
    return PolicyPrefix(
        sample_id=f"{case['id']}:{_clause_tag(clause)}",
        kind=kind,
        messages=messages,
        chosen_bash=chosen,
        rejected_bashes=_rejected_bashes(
            clause, path, echo_only=kind in {"after_submit", "late_after_submit"}
        ),
        submit_command=clause,
        submit_marker=marker,
        file=path,
        function=func,
    )


def iter_prefixes() -> list[PolicyPrefix]:
    return [build_prefix(case, clause, marker) for case in CASES for clause, marker in CLAUSES]


def iter_train_prefixes(*, narrow: bool = False) -> list[PolicyPrefix]:
    if narrow:
        return [build_prefix(case, clause, marker) for case in NARROW_CASES for clause, marker in CLAUSES]
    prefixes: list[PolicyPrefix] = []
    for case in (*CASES, *TRAIN_EXTRA):
        for clause, marker in CLAUSES:
            prefix = build_prefix(case, clause, marker)
            copies = _KIND_COPIES.get(case["kind"], 1)
            prefixes.extend([prefix] * copies)
    return prefixes


def pack_dpo(
    *,
    out_dir: Path = DEFAULT_PACK_DIR,
    seed: str = "dpo-live-fault-v16",
    tokenizer_path: Path | None = None,
    narrow: bool = False,
) -> Path:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    tokenizer = str(tokenizer_path or TOKENIZER_DIR)
    rows: list[dict[str, Any]] = []
    for prefix in iter_train_prefixes(narrow=narrow):
        prompt = format_messages(
            prefix.messages, tokenizer_path=tokenizer, enable_thinking=True
        )
        chosen = strip_open_think(
            prompt,
            _wrap(
                _chosen_thought(prefix.file, prefix.function, prefix.submit_command, prefix.kind),
                prefix.chosen_bash,
            ),
        )
        for index, rejected_bash in enumerate(prefix.rejected_bashes):
            if rejected_bash == prefix.chosen_bash:
                continue
            rejected = strip_open_think(
                prompt, _wrap(_wrong_thought(prefix.submit_command, prefix.kind), rejected_bash)
            )
            rows.append(
                {
                    "sample_id": f"{prefix.sample_id}:rej{index}",
                    "prompt": prompt,
                    "chosen": chosen,
                    "rejected": rejected,
                    "kind": prefix.kind,
                    "submit_command": prefix.submit_command,
                    "submit_marker": prefix.submit_marker,
                    "file": prefix.file,
                    "chosen_bash": prefix.chosen_bash,
                    "rejected_bash": rejected_bash,
                }
            )
    pack_path = out_dir / f"{seed}.jsonl"
    with pack_path.open("w") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    kinds: dict[str, int] = {}
    for row in rows:
        kinds[row["kind"]] = kinds.get(row["kind"], 0) + 1
    meta = {"n": len(rows), "kinds": kinds, "pack": str(pack_path)}
    (out_dir / f"{seed}.meta.json").write_text(json.dumps(meta, indent=2) + "\n")
    print(json.dumps(meta, indent=2), flush=True)
    print(f"dpo pack: {pack_path}", flush=True)
    return pack_path