learningcircuit--local-deep-research/tests/performance/_shared/build_eval_dataset.py

"""Build an evaluation dataset: cross-product of queries × search engines × judge LLMs.

For each (query, engine, model) cell, runs `quick_summary` end-to-end and
saves (a) the formatted report and (b) a verbose log capturing the
filter's KEPT/REMOVED decisions. Writes a summary CSV so you can rank
runs by source count, timing, and visually inspect outliers.

Resumable: cells whose report file already exists are skipped unless
--force is passed. Use this to recover from interrupted runs or to
extend an existing dataset with new rows.

Usage (defaults produce a ~12-cell grid ≈ 60-90min total on qwen3.5:9b):
    LDR_BOOTSTRAP_ALLOW_UNENCRYPTED=true LDR_TESTING_WITH_MOCKS=false \\
    pdm run python tests/performance/_shared/build_eval_dataset.py \\
        --output-dir ./ldr_eval_output

Override any axis from the CLI:
    ... --queries "q1|q2" --engines "arxiv,openalex" --models "qwen3.5:9b,gemma3:12b"

Be conservative with --parallel — each parallel slot loads an LLM on the
Ollama endpoint; running 2+ different model tags concurrently may OOM
the server. Same-model parallel is usually fine (Ollama queues).
"""

from __future__ import annotations

import argparse
import csv
import os
import re
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path

DEFAULT_QUERIES = [
    "LLM interpretability latest research",
    "mechanistic interpretability of transformer language models",
    "sparse autoencoders for neural network feature discovery",
    "safety alignment and refusal in large language models",
]
DEFAULT_ENGINES = ["arxiv", "openalex"]
DEFAULT_MODELS = ["qwen3.5:9b", "gemma3:12b", "ministral-3:14b"]


def slugify(text: str, max_len: int = 50) -> str:
    s = re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")
    return s[:max_len]


@dataclass
class Cell:
    query: str
    engine: str
    model: str

    def key(self) -> str:
        return f"{self.engine}__{slugify(self.model)}__{slugify(self.query)}"


def already_done(cell: Cell, out_dir: Path) -> bool:
    return (out_dir / "reports" / f"{cell.key()}.md").exists()


def run_cell(cell: Cell, out_dir: Path, ollama_url: str) -> dict:
    """Invoke run_full_search.py as a subprocess for isolation between cells."""
    report_path = out_dir / "reports" / f"{cell.key()}.md"
    log_path = out_dir / "logs" / f"{cell.key()}.log"

    script = Path(__file__).parent / "run_full_search.py"
    cmd = [
        sys.executable,
        str(script),
        "--query",
        cell.query,
        "--engine",
        cell.engine,
        "--model",
        cell.model,
        "--ollama-url",
        ollama_url,
        "--verbose",
        "--output",
        str(report_path),
    ]

    env = os.environ.copy()
    env.setdefault("LDR_BOOTSTRAP_ALLOW_UNENCRYPTED", "true")
    env.setdefault("LDR_TESTING_WITH_MOCKS", "false")

    t0 = time.monotonic()
    try:
        with log_path.open("w") as logf:
            proc = subprocess.run(
                cmd,
                env=env,
                stdout=logf,
                stderr=subprocess.STDOUT,
                timeout=60 * 30,  # 30-min hard cap per cell
                check=False,
            )
        elapsed = time.monotonic() - t0
        exit_code = proc.returncode
    except subprocess.TimeoutExpired:
        elapsed = time.monotonic() - t0
        exit_code = -1
        # A partially-written report would otherwise cause ``already_done``
        # to skip this cell on the next run even though it never finished.
        report_path.unlink(missing_ok=True)
    except Exception as exc:
        elapsed = time.monotonic() - t0
        exit_code = -2
        log_path.write_text(f"Exception: {exc!r}\n")
        report_path.unlink(missing_ok=True)

    result = {
        "query": cell.query,
        "engine": cell.engine,
        "model": cell.model,
        "elapsed_s": round(elapsed, 1),
        "exit_code": exit_code,
        "report_path": str(report_path) if report_path.exists() else "",
        "sources": 0,
        "log_path": str(log_path),
    }

    if report_path.exists():
        text = report_path.read_text()
        m = re.search(r"^\s*-\s*\*\*Sources:\*\*\s*(\d+)", text, re.M)
        if m:
            result["sources"] = int(m.group(1))

    return result


def parse_list(
    s: str, sep_primary: str = "|", sep_fallback: str = ","
) -> list[str]:
    """Parse a delimiter-separated list — '|' preferred for queries with commas."""
    sep = sep_primary if sep_primary in s else sep_fallback
    return [x.strip() for x in s.split(sep) if x.strip()]


