scripts/build_sft_messages.py

#!/usr/bin/env python3
"""
build_sft_messages.py

Build Hugging Face SFT conversational rows from the canonical Mythos-Coder dataset.

Response style is chosen by task intent:
- Code-generation prompts -> "Here is the complete code:" + full solution
- bug_fix / terminal_debug / ui_repair / migration / refactor -> Diagnosis format
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

from sft_pipeline_utils import (
    asks_for_code,
    has_code,
    is_weak_solution,
    uses_diagnosis_format,
)
from dataset_quality_utils import (
    SCORE_KEEP,
    asks_single_file_html,
    has_external_assets,
    has_inline_style_script,
    score_row,
)

SYSTEM_PROMPT = (
    "You are Mythos-Coder, a coding agent that inspects the task, makes a "
    "concise plan, edits safely, verifies results, and explains fixes clearly. "
    "When the user asks for code, output complete runnable code."
)

REQUIRED_FIELDS = (
    "user_prompt",
    "failure_log",
    "investigation_steps",
    "plan",
    "solution",
    "verification",
    "lesson",
)


def project_root() -> Path:
    return Path(__file__).resolve().parent.parent


def field_text(row: dict, field: str, default: str = "") -> str:
    return str(row.get(field, "")).strip() or default


def short_diagnosis(failure_log: str, user_prompt: str) -> str:
    """Diagnosis must be short and must not copy the full user_prompt."""
    fl = str(failure_log or "").strip()
    user = str(user_prompt or "").strip()

    if fl.startswith("Initial problem:"):
        if "Common pitfall:" in fl:
            return 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 rest.split("Common pitfall:", 1)[1].strip()
        if rest:
            return rest[:240]
        return "Reproduce the reported behavior in the affected file before editing."

    if user and user.lower() in fl.lower():
        return "Reproduce the reported behavior in the affected file before editing."

    return fl[:240] if fl else "Reproduce the reported behavior in the affected file before editing."


def format_verification_checklist(verification: str) -> str:
    text = str(verification or "").strip()
    if not text:
        return "- Open the file or app and confirm the expected behavior works."
    numbered = re.split(r"\d+\)\s*", text)
    items = [part.strip(" ;.") for part in numbered if part.strip()]
    if len(items) >= 2:
        return "\n".join(f"- {item}" for item in items[:5])
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    if len(lines) >= 2:
        return "\n".join(f"- {ln.lstrip('- ')}" for ln in lines[:5])
    return f"- {text}"


def is_code_generation_row(row: dict) -> bool:
    solution = str(row.get("solution", ""))
    if not has_code(solution):
        return False
    if asks_for_code(row.get("user_prompt", "")):
        return True
    if "```" in solution or "<!DOCTYPE" in solution or "<html" in solution.lower():
        return True
    return False


def build_code_generation_content(row: dict) -> str:
    solution = field_text(row, "solution")
    verification = format_verification_checklist(field_text(row, "verification"))
    return (
        "Here is the complete code:\n"
        f"{solution}\n\n"
        f"Verification:\n{verification}"
    )


def build_diagnosis_content(row: dict) -> str:
    user_prompt = field_text(row, "user_prompt")
    diagnosis = short_diagnosis(field_text(row, "failure_log"), user_prompt)
    plan = field_text(row, "plan", "Outline the smallest safe change before editing.")
    implementation = field_text(row, "solution")
    verification = field_text(row, "verification", "Re-test the affected workflow and confirm the fix holds.")
    lesson = field_text(row, "lesson", "Prefer small verified edits over broad rewrites.")

    return (
        f"Diagnosis:\n{diagnosis}\n\n"
        f"Plan:\n{plan}\n\n"
        f"Implementation:\n{implementation}\n\n"
        f"Verification:\n{verification}\n\n"
        f"Lesson:\n{lesson}"
    )


def build_assistant_content(row: dict) -> str:
    if is_code_generation_row(row):
        return build_code_generation_content(row)
    if uses_diagnosis_format(row):
        return build_diagnosis_content(row)
    # feature_build / website_vibe without explicit code ask — still prefer code if present
    if has_code(row.get("solution", "")):
        return build_code_generation_content(row)
    return build_diagnosis_content(row)


def row_quality_ok(row: dict) -> bool:
    """Skip weak canonical rows before SFT conversion."""
    solution = field_text(row, "solution")
    user = field_text(row, "user_prompt")

    score, issues = score_row(row)
    if score < SCORE_KEEP:
        return False

    if is_weak_solution(solution):
        return False
    if asks_for_code(user) and not has_code(solution):
        return False

    if asks_single_file_html(user):
        if has_external_assets(solution):
            return False
        if "inline" in user.lower() or "inside the file" in user.lower() or "single file" in user.lower():
            if not has_inline_style_script(solution):
                return False

    failure_log = field_text(row, "failure_log")
    if failure_log.startswith("Initial problem:") and user in failure_log and not has_code(solution):
        return False

    if "solution_repeats_prompt" in issues or "failure_log_repeats_prompt" in issues:
        if not (asks_for_code(user) and has_code(solution)):
            return False

    return True


def is_valid_row(row: dict) -> bool:
    if not isinstance(row, dict):
        return False
    user_prompt = str(row.get("user_prompt", "")).strip()
    if not user_prompt:
        return False
    for field in REQUIRED_FIELDS:
        if field not in row:
            return False
    if not isinstance(row.get("investigation_steps"), list):
        return False
    return True


def load_jsonl_file(path: Path) -> tuple[list[dict], int]:
    rows: list[dict] = []
    skipped = 0
    with open(path, "r", encoding="utf-8") as handle:
        for line_num, line in enumerate(handle, 1):
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                skipped += 1
                print(f"Skip invalid JSON: {path.name}:{line_num} ({exc})", file=sys.stderr)
                continue
            if not is_valid_row(row):
                skipped += 1
                print(f"Skip invalid row: {path.name}:{line_num}", file=sys.stderr)
                continue
            rows.append(row)
    return rows, skipped


def load_train_rows(train_path: Path) -> tuple[list[dict], int, int]:
    rows, skipped = load_jsonl_file(train_path)
    return rows, skipped, len(rows) + skipped


def load_converted_rows(converted_dir: Path) -> tuple[list[dict], int, int]:
    rows: list[dict] = []
    skipped = 0
    input_total = 0

    jsonl_files = sorted(converted_dir.glob("*.jsonl"))
    if not jsonl_files:
        print(f"Warning: no JSONL files found in {converted_dir}", file=sys.stderr)

    for path in jsonl_files:
        file_rows, file_skipped = load_jsonl_file(path)
        rows.extend(file_rows)
        skipped += file_skipped
        input_total += len(file_rows) + file_skipped

    return rows, skipped, input_total


def dedupe_rows(rows: list[dict]) -> tuple[list[dict], int]:
    seen: set[str] = set()
    unique: list[dict] = []
    duplicates = 0

    for row in rows:
        key = str(row["user_prompt"]).strip()
        if key in seen:
            duplicates += 1
            continue
        seen.add(key)
        unique.append(row)

    return unique, duplicates


def to_sft_message(row: dict) -> dict:
    user_content = str(row["user_prompt"]).strip()
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": build_assistant_content(row)},
        ],
        "_source_id": row.get("id", ""),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build Mythos SFT messages JSONL.")
    parser.add_argument(
        "--train-file",
        default=None,
        help="Canonical train JSONL (default: data/train/mythos_coder_clean_canonical.jsonl if exists, else datasets/mythos_coder_train.jsonl)",
    )
    parser.add_argument(
        "--converted-dir",
        default=None,
        help="Optional converted JSONL directory instead of train file",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Output path (default: data/train/mythos_sft_messages.jsonl)",
    )
    parser.add_argument(
        "--include-weak",
        action="store_true",
        help="Include low-quality rows (default: skip weak/planning-only rows)",
    )
    parser.add_argument(
        "--extra-jsonl",
        action="append",
        default=[],
        help="Additional canonical JSONL files to merge (e.g. code_output correction batch)",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    root = project_root()
    output_path = Path(args.output) if args.output else root / "data" / "train" / "mythos_sft_messages_clean.jsonl"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if args.converted_dir:
        source_label = args.converted_dir
        rows, skipped, input_total = load_converted_rows(Path(args.converted_dir))
    else:
        clean_path = root / "data" / "train" / "mythos_coder_clean_canonical.jsonl"
        if args.train_file:
            train_path = Path(args.train_file)
            if not train_path.is_absolute():
                train_path = root / train_path
        elif clean_path.exists():
            train_path = clean_path
        else:
            train_path = root / "datasets" / "mythos_coder_train.jsonl"
        if not train_path.exists():
            print(f"Error: train file not found: {train_path}", file=sys.stderr)
            return 1
        source_label = str(train_path)
        rows, skipped, input_total = load_train_rows(train_path)

    unique_rows, duplicates = dedupe_rows(rows)

    for extra in args.extra_jsonl:
        extra_path = Path(extra)
        if not extra_path.is_absolute():
            extra_path = root / extra_path
        if not extra_path.exists():
            print(f"Warning: extra JSONL not found: {extra_path}", file=sys.stderr)
            continue
        extra_rows, extra_skipped = load_jsonl_file(extra_path)
        skipped += extra_skipped
        input_total += len(extra_rows) + extra_skipped
        rows_before = len(unique_rows)
        unique_rows, extra_dupes = dedupe_rows(unique_rows + extra_rows)
        duplicates += extra_dupes
        print(f"Merged {len(extra_rows)} rows from {extra_path.name} (+{len(unique_rows)-rows_before} unique)")

    quality_skipped = 0
    if not args.include_weak:
        kept = []
        for row in unique_rows:
            if row_quality_ok(row):
                kept.append(row)
            else:
                quality_skipped += 1
        unique_rows = kept

    sft_rows = []
    for row in unique_rows:
        msg = to_sft_message(row)
        msg.pop("_source_id", None)
        sft_rows.append(msg)

    assistant_lengths = [len(row["messages"][2]["content"]) for row in sft_rows]
    avg_len = sum(assistant_lengths) / len(assistant_lengths) if assistant_lengths else 0
    with_code = sum(1 for row in sft_rows if has_code(row["messages"][2]["content"]))

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

    print(f"Source:              {source_label}")
    print(f"Input rows:          {input_total}")
    print(f"Skipped invalid:     {skipped}")
    print(f"Duplicate rows:      {duplicates}")
    print(f"Quality skipped:     {quality_skipped}")
    print(f"Output rows:         {len(sft_rows)}")
    print(f"Assistant with code: {with_code}")
    print(f"Avg assistant chars: {avg_len:.0f}")
    print(f"Wrote:               {output_path}")
    return 0


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