scripts/merge_v8_training_corpus.py

"""Merge Figment v8 delta rows with the v7 corpus and prepare Modal splits."""

from __future__ import annotations

import argparse
from collections import Counter
from datetime import UTC
from datetime import datetime
import hashlib
import json
from pathlib import Path
import subprocess
import sys
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from figment.focused_repair import build_focused_repair_prompts  # noqa: E402
from figment.eval_metrics import bucket_expected_observation_cues  # noqa: E402
from figment.observation_targets import required_observation_targets  # noqa: E402
from figment.prompt_builder import build_prompt  # noqa: E402
from figment.retrieval import known_card_ids  # noqa: E402
from figment.retrieval import load_protocol_cards  # noqa: E402
from figment.retrieval import query_from_intake  # noqa: E402
from figment.retrieval import search_protocol_cards  # noqa: E402
from figment.rules import run_red_flag_checks  # noqa: E402
from figment.trace import stable_hash  # noqa: E402
from figment.validators import urgency_floor_from_rules  # noqa: E402
from figment.validators import validate_navigator_output  # noqa: E402
from scripts.augment_finetune_repair_rows import _corrupt_output  # noqa: E402
from scripts.augment_finetune_repair_rows import _extra_failures_for_scope  # noqa: E402
from scripts.generate_finetune_data import _required_retrieved_ids  # noqa: E402
from scripts.generate_finetune_data import _expected_candidate_cards  # noqa: E402
from scripts.generate_finetune_data import _expected_missing_observations  # noqa: E402
from scripts.generate_finetune_data import _expected_source_cards  # noqa: E402
from scripts.generate_finetune_data import ensure_retrieved_cards  # noqa: E402
from scripts.generate_finetune_data import uses_v7_source_card_policy  # noqa: E402


DEFAULT_BASE = Path("data/finetune/figment_sft_v7.jsonl")
DEFAULT_BASE_CASE_SPECS = Path("data/finetune/figment_sft_v7_case_specs.jsonl")
DEFAULT_DELTA = Path("data/finetune/figment_sft_v8_delta.jsonl")
DEFAULT_DELTA_CASE_SPECS = Path("data/finetune/figment_sft_v8_delta_case_specs.jsonl")
DEFAULT_OUTPUT = Path("data/finetune/figment_sft_v8.jsonl")
DEFAULT_CASE_SPECS = Path("data/finetune/figment_sft_v8_case_specs.jsonl")
DEFAULT_MANIFEST = Path("data/finetune/figment_sft_v8_manifest.json")
DEFAULT_MODAL_DIR = Path("data/finetune/modal/figment_sft_v8")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base", type=Path, default=DEFAULT_BASE)
    parser.add_argument("--base-case-specs", type=Path, default=DEFAULT_BASE_CASE_SPECS)
    parser.add_argument("--delta", type=Path, default=DEFAULT_DELTA)
    parser.add_argument("--delta-case-specs", type=Path, default=DEFAULT_DELTA_CASE_SPECS)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--case-specs", type=Path, default=DEFAULT_CASE_SPECS)
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    parser.add_argument("--modal-output-dir", type=Path, default=DEFAULT_MODAL_DIR)
    parser.add_argument("--dataset-version", default="figment_sft_v8")
    parser.add_argument("--validation-fraction", type=float, default=0.1)
    parser.add_argument("--seed", default="figment-modal-sft-v8")
    parser.add_argument("--min-validation-group-size", type=int, default=5)
    parser.add_argument("--skip-verify", action="store_true")
    parser.add_argument("--skip-modal-prep", action="store_true")
    args = parser.parse_args(argv)

    summary = merge_v8_corpus(
        base_path=args.base,
        base_case_specs_path=args.base_case_specs,
        delta_path=args.delta,
        delta_case_specs_path=args.delta_case_specs,
        output_path=args.output,
        case_specs_path=args.case_specs,
        dataset_version=args.dataset_version,
    )

    verify_summary = None
    if not args.skip_verify:
        verify_summary = _run_json_command(
            [
                sys.executable,
                "scripts/verify_finetune_harness_alignment.py",
                "--dataset",
                str(args.output),
                "--case-specs",
                str(args.case_specs),
            ]
        )
        if verify_summary.get("passed") is not True:
            raise SystemExit(f"harness verification failed: {json.dumps(verify_summary, sort_keys=True)}")

    modal_summary = None
    if not args.skip_modal_prep:
        modal_summary = _run_json_command(
            [
                sys.executable,
                "scripts/prepare_modal_finetune_dataset.py",
                "--dataset",
                str(args.output),
                "--dataset-version",
                args.dataset_version,
                "--output-dir",
                str(args.modal_output_dir),
                "--validation-fraction",
                str(args.validation_fraction),
                "--seed",
                args.seed,
                "--min-validation-group-size",
                str(args.min_validation_group_size),
            ]
        )

    summary["verify"] = verify_summary
    summary["modal"] = modal_summary
    args.manifest.parent.mkdir(parents=True, exist_ok=True)
    args.manifest.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


