scripts/repair_weak_rows.py

#!/usr/bin/env python3
"""Mechanically repair weak canonical rows; reject unsafe ones."""

from __future__ import annotations

import json
import re
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from dataset_quality_utils import (
    SCORE_REPAIR,
    asks_for_code,
    asks_single_file_html,
    asks_three_feature_cards,
    has_code,
    nav_selector_issues,
    score_row,
)

ROOT = Path(__file__).resolve().parent.parent
WEAK_PATH = ROOT / "data" / "audit" / "weak_rows.jsonl"
REPAIRED_DIR = ROOT / "data" / "repaired"
REJECTED_DIR = ROOT / "data" / "rejected"
REPAIRED_PATH = REPAIRED_DIR / "repaired_rows.jsonl"
REJECTED_PATH = REJECTED_DIR / "rejected_rows.jsonl"
REPORT_PATH = ROOT / "data" / "audit" / "repair_report.md"

BANNED_RE = re.compile(
    r"Plus \d+ additional targeted edits|Plus additional targeted edits|"
    r"Make targeted edits|Read the HTML structure|Fix the issue|Verify it works",
    re.I,
)

VERIFICATION_BY_TYPE = {
    "bug_fix": "Open the affected file, reproduce the bug, apply the fix, and confirm the error or broken behavior is gone.",
    "terminal_debug": "Run the command or browser check again and confirm output matches the expected result with no new errors.",
    "ui_repair": "Resize the viewport and confirm the visual issue is fixed on both mobile and desktop.",
    "refactor": "Run existing tests or smoke the feature path and confirm behavior is unchanged after the refactor.",
    "feature_build": "Load the feature in the app and confirm the new behavior works end-to-end.",
    "website_vibe": "Open the page in a browser and confirm layout, links, and interactions match the request.",
    "test_writing": "Run the test suite and confirm the new test fails before the fix and passes after.",
    "migration": "Run migration on a copy of the environment and verify schema/data changes match the plan.",
}


def trim_failure_log(failure_log: str, user_prompt: str) -> str:
    fl = str(failure_log or "").strip()
    user = str(user_prompt or "").strip()
    if fl.startswith("Initial problem:"):
        if "Common pitfall:" in fl:
            return f"Symptom: {fl.split('Common pitfall:', 1)[1].strip()}"
        rest = fl[len("Initial problem:") :].strip()
        if user and rest.startswith(user):
            rest = rest[len(user) :].strip(" .")
        if rest.startswith("Common pitfall:"):
            return f"Symptom: {rest.split('Common pitfall:', 1)[1].strip()}"
        return rest[:200] if rest else "Reproduce the issue before editing."
    if user and user in fl:
        return "Reproduce the issue in the named file or handler before editing."
    return fl


def concrete_verification(task_type: str, verification: str) -> str:
    v = str(verification or "").strip()
    if v and not v.lower().startswith("verify it works") and len(v) > 40:
        cleaned = BANNED_RE.sub("", v).strip()
        if cleaned:
            return cleaned
    return VERIFICATION_BY_TYPE.get(task_type, "Open the project and confirm the requested behavior works.")


def sharpen_lesson(lesson: str, task_type: str) -> str:
    l = str(lesson or "").strip()
    if len(l) > 60 and "avoid:" not in l.lower():
        return l
    defaults = {
        "bug_fix": "Match the selector or state variable named in the task; do not toggle classes on the wrong element.",
        "ui_repair": "Verify responsive behavior at the exact breakpoint mentioned in the prompt.",
    }
    return defaults.get(task_type, l or "Prefer the smallest verified change over a broad rewrite.")


def fix_selector_in_solution(solution: str, user_prompt: str) -> str | None:
    """Only fix when prompt explicitly states correct target."""
    p = user_prompt.lower()
    if ".nav-links.open" not in p and "toggle .open on .nav-links" not in p:
        return None
    if "do not toggle .open on .hamburger" not in p and "hamburger button" not in p:
        return None
    sol = solution
    if re.search(r"\.hamburger[^;\n]*classList\.toggle\(['\"]open", sol, re.I):
        fixed = re.sub(
            r"document\.querySelector\(['\"]\.hamburger['\"]\)\.classList\.toggle\(['\"]open['\"]\)",
            "document.querySelector('.nav-links').classList.toggle('open')",
            sol,
            flags=re.I,
        )
        if fixed != sol:
            return fixed
    return None