def main() -> int:
    p = argparse.ArgumentParser(
        description="Build a cross-product eval dataset (query × engine × model).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument(
        "--queries",
        default=None,
        help="Queries separated by '|'. Default: a 4-query interpretability set.",
    )
    p.add_argument(
        "--engines",
        default=",".join(DEFAULT_ENGINES),
        help="Comma-separated engine names.",
    )
    p.add_argument(
        "--models",
        default=",".join(DEFAULT_MODELS),
        help="Comma-separated Ollama model tags.",
    )
    p.add_argument(
        "--output-dir",
        default="./ldr_eval_output",
        help="Where reports/, logs/, summary.csv land.",
    )
    p.add_argument(
        "--ollama-url",
        default=os.environ.get(
            "LDR_TEST_OLLAMA_BASE_URL", "http://localhost:11434"
        ),
    )
    p.add_argument(
        "--parallel",
        type=int,
        default=1,
        help="Cells to run concurrently. 1 = sequential. 2+ may OOM Ollama if different models are loaded simultaneously.",
    )
    p.add_argument(
        "--force",
        action="store_true",
        help="Re-run cells even if their report file already exists.",
    )
    args = p.parse_args()

    queries = (
        parse_list(args.queries) if args.queries else list(DEFAULT_QUERIES)
    )
    engines = parse_list(args.engines)
    models = parse_list(args.models)

    out_dir = Path(args.output_dir).resolve()
    (out_dir / "reports").mkdir(parents=True, exist_ok=True)
    (out_dir / "logs").mkdir(parents=True, exist_ok=True)

    cells = [
        Cell(query=q, engine=e, model=m)
        for q in queries
        for e in engines
        for m in models
    ]

    total = len(cells)
    todo = [c for c in cells if args.force or not already_done(c, out_dir)]
    skipped = total - len(todo)

    print(
        f"Grid: {len(queries)} queries × {len(engines)} engines × {len(models)} models = {total} cells",
        file=sys.stderr,
    )
    print(f"  {skipped} already done (use --force to re-run)", file=sys.stderr)
    print(f"  {len(todo)} to run, parallel={args.parallel}", file=sys.stderr)
    print(f"  Output: {out_dir}", file=sys.stderr)

    summary_path = out_dir / "summary.csv"
    summary_rows: list[dict] = []
    if summary_path.exists():
        with summary_path.open() as f:
            summary_rows = list(csv.DictReader(f))

    def run_and_log(cell: Cell) -> dict:
        print(
            f"  >>> {cell.engine} | {cell.model} | {cell.query[:60]}",
            file=sys.stderr,
        )
        r = run_cell(cell, out_dir, args.ollama_url)
        status = "ok" if r["exit_code"] == 0 else f"exit={r['exit_code']}"
        print(
            f"      {status}  {r['elapsed_s']}s  sources={r['sources']}  {cell.key()}",
            file=sys.stderr,
        )
        return r

    t0 = time.monotonic()
    if args.parallel <= 1:
        for cell in todo:
            summary_rows = _upsert(summary_rows, run_and_log(cell))
            _write_summary(summary_path, summary_rows)
    else:
        with ThreadPoolExecutor(max_workers=args.parallel) as pool:
            futures = {pool.submit(run_and_log, c): c for c in todo}
            for fut in as_completed(futures):
                summary_rows = _upsert(summary_rows, fut.result())
                _write_summary(summary_path, summary_rows)

    elapsed = time.monotonic() - t0
    print(
        f"\nDataset build complete in {elapsed / 60:.1f}min.", file=sys.stderr
    )
    print(f"  Summary: {summary_path}", file=sys.stderr)
    _print_table(summary_rows)
    return 0


def _upsert(rows: list[dict], new: dict) -> list[dict]:
    """Replace a row with the same (query, engine, model) or append."""
    out = [
        r
        for r in rows
        if not (
            r.get("query") == new["query"]
            and r.get("engine") == new["engine"]
            and r.get("model") == new["model"]
        )
    ]
    out.append({k: str(v) for k, v in new.items()})
    return out


def _write_summary(path: Path, rows: list[dict]) -> None:
    if not rows:
        return
    fieldnames = [
        "query",
        "engine",
        "model",
        "sources",
        "elapsed_s",
        "exit_code",
        "report_path",
        "log_path",
    ]
    with path.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k, "") for k in fieldnames})


def _print_table(rows: list[dict]) -> None:
    if not rows:
        return
    print("\n=== Dataset summary ===")
    print(
        f"{'engine':10s} {'model':20s} {'sources':>8s} {'elapsed_s':>10s}  query"
    )
    for r in rows:
        print(
            f"{r.get('engine', ''):10s} {r.get('model', ''):20s} "
            f"{str(r.get('sources', '')):>8s} {str(r.get('elapsed_s', '')):>10s}  "
            f"{r.get('query', '')[:60]}"
        )


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