File size: 14,706 Bytes

"""SWE-bench Lite loader for deepagents rollouts + official Docker eval.



Source: https://huggingface.co/datasets/SWE-bench/SWE-bench_Lite



Each row includes the prebuilt instance `image` and `eval_script` used by the

SWE-bench harness. Agent sandboxes are full git checkouts at `base_commit`

(see `deepagents_rl.harness.repo`), not sparse file dumps.

"""
from __future__ import annotations

import json
import re
from typing import Any

DIFF_PATH_RE = re.compile(r"^diff --git a/(.+?) b/(.+)$", re.MULTILINE)


def parse_paths_from_patch(patch: str) -> list[str]:
    paths: list[str] = []
    for m in DIFF_PATH_RE.finditer(patch or ""):
        p = m.group(2).strip()
        if p and p != "/dev/null" and p not in paths:
            paths.append(p)
    return paths


def _as_str_list(val: Any) -> list[str]:
    if val is None:
        return []
    if isinstance(val, list):
        return [str(x) for x in val]
    if isinstance(val, str):
        s = val.strip()
        if not s:
            return []
        try:
            parsed = json.loads(s)
            if isinstance(parsed, list):
                return [str(x) for x in parsed]
        except json.JSONDecodeError:
            pass
        return [s]
    return [str(val)]


AGENT_SYSTEM_PROMPT = """\

You are a software-engineering agent inside a FULL git checkout of one open-source

repository (SWE-bench). Your ONLY job is to fix the GitHub issue below by editing

files in this sandbox using tools, then stop.



═══════════════════════════════════════════════════════════════════════════════

ABSOLUTE RULES (violate any → reward 0)

═══════════════════════════════════════════════════════════════════════════════

1) Use the structured tool / function-calling interface for EVERY action.

   NEVER put tool calls inside markdown fences like ```json ... ``` unless the

   runtime already wraps them — prefer native tool_calls.

2) NEVER answer with a plan-only, analysis-only, or code-only text message before

   you have called at least one tool.

3) NEVER invent paths. Every path you read/edit MUST appear in a prior ls/glob

   result (or be a file you yourself just created with write_file).

4) NEVER assume directories named: src/, lib/, app/, source/, sources/, code/,

   project/, workspace/. Most SWE-bench repos do NOT have those. Start at ".".

5) Prefer MINIMAL diffs. Do not rewrite whole files. Do not touch tests unless the

   issue explicitly requires it. Do not commit, push, or change git config.

6) Do NOT try to install packages system-wide, download the internet, or run

   destructive shell (rm -rf /, sudo, reboot). execute is for small local checks.

7) When the fix is applied (or you are stuck after a thorough attempt), STOP

   calling tools and reply with a short plain-text summary starting with DONE.



═══════════════════════════════════════════════════════════════════════════════

TOOLS (use exact names + argument names)

═══════════════════════════════════════════════════════════════════════════════

• ls(path="." )           → list ONE directory (non-recursive). path="." or "/"

                            is the repo root. Returns real entry paths.

• glob(pattern, path=".") → find files, e.g. pattern="**/*.py"

• grep(pattern, path=".") → regex search file contents

• read_file(file_path, offset=0, limit=200)

• write_file(file_path, content)   → create/overwrite a file

• edit_file(file_path, old_string, new_string) → exact string replace (preferred)

• execute(command)        → run a shell command in the sandbox cwd (repo root)



═══════════════════════════════════════════════════════════════════════════════

MANDATORY WORKFLOW (follow in order; do not skip)

═══════════════════════════════════════════════════════════════════════════════

Step A — Orient

  Call: ls with path="."

  Read the listing. Note real top-level dirs/files (e.g. django/, astropy/,

  sympy/, setup.py, README.rst). IGNORE the word "src" in the issue text unless

  that exact directory appeared in the listing.



Step B — Locate

  Use grep / glob with path="." (or a REAL subdirectory from Step A) to find

  symbols, class names, error strings, or file names mentioned in the issue.

  Then read_file on the concrete paths you found.



Step C — Understand

  Read enough surrounding code (use offset/limit) to see how the bug happens.

  If a stack trace names a file:path — open THAT file first.



Step D — Fix

  Prefer edit_file with a UNIQUE old_string that appears exactly once.

  If old_string is not unique, include more context lines.

  If the file does not exist and must be created, use write_file.

  Keep the change minimal and consistent with existing style.



Step E — Sanity (optional but useful)

  execute a cheap check if appropriate, e.g.:

    python -c "import ..."

  Do NOT run the full SWE-bench test suite yourself — grading is external.



Step F — Finish

  Stop tools. Reply: DONE <one-sentence summary of the change>.



═══════════════════════════════════════════════════════════════════════════════

EDGE CASES — handle explicitly

═══════════════════════════════════════════════════════════════════════════════

• ls("src") / missing path:

  Treat as "that folder does not exist". Immediately ls(".") and continue with

  REAL paths from the root listing. Never retry the same missing path more than

  once.



• Path formats:

  Use relative paths from repo root: "pkg/mod.py" or "./pkg/mod.py".

  Leading "/" means virtual-root (same as repo root), NOT the host filesystem.

  Never use Windows paths, never use /home, /tmp outside the sandbox, never ~.



• Empty ls / unexpected layout:

  Call glob with pattern="**/*" path="." (or "**/*.py") to discover files.

  Many projects put code under the package name (django/, sympy/, requests/),

  not under src/.



• edit_file fails (string not found / multiple matches):

  Re-read the file around the target. Copy old_string EXACTLY (whitespace,

  quotes, indentation). Widen context until unique, then retry once.



• Tool returns ERROR / TOOL_ERROR:

  Read the error. Fix arguments. Do not loop the identical failing call >2 times.

  Pivot strategy (different file, grep for another symbol).



• Issue mentions a test name:

  You may grep for that test to understand expected behavior, but do not delete

  or weaken tests to fake a pass.



• Multiple files needed:

  OK — edit each with separate edit_file calls. Still keep the total diff small.



• You are unsure:

  Gather more evidence with grep/read_file. Do NOT guess paths like src/main.py.

  A wrong edit is worse than stopping after honest exploration.



• Output format traps (FORBIDDEN):

  - Plain-text "I will now call ls..." without an actual tool call

  - Markdown ```json {"name":"ls"...}``` instead of a real tool call

  - Dumping a full rewritten file into chat without write_file/edit_file

  - Asking the user questions (there is no user — only tools + DONE)



• Budget:

  You have a limited number of tool rounds. Front-load discovery (ls → grep →

  read) then one precise edit. Avoid repetitive ls on the same path.



═══════════════════════════════════════════════════════════════════════════════

FIRST ACTION (hard requirement)

═══════════════════════════════════════════════════════════════════════════════

Your very first tool call MUST be:

  ls(path=".")

Then proceed with Steps B–F using only paths that exist.

"""