def merge_v8_corpus(
    *,
    base_path: Path,
    base_case_specs_path: Path,
    delta_path: Path,
    delta_case_specs_path: Path,
    output_path: Path,
    case_specs_path: Path,
    dataset_version: str,
) -> dict[str, Any]:
    base_rows = _read_jsonl(base_path)
    delta_rows = _read_jsonl(delta_path)
    rows = _dedupe_rows(base_rows + delta_rows)
    specs = _case_specs_for_rows(
        rows,
        source_case_specs=[base_case_specs_path, delta_case_specs_path],
    )
    spec_refresh_summary = _refresh_case_specs_for_current_harness(specs)
    refresh_summary = _refresh_prompts_for_current_harness(rows, specs)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    case_specs_path.parent.mkdir(parents=True, exist_ok=True)
    _write_jsonl(output_path, rows)
    _write_jsonl(case_specs_path, specs)

    return {
        "dataset_version": dataset_version,
        "merged_at": datetime.now(UTC).isoformat(),
        "row_count": len(rows),
        "base_rows": len(base_rows),
        "delta_rows": len(delta_rows),
        "case_spec_count": len(specs),
        "output_path": str(output_path),
        "case_specs_path": str(case_specs_path),
        "output_sha256": _sha256_path(output_path),
        "case_specs_sha256": _sha256_path(case_specs_path),
        "task_type_counts": dict(sorted(Counter(_task_type(row) for row in rows).items())),
        "category_counts": dict(sorted(Counter(str(row.get("category") or "unknown") for row in rows).items())),
        "case_spec_refresh": spec_refresh_summary,
        "prompt_refresh": refresh_summary,
    }


def _refresh_case_specs_for_current_harness(specs: list[dict[str, Any]]) -> dict[str, Any]:
    cards_by_id = {str(card["card_id"]): card for card in load_protocol_cards()}
    changed: Counter[str] = Counter()
    changed_specs = 0
    for spec in specs:
        spec_changed = False
        harness = _harness_for_spec(spec, {"version": spec.get("dataset_version")}, cards_by_id)
        synthetic_spec = _synthetic_spec_for_case_spec(spec)
        expected_source = _expected_source_cards(
            synthetic_spec,
            harness["rule_results"],
            harness["retrieved_ids"],
        )
        expected_candidates = _expected_candidate_cards(synthetic_spec, harness["rule_results"])
        expected_missing = _expected_missing_observations(
            synthetic_spec,
            [card_id for card_id in expected_source if card_id in harness["retrieved_ids"]],
            cards_by_id,
        )
        updates = {
            "expected_red_flag_rule_ids": [str(rule["rule_id"]) for rule in harness["rule_results"]],
            "expected_min_protocol_urgency": harness["urgency_floor"],
            "expected_source_card_ids": expected_source,
            "expected_candidate_pathway_card_ids": expected_candidates,
            "expected_missing_observations": expected_missing,
            "retrieved_card_ids": harness["retrieved_ids"],
        }
        cue_buckets = bucket_expected_observation_cues(expected_missing)
        updates.update(
            {
                "expected_model_observation_cues": cue_buckets["model"],
                "expected_handoff_cues": cue_buckets["handoff"],
                "expected_harness_evidence_cues": cue_buckets["harness"],
            }
        )
        for key, value in updates.items():
            if spec.get(key) != value:
                changed[key] += 1
                spec[key] = value
                spec_changed = True
        if spec_changed:
            changed_specs += 1
    return {
        "policy_version": 1,
        "updated_field_counts": dict(sorted(changed.items())),
        "updated_specs": changed_specs,
    }


def _synthetic_spec_for_case_spec(spec: dict[str, Any]) -> Any:
    return type(
        "SyntheticSpecForV8MergeCaseSpecRefresh",
        (),
        {
            "target_protocol_card_id": str(spec.get("target_protocol_card_id") or ""),
            "dataset_version": str(spec.get("dataset_version") or ""),
            "failure_class": str(spec.get("failure_class") or ""),
        },
    )()


def _refresh_prompts_for_current_harness(rows: list[dict[str, Any]], specs: list[dict[str, Any]]) -> dict[str, Any]:
    specs_by_id = {str(spec.get("case_id") or ""): spec for spec in specs}
    rows_by_id = {str(row.get("case_id") or row.get("uuid") or ""): row for row in rows}
    cards_by_id = {str(card["card_id"]): card for card in load_protocol_cards()}
    refreshed: Counter[str] = Counter()
    skipped: Counter[str] = Counter()

    for row in rows:
        row_id = str(row.get("case_id") or row.get("uuid") or "")
        metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
        base_case_id = str(metadata.get("base_case_id") or row_id)
        spec = specs_by_id.get(base_case_id)
        messages = row.get("messages")
        if spec is None or not isinstance(messages, list) or len(messages) < 2:
            skipped["missing_spec_or_messages"] += 1
            continue

        harness = _harness_for_spec(spec, row, cards_by_id)
        task_type = str(metadata.get("task_type") or "navigator_full")
        if task_type == "focused_repair":
            prompt = _focused_repair_prompt_for_row(
                row=row,
                rows_by_id=rows_by_id,
                spec=spec,
                harness=harness,
            )
            if prompt is None:
                skipped["focused_repair_prompt_not_refreshed"] += 1
                continue
            prompt_text = prompt
        else:
            prompt_text = harness["prompt"]

        if not isinstance(row.get("metadata"), dict):
            row["metadata"] = metadata
        if not isinstance(messages[0], dict):
            skipped["bad_user_message_shape"] += 1
            continue
        messages[0]["content"] = prompt_text
        metadata["prompt_hash"] = stable_hash(prompt_text)
        metadata["prompt_template_hash"] = harness["prompt_template_hash"]
        refreshed[task_type] += 1

    return {
        "refreshed_rows": sum(refreshed.values()),
        "refreshed_by_task_type": dict(sorted(refreshed.items())),
        "skipped": dict(sorted(skipped.items())),
        "policy_version": 1,
    }


def _harness_for_spec(spec: dict[str, Any], row: dict[str, Any], cards_by_id: dict[str, dict[str, Any]]) -> dict[str, Any]:
    intake = spec["structured_intake"]
    rule_results = [rule.to_dict() for rule in run_red_flag_checks(intake)]
    floor = urgency_floor_from_rules(rule_results)
    retrieved = search_protocol_cards(query_from_intake(intake), limit=6)
    spec_dataset_version = str(spec.get("dataset_version") or row.get("version") or "")
    if uses_v7_source_card_policy(spec_dataset_version):
        synthetic_spec = type(
            "SyntheticSpecForV8MergeRefresh",
            (),
            {
                "target_protocol_card_id": str(spec.get("target_protocol_card_id") or ""),
                "dataset_version": spec_dataset_version,
            },
        )()
        retrieved = ensure_retrieved_cards(
            retrieved,
            required_ids=_required_retrieved_ids(synthetic_spec, rule_results),
            cards_by_id=cards_by_id,
            limit=6,
        )
    prompt, prompt_template_hash = build_prompt(intake, retrieved, rule_results, floor)
    return {
        "prompt": prompt,
        "prompt_template_hash": prompt_template_hash,
        "rule_results": rule_results,
        "urgency_floor": floor,
        "retrieved": retrieved,
        "retrieved_ids": [str(item.get("card_id", "")) for item in retrieved if item.get("card_id")],
    }


def _focused_repair_prompt_for_row(
    *,
    row: dict[str, Any],
    rows_by_id: dict[str, dict[str, Any]],
    spec: dict[str, Any],
    harness: dict[str, Any],
) -> str | None:
    metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
    base_case_id = str(metadata.get("base_case_id") or "")
    base_row = rows_by_id.get(base_case_id)
    if base_row is None:
        return None
    repair_scope = str(metadata.get("repair_scope") or "")
    if not repair_scope:
        return None
    try:
        base_gold = json.loads(base_row["messages"][1]["content"])
    except (KeyError, TypeError, json.JSONDecodeError):
        return None
    previous_output = _corrupt_output(base_gold, repair_scope, str(harness["urgency_floor"]))
    if previous_output is None:
        return None
    validation = validate_navigator_output(
        previous_output,
        known_card_ids=known_card_ids(),
        urgency_floor=str(harness["urgency_floor"]),
        confirmed_intake=spec["structured_intake"],
        rule_results=harness["rule_results"],
        retrieved_card_ids=set(harness["retrieved_ids"]),
        retrieved_cards=harness["retrieved"],
        strict_schema=True,
    ).to_dict()
    focused_prompt = next(
        (
            item
            for item in build_focused_repair_prompts(
                original_prompt=str(harness["prompt"]),
                previous_output=previous_output,
                failures=list(validation.get("failures", []))
                + _extra_failures_for_scope(previous_output, spec, repair_scope),
                urgency_floor=str(harness["urgency_floor"]),
                required_observation_targets=required_observation_targets(harness["retrieved"]),
            )
            if item.scope.name == repair_scope
        ),
        None,
    )
    return focused_prompt.prompt if focused_prompt is not None else None


def _case_specs_for_rows(rows: list[dict[str, Any]], *, source_case_specs: list[Path]) -> list[dict[str, Any]]:
    specs_by_id: dict[str, dict[str, Any]] = {}
    for path in source_case_specs:
        for spec in _read_jsonl(path):
            case_id = str(spec.get("case_id") or "")
            if case_id:
                specs_by_id[case_id] = spec

    needed_ids = {_row_case_id(row) for row in rows}
    missing = sorted(needed_ids - set(specs_by_id))
    if missing:
        raise ValueError(f"missing case specs for rows: {', '.join(missing[:10])}")
    return [specs_by_id[case_id] for case_id in sorted(needed_ids)]


def _row_case_id(row: dict[str, Any]) -> str:
    metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
    return str(metadata.get("base_case_id") or row.get("case_id") or row.get("uuid") or "")


def _dedupe_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    rows_by_id: dict[str, dict[str, Any]] = {}
    for row in rows:
        row_id = str(row.get("case_id") or row.get("uuid") or "")
        if not row_id:
            raise ValueError("row missing case_id/uuid")
        if row_id in rows_by_id:
            raise ValueError(f"duplicate row id: {row_id}")
        rows_by_id[row_id] = row
    return [rows_by_id[row_id] for row_id in sorted(rows_by_id)]


def _task_type(row: dict[str, Any]) -> str:
    metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
    return str(metadata.get("task_type") or "navigator_full")


def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        raise FileNotFoundError(path)
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.write_text("".join(json.dumps(row, sort_keys=True) + "\n" for row in rows), encoding="utf-8")


def _sha256_path(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for chunk in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _run_json_command(cmd: list[str]) -> dict[str, Any]:
    completed = subprocess.run(cmd, check=True, text=True, capture_output=True)
    return json.loads(completed.stdout)


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