scripts/run_experiments/run_doc_to_model_experiment.py

#!/usr/bin/env python3
"""
Ablation experiment: Direct spec vs. Document-preprocessed spec → UPPAAL model.

For each spec in spec.csv we run two conditions:

  (A) DIRECT   — clean nl_spec text → FormalModelBuildingPipeline
  (B) DOC-RAG  — synthetic noisy document (nl_spec wrapped in boilerplate)
                 → DocumentSpecRewriterAgent → FormalModelBuildingPipeline

The synthetic document simulates a real requirements doc that contains the
relevant spec information embedded in surrounding noise (background, goals,
references, implementation notes).  This tests whether the rewriter can
faithfully recover the spec from a messier input.

Output: artifacts/doc_to_model_experiment/results_<timestamp>.json
        artifacts/doc_to_model_experiment/summary_<timestamp>.txt

Usage:
  python scripts/run_experiments/run_doc_to_model_experiment.py
  python scripts/run_experiments/run_doc_to_model_experiment.py --spec-ids 1 2 3
  python scripts/run_experiments/run_doc_to_model_experiment.py --conditions direct
  python scripts/run_experiments/run_doc_to_model_experiment.py --conditions doc
  python scripts/run_experiments/run_doc_to_model_experiment.py --fast
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
import textwrap
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))

from frame.pipeline.document_to_formal_pipeline import DocumentToFormalPipeline
from frame.pipeline.formal_model_building_pipeline import FormalModelBuildingPipeline

SPEC_CSV = ROOT / "datasets" / "spec.csv"
OUT_ROOT = ROOT / "artifacts" / "doc_to_model_experiment"

# ---------------------------------------------------------------------------
# Synthetic document builder
# ---------------------------------------------------------------------------

_DOC_TEMPLATE = textwrap.dedent("""\
    TECHNICAL REPORT — {system_name}
    Version 1.0 | Domain: {domain}
    ============================================================

    1. INTRODUCTION

    This report documents the design and requirements for the {system_name}.
    The system has been developed to address operational needs in the
    {domain} domain.  This document is intended for system engineers and
    verification teams.

    2. BACKGROUND AND MOTIVATION

    Automated systems in the {domain} domain must satisfy strict correctness
    and timing requirements.  Prior approaches relied on manual inspection,
    which is error-prone and does not scale.  This system provides a
    structured solution with well-defined behaviour.

    3. SCOPE

    This document covers functional requirements, timing constraints, and
    interface definitions for the {system_name}.  Security, deployment, and
    maintenance procedures are out of scope.

    4. SYSTEM REQUIREMENTS AND BEHAVIOUR

    The following section describes the functional behaviour of the system.

    {nl_spec}

    5. IMPLEMENTATION NOTES

    The system shall be implemented using industry-standard components.
    Integration testing shall verify all specified transitions.  Performance
    benchmarks are defined in a separate document.

    6. REFERENCES

    [1] ISO 26262 — Functional Safety
    [2] IEC 61508 — Safety Instrumented Systems
    [3] Internal design document v0.9 (restricted)

    ============================================================
    END OF DOCUMENT