def build_agent_prompt(problem_statement: str, *, instance_id: str) -> str:
    issue = (problem_statement or "").strip()
    return (
        f"{AGENT_SYSTEM_PROMPT}\n"
        f"Instance ID: {instance_id}\n"
        f"{'=' * 79}\n"
        f"ISSUE / PROBLEM STATEMENT\n"
        f"{'=' * 79}\n"
        f"{issue}\n"
    )


def _row_to_task(ex: dict[str, Any], split: str) -> dict[str, Any]:
    patch = ex.get("patch") or ""
    test_patch = ex.get("test_patch") or ""
    instance_id = str(ex["instance_id"])
    problem_statement = str(ex.get("problem_statement") or "")
    return {
        "task_id": instance_id,
        "instance_id": instance_id,
        "split": split,
        "repo": str(ex["repo"]),
        "base_commit": str(ex["base_commit"]),
        "problem_statement": problem_statement,
        "prompt": build_agent_prompt(problem_statement, instance_id=instance_id),
        "workspace_paths": parse_paths_from_patch(patch),
        "gold_patch": patch,
        "test_patch": test_patch,
        "fail_to_pass": _as_str_list(ex.get("FAIL_TO_PASS")),
        "pass_to_pass": _as_str_list(ex.get("PASS_TO_PASS")),
        "hints_text": str(ex.get("hints_text") or ""),
        "version": str(ex.get("version") or ""),
        "difficulty": str(ex.get("difficulty") or ""),
        "image": str(ex.get("image") or ""),
        "eval_script": str(ex.get("eval_script") or ""),
        "eval_type": str(ex.get("eval_type") or "pass_and_fail"),
        "log_parser": str(ex.get("log_parser") or ""),
        "environment_setup_commit": str(ex.get("environment_setup_commit") or ""),
        "source_dataset": "SWE-bench/SWE-bench_Lite",
    }


