src/spreadsheetbench_env/interface.py

"""Framework-neutral SpreadsheetBench actor-facing interface helpers."""

from __future__ import annotations

from pathlib import Path
import re

import openpyxl


# SkillOpt SpreadsheetBench mode=multi: up to 30 code-repair turns per task, a 600 s wall-clock
# budget per task (LLM latency included), and each execution feedback truncated to 3000 chars.
STEP_LIMIT = 30
TASK_TIME_BUDGET_S = 600.0
# SkillOpt stops issuing turns once fewer than this many seconds of the budget remain.
DEADLINE_MARGIN_S = 10.0
_ERROR_FEEDBACK_CHARS = 3000
_REPORT_LINE = re.compile(r"(\S+!?\w+):\s*got=(.+?),\s*expected=.+?\s*(PASS|FAIL)$")


def environment_description() -> str:
    """The SkillOpt SpreadsheetBench codegen system prompt, adapted only in its last sentence to
    the native-tool transport (the script travels in the tool call's ``code`` field, not a fence)."""

    return (
        "You are an expert Python programmer specializing in spreadsheet manipulation. You will "
        "be given a user instruction together with a preview of an input .xlsx file. Your job is "
        "to write a single self-contained Python script that reads the input file at the path "
        "stored in the variable INPUT_PATH, performs the requested manipulation, and saves the "
        "result to OUTPUT_PATH. Use only the standard library, openpyxl, and pandas. Do not print "
        "anything. Do not use input(). Do not hardcode file paths. Return ONLY the Python code in "
        "the `code` field of the submit_script tool call."
    )


def preview_workbook(
    path: str | Path,
    *,
    max_rows: int = 5,
    max_columns: int = 20,
) -> str:
    """The SkillOpt workbook preview: first 5 rows, 20 columns per sheet, coordinates and values."""

    workbook = openpyxl.load_workbook(path, data_only=False)
    chunks: list[str] = []
    try:
        for sheet_name in workbook.sheetnames:
            sheet = workbook[sheet_name]
            chunks.append(
                f"## Sheet: {sheet_name}  "
                f"(dim={sheet.dimensions}, max_row={sheet.max_row}, max_col={sheet.max_column})"
            )
            for row in sheet.iter_rows(
                min_row=1,
                max_row=min(sheet.max_row, max_rows),
                max_col=min(sheet.max_column, max_columns),
                values_only=False,
            ):
                cells: list[str] = []
                for cell in row:
                    value = "" if cell.value is None else str(cell.value)
                    if len(value) > 40:
                        value = value[:37] + "..."
                    cells.append(f"{cell.coordinate}={value}")
                chunks.append(" | ".join(cells))
            if sheet.max_row > max_rows:
                chunks.append(f"... ({sheet.max_row - max_rows} more rows)")
            chunks.append("")
    finally:
        workbook.close()
    return "\n".join(chunks)


def build_task_prompt(
    instruction: str,
    instruction_type: str,
    answer_position: str,
    preview: str,
) -> str:
    """The SkillOpt codegen user prompt, verbatim except the transport sentence."""

    return (
        f"# Instruction\n{instruction}\n"
        f"\nInstruction type: {instruction_type}"
        f"\nExpected answer position: {answer_position}\n\n"
        f"# Input spreadsheet preview\n{preview}\n\n"
        "# Task\n"
        "Write a Python script that reads the workbook from the variable `INPUT_PATH`, "
        "applies the instruction, and writes the modified workbook to `OUTPUT_PATH`. "
        "Preserve all other cells unchanged. "
        "The preview may be truncated \u2014 do not hardcode row counts or assume the data ends at "
        "the last previewed row; iterate over all actual rows in the workbook instead. "
        "Return only the Python code in the `code` field."
    )


def extract_script(value: str) -> str:
    """Extract the first fenced code block, preserving legacy malformed-fence behavior."""

    text = str(value)
    if "```" not in text:
        return text.strip()
    start = text.find("```")
    newline = text.find("\n", start)
    end = text.find("```", newline + 1)
    if newline == -1 or end == -1:
        code = text.strip()
        if code.startswith("```"):
            first_newline = code.find("\n")
            code = code[first_newline + 1 :] if first_newline >= 0 else ""
        return code.strip()
    return text[newline + 1 : end].strip()


def execution_error_feedback(error: str) -> str:
    """SkillOpt's execution-failure feedback turn, verbatim except the transport sentence."""

    return (
        "The code raised an error during execution:\n\n"
        f"```\n{error[:_ERROR_FEEDBACK_CHARS]}\n```\n\n"
        "Please fix the code and return a complete corrected Python script in the `code` field "
        "of the submit_script tool call."
    )


def eval_feedback(verification: str) -> str:
    """SkillOpt's training-only wrong-output feedback: the agent's own cell values marked WRONG,
    expected values hidden, correct cells counted. Lines that are not per-cell comparisons
    (missing sheets, truncation markers) are dropped, as in the original."""

    wrong_lines: list[str] = []
    n_correct = 0
    for raw_line in verification.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        match = _REPORT_LINE.match(line)
        if match is None:
            continue
        cell, got_value, mark = match.groups()
        if mark == "FAIL":
            wrong_lines.append(f"  {cell}: your output = {got_value} (WRONG)")
        else:
            n_correct += 1
    lines = [
        "Your code executed successfully but produced incorrect results.",
        "The following cells have wrong values:",
    ]
    lines.extend(wrong_lines)
    if n_correct:
        lines.append(f"  ({n_correct} other cells are correct.)")
    lines.append(
        "\nPlease analyze the spreadsheet data more carefully and fix the code. "
        "Return a complete corrected Python script in the `code` field of the submit_script "
        "tool call."
    )
    return "\n".join(lines)