def rejection_reason(issues: list[str], row: dict) -> str:
    user = row.get("user_prompt", "")
    solution = row.get("solution", "")
    if "code_request_no_code" in issues or (asks_for_code(user) and not has_code(solution)):
        return "needs_code_output_rewrite"
    if "single_file_external_assets" in issues or "single_file_missing_inline_style_script" in issues:
        return "single_file_html_violation"
    if "missing_three_feature_cards" in issues:
        return "missing_required_feature"
    if any(i.startswith("toggles_open") or i.startswith("invents_") for i in issues):
        return "selector_mismatch"
    if "solution_repeats_prompt" in issues or "failure_log_repeats_prompt" in issues:
        return "prompt_repetition"
    if "generic_bullets_only" in issues:
        return "generic_solution"
    return "unsafe_to_auto_repair"


def can_repair(entry: dict) -> bool:
    issues = entry["issues"]
    row = entry["row"]
    user = row.get("user_prompt", "")
    solution = row.get("solution", "")

    hard_reject = (
        asks_for_code(user) and not has_code(solution),
        asks_single_file_html(user) and ("single_file_external_assets" in issues or "single_file_missing_inline_style_script" in issues),
        asks_three_feature_cards(user) and "missing_three_feature_cards" in issues,
        any(i.startswith("invents_") for i in issues),
        "generic_bullets_only" in issues and not has_code(solution),
        entry["score"] < SCORE_REPAIR,
    )
    if any(hard_reject):
        return False
    safe = {
        "failure_log_repeats_prompt", "failure_log_copies_instruction",
        "vague_verification", "generic_lesson", "banned_phrase",
    }
    return any(any(i.startswith(s) for s in safe) for i in issues) or entry["score"] >= SCORE_REPAIR


def repair_row(entry: dict) -> tuple[dict | None, str | None]:
    row = dict(entry["row"])
    issues = entry["issues"]
    user = row.get("user_prompt", "")

    if not can_repair(entry):
        return None, rejection_reason(issues, row)

    row["failure_log"] = trim_failure_log(row.get("failure_log", ""), user)
    row["verification"] = concrete_verification(row.get("task_type", ""), row.get("verification", ""))
    row["lesson"] = sharpen_lesson(row.get("lesson", ""), row.get("task_type", ""))

    for field in ("solution", "plan", "verification", "lesson", "failure_log"):
        if field in row and isinstance(row[field], str):
            row[field] = BANNED_RE.sub("", row[field]).strip()

    selector_fix = fix_selector_in_solution(row.get("solution", ""), user)
    if selector_fix:
        row["solution"] = selector_fix

    new_score, new_issues = score_row(row)
    row["quality_score"] = min(10, max(1, new_score // 10))
    if new_score < SCORE_REPAIR:
        return None, rejection_reason(new_issues, row)
    return row, None


def main() -> int:
    REPAIRED_DIR.mkdir(parents=True, exist_ok=True)
    REJECTED_DIR.mkdir(parents=True, exist_ok=True)

    if not WEAK_PATH.exists():
        print(f"Missing {WEAK_PATH} — run audit_dataset_quality.py first", file=sys.stderr)
        return 1

    repaired: list[dict] = []
    rejected: list[dict] = []
    reasons: Counter = Counter()

    with WEAK_PATH.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            entry = json.loads(line)
            fixed, reason = repair_row(entry)
            if fixed:
                repaired.append(fixed)
            else:
                reasons[reason or "unsafe_to_auto_repair"] += 1
                rejected.append({
                    "rejection_reason": reason,
                    "source_file": entry.get("source_file"),
                    "id": entry.get("id"),
                    "score": entry.get("score"),
                    "issues": entry.get("issues"),
                    "row": entry.get("row"),
                })

    with REPAIRED_PATH.open("w", encoding="utf-8") as f:
        for row in repaired:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    with REJECTED_PATH.open("w", encoding="utf-8") as f:
        for item in rejected:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    lines = [
        "# Repair Report",
        "",
        f"Weak rows input: {len(repaired) + len(rejected)}",
        f"Repaired: **{len(repaired)}** -> `{REPAIRED_PATH}`",
        f"Rejected: **{len(rejected)}** -> `{REJECTED_PATH}`",
        "",
        "## Rejection reasons",
        "",
    ]
    for r, c in reasons.most_common():
        lines.append(f"- `{r}`: {c}")

    REPORT_PATH.write_text("\n".join(lines), encoding="utf-8")
    print(f"Repaired: {len(repaired)}, Rejected: {len(rejected)}")
    print(f"Wrote {REPORT_PATH}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())