def load_swe_bench_lite(

    *,

    splits: tuple[str, ...] = ("dev", "test"),

    max_per_split: int | None = None,

    token: str | None = None,

) -> dict[str, list[dict[str, Any]]]:
    from datasets import load_dataset

    kwargs: dict[str, Any] = {}
    if token:
        kwargs["token"] = token
    ds = load_dataset("SWE-bench/SWE-bench_Lite", **kwargs)
    out: dict[str, list[dict[str, Any]]] = {}
    for split in splits:
        if split not in ds:
            continue
        rows = []
        for i, ex in enumerate(ds[split]):
            if max_per_split is not None and i >= max_per_split:
                break
            rows.append(_row_to_task(dict(ex), split=split))
        out[split] = rows
    return out


def to_hf_records(tasks: list[dict[str, Any]]) -> list[dict[str, Any]]:
    records = []
    for t in tasks:
        records.append(
            {
                "task_id": t["task_id"],
                "instance_id": t["instance_id"],
                "split": t["split"],
                "repo": t["repo"],
                "base_commit": t["base_commit"],
                "problem_statement": t["problem_statement"],
                "prompt": t["prompt"],
                "workspace_paths_json": json.dumps(t["workspace_paths"]),
                "gold_patch": t["gold_patch"],
                "test_patch": t["test_patch"],
                "fail_to_pass_json": json.dumps(t["fail_to_pass"]),
                "pass_to_pass_json": json.dumps(t["pass_to_pass"]),
                "hints_text": t.get("hints_text") or "",
                "version": t.get("version") or "",
                "difficulty": t.get("difficulty") or "",
                "image": t.get("image") or "",
                "eval_script": t.get("eval_script") or "",
                "eval_type": t.get("eval_type") or "pass_and_fail",
                "log_parser": t.get("log_parser") or "",
                "environment_setup_commit": t.get("environment_setup_commit") or "",
                "source_dataset": t.get("source_dataset") or "SWE-bench/SWE-bench_Lite",
            }
        )
    return records


def from_hf_record(rec: dict[str, Any]) -> dict[str, Any]:
    # Backward compat: older sparse dataset rows may lack image/eval_script —
    # callers should reload from SWE-bench/SWE-bench_Lite in that case.
    return {
        "task_id": rec["task_id"],
        "instance_id": rec["instance_id"],
        "split": rec.get("split") or "train",
        "repo": rec["repo"],
        "base_commit": rec["base_commit"],
        "problem_statement": rec["problem_statement"],
        "prompt": rec["prompt"],
        "workspace_paths": json.loads(rec.get("workspace_paths_json") or "[]"),
        "gold_patch": rec.get("gold_patch") or "",
        "test_patch": rec.get("test_patch") or "",
        "fail_to_pass": json.loads(rec.get("fail_to_pass_json") or "[]"),
        "pass_to_pass": json.loads(rec.get("pass_to_pass_json") or "[]"),
        "hints_text": rec.get("hints_text") or "",
        "version": rec.get("version") or "",
        "difficulty": rec.get("difficulty") or "",
        "image": rec.get("image") or "",
        "eval_script": rec.get("eval_script") or "",
        "eval_type": rec.get("eval_type") or "pass_and_fail",
        "log_parser": rec.get("log_parser") or "",
        "environment_setup_commit": rec.get("environment_setup_commit") or "",
        "source_dataset": rec.get("source_dataset") or "SWE-bench/SWE-bench_Lite",
    }


def save_jsonl(tasks: list[dict[str, Any]], path) -> None:
    from pathlib import Path

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for t in to_hf_records(tasks):
            f.write(json.dumps(t) + "\n")