""")


def make_synthetic_doc(row: Dict[str, str]) -> str:
    return _DOC_TEMPLATE.format(
        system_name=row.get("system_name", "System"),
        domain=row.get("domain", "general"),
        nl_spec=row.get("nl_spec", "").strip(),
    )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def load_specs(spec_csv: Path, spec_ids: Optional[List[int]] = None) -> List[Dict]:
    rows = []
    with spec_csv.open(encoding="utf-8", newline="") as f:
        for row in csv.DictReader(f):
            sid = int(row["spec_id"])
            if spec_ids is None or sid in spec_ids:
                rows.append(row)
    return rows


def smoke_passed(summary) -> bool:
    if summary is None:
        return False
    return bool(summary.smoke_ok)


def xml_written(summary) -> bool:
    if summary is None:
        return False
    return summary.xml_path is not None


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--spec-ids", type=int, nargs="+", default=None,
                    metavar="ID", help="Subset of spec IDs to run (default: all 20)")
    ap.add_argument("--conditions", nargs="+", choices=["direct", "doc"],
                    default=["direct", "doc"],
                    help="Which conditions to run (default: both)")
    ap.add_argument("--llm-model", type=str, default="gpt-4.1",
                    help="LLM for spec analysis + rewriting")
    ap.add_argument("--builder-model", type=str, default="gpt-4o-mini",
                    help="LLM for IR building")
    ap.add_argument("--rewriter-model", type=str, default=None,
                    help="LLM for doc→spec rewriting (default: same as --llm-model)")
    ap.add_argument("--fast", action="store_true",
                    help="Use gpt-4o-mini for everything (faster, cheaper)")
    ap.add_argument("--no-smoke", action="store_true",
                    help="Skip verifyta smoke test")
    ap.add_argument("--out-dir", type=Path, default=OUT_ROOT)
    args = ap.parse_args()

    if args.fast:
        args.llm_model = "gpt-4o-mini"
        args.builder_model = "gpt-4o-mini"

    rewriter_model = args.rewriter_model or args.llm_model

    specs = load_specs(SPEC_CSV, args.spec_ids)
    if not specs:
        print("No specs found.", file=sys.stderr)
        return 1

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = args.out_dir / timestamp
    run_dir.mkdir(parents=True, exist_ok=True)

    print(f"Specs: {len(specs)}  |  Conditions: {args.conditions}")
    print(f"Output: {run_dir}\n")

    # Build pipelines once (shared across specs)
    direct_pipeline: Optional[FormalModelBuildingPipeline] = None
    doc_pipeline: Optional[DocumentToFormalPipeline] = None

    common_formal_kwargs = dict(
        stub_analyzer=False,
        stub_agents=False,
        llm_model=args.llm_model,
        builder_model_name=args.builder_model,
        builder_max_tokens=12000,
        enable_verifyta_smoke=not args.no_smoke,
    )

    if "direct" in args.conditions:
        direct_pipeline = FormalModelBuildingPipeline(**common_formal_kwargs)

    if "doc" in args.conditions:
        doc_pipeline = DocumentToFormalPipeline(
            rewriter_model=rewriter_model,
            rewriter_max_tokens=2048,
            **common_formal_kwargs,
        )

    results = []

    for row in specs:
        sid = int(row["spec_id"])
        name = row.get("system_name", f"S{sid:02d}")
        nl_spec = row.get("nl_spec", "").strip()

        print(f"\n{'='*60}")
        print(f"Spec {sid:02d} — {name}")
        print("="*60)

        row_result: Dict = {"spec_id": sid, "system_name": name}

        # --- Condition A: direct ---
        if "direct" in args.conditions and direct_pipeline is not None:
            out_dir_direct = run_dir / f"S{sid:02d}" / "direct"
            print(f"\n[Direct] {name}")
            try:
                fs = direct_pipeline.run(nl_spec, out_dir_direct)
                row_result["direct"] = {
                    "xml_written": xml_written(fs),
                    "smoke_ok": smoke_passed(fs),
                    "validator_ok": fs.validator_ok,
                    "degraded_export": fs.degraded_export,
                    "error": fs.error,
                    "xml_path": fs.xml_path,
                }
                status = "SMOKE OK" if smoke_passed(fs) else ("XML" if xml_written(fs) else "FAIL")
                print(f"    → {status}")
            except Exception as exc:
                print(f"    → ERROR: {exc}", file=sys.stderr)
                row_result["direct"] = {"xml_written": False, "smoke_ok": False,
                                         "error": str(exc)}

        # --- Condition B: doc preprocessing ---
        if "doc" in args.conditions and doc_pipeline is not None:
            out_dir_doc = run_dir / f"S{sid:02d}" / "doc"
            synthetic_doc = make_synthetic_doc(row)
            print(f"\n[Doc→Spec] {name}  ({len(synthetic_doc)} chars)")
            try:
                ds = doc_pipeline.run_from_text(synthetic_doc, out_dir_doc)
                fs = ds.formal
                row_result["doc"] = {
                    "rewriter_ok": ds.rewriter_ok,
                    "rewriter_error": ds.rewriter_error,
                    "used_map_reduce": ds.used_map_reduce,
                    "extracted_spec_path": ds.extracted_spec_path,
                    "xml_written": xml_written(fs),
                    "smoke_ok": smoke_passed(fs),
                    "validator_ok": fs.validator_ok if fs else None,
                    "degraded_export": fs.degraded_export if fs else None,
                    "error": fs.error if fs else ds.rewriter_error,
                    "xml_path": fs.xml_path if fs else None,
                }
                if ds.rewriter_ok:
                    status = "SMOKE OK" if smoke_passed(fs) else ("XML" if xml_written(fs) else "FAIL")
                else:
                    status = f"REWRITE FAIL: {ds.rewriter_error}"
                print(f"    → {status}")
            except Exception as exc:
                print(f"    → ERROR: {exc}", file=sys.stderr)
                row_result["doc"] = {"rewriter_ok": False, "xml_written": False,
                                      "smoke_ok": False, "error": str(exc)}

        results.append(row_result)

    # --- Save results ---
    results_path = run_dir / f"results_{timestamp}.json"
    results_path.write_text(json.dumps(results, indent=2, ensure_ascii=False),
                            encoding="utf-8")
    print(f"\nResults JSON: {results_path}")

    # --- Print summary table ---
    summary_lines = _build_summary(results, args.conditions)
    summary_text = "\n".join(summary_lines)
    print("\n" + summary_text)

    summary_path = run_dir / f"summary_{timestamp}.txt"
    summary_path.write_text(summary_text, encoding="utf-8")
    print(f"Summary: {summary_path}")

    return 0


def _build_summary(results: List[Dict], conditions: List[str]) -> List[str]:
    lines = []
    header_parts = ["Spec", "Name".ljust(30)]
    if "direct" in conditions:
        header_parts += ["Direct XML", "Direct Smoke"]
    if "doc" in conditions:
        header_parts += ["Doc XML", "Doc Smoke", "Rewrite OK"]
    lines.append("  ".join(header_parts))
    lines.append("-" * 80)

    direct_xml_total = direct_smoke_total = 0
    doc_xml_total = doc_smoke_total = doc_rewrite_total = 0
    n = len(results)

    for r in results:
        sid = r["spec_id"]
        name = r.get("system_name", "")[:28].ljust(30)
        parts = [f"S{sid:02d}", name]

        if "direct" in conditions:
            d = r.get("direct", {})
            dx = "YES" if d.get("xml_written") else "NO "
            ds = "YES" if d.get("smoke_ok") else "NO "
            parts += [dx.center(10), ds.center(12)]
            direct_xml_total += int(bool(d.get("xml_written")))
            direct_smoke_total += int(bool(d.get("smoke_ok")))

        if "doc" in conditions:
            doc = r.get("doc", {})
            dx = "YES" if doc.get("xml_written") else "NO "
            ds = "YES" if doc.get("smoke_ok") else "NO "
            rw = "YES" if doc.get("rewriter_ok") else "NO "
            parts += [dx.center(7), ds.center(9), rw.center(10)]
            doc_xml_total += int(bool(doc.get("xml_written")))
            doc_smoke_total += int(bool(doc.get("smoke_ok")))
            doc_rewrite_total += int(bool(doc.get("rewriter_ok")))

        lines.append("  ".join(parts))

    lines.append("-" * 80)
    totals = ["TOTAL", f"n={n}".ljust(30)]
    if "direct" in conditions:
        totals += [
            f"{direct_xml_total}/{n}".center(10),
            f"{direct_smoke_total}/{n}".center(12),
        ]
    if "doc" in conditions:
        totals += [
            f"{doc_xml_total}/{n}".center(7),
            f"{doc_smoke_total}/{n}".center(9),
            f"{doc_rewrite_total}/{n}".center(10),
        ]
    lines.append("  ".join(totals))

    if "direct" in conditions and "doc" in conditions and n > 0:
        lines.append("")
        lines.append("MCR comparison:")
        lines.append(f"  Direct  XML rate:  {direct_xml_total/n:.2%}")
        lines.append(f"  Doc     XML rate:  {doc_xml_total/n:.2%}")
        lines.append(f"  Direct  Smoke rate: {direct_smoke_total/n:.2%}")
        lines.append(f"  Doc     Smoke rate: {doc_smoke_total/n:.2%}")

    return lines


if __name__ == "__main__":
    sys.exit(main())