PlantAI-rag/scripts/run_qa_v1_0_0_ablation.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from typing import Any, Dict, Iterable, List

import requests


VARIANTS = {
    "v1_no_draft_one_round": {
        "use_provisional_draft": False,
        "max_rounds": 1,
    },
    "v1_draft_one_round": {
        "use_provisional_draft": True,
        "max_rounds": 1,
    },
    "v1_draft_two_rounds": {
        "use_provisional_draft": True,
        "max_rounds": 2,
    },
}


def load_questions(path: Path) -> List[Dict[str, Any]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(payload, dict):
        for key in ("questions", "items", "data"):
            if isinstance(payload.get(key), list):
                payload = payload[key]
                break
    if not isinstance(payload, list):
        raise ValueError("Question JSON must be a list or contain questions/items/data")
    output = []
    for index, item in enumerate(payload, 1):
        if isinstance(item, str):
            output.append({"question_id": f"Q{index:03d}", "question": item})
            continue
        if not isinstance(item, dict):
            continue
        question = str(
            item.get("question")
            or item.get("query")
            or item.get("text")
            or ""
        ).strip()
        if question:
            output.append(
                {
                    **item,
                    "question_id": str(
                        item.get("question_id")
                        or item.get("id")
                        or f"Q{index:03d}"
                    ),
                    "question": question,
                }
            )
    return output


def run_one(
    endpoint: str,
    question: Dict[str, Any],
    variant_id: str,
    timeout: float,
) -> Dict[str, Any]:
    config = VARIANTS[variant_id]
    started = time.perf_counter()
    try:
        response = requests.post(
            endpoint,
            json={
                "question": question["question"],
                "debug": False,
                **config,
            },
            timeout=timeout,
        )
        payload = response.json()
        return {
            "question_id": question["question_id"],
            "question": question["question"],
            "variant_id": variant_id,
            "http_status": response.status_code,
            "success": bool(payload.get("success")) and response.status_code == 200,
            "answer": payload.get("answer", ""),
            "source_count": payload.get("source_count", 0),
            "round_count": len(payload.get("round_history", []) or []),
            "audit_path": payload.get("audit_path", ""),
            "timing": payload.get("timing", {}),
            "retrieval_hypothesis": payload.get("retrieval_hypothesis", {}),
            "claim_verification": payload.get(
                "provisional_claim_verification", []
            ),
            "evidence_assessment": payload.get("evidence_assessment", {}),
            "sources": payload.get("sources", []),
            "error": payload.get("detail", ""),
            "client_elapsed": time.perf_counter() - started,
        }
    except Exception as exc:
        return {
            "question_id": question["question_id"],
            "question": question["question"],
            "variant_id": variant_id,
            "success": False,
            "error": f"{type(exc).__name__}: {exc}",
            "client_elapsed": time.perf_counter() - started,
        }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--questions", type=Path, required=True)
    parser.add_argument(
        "--endpoint",
        default="http://127.0.0.1:7909/api/v1/qa-v1/chat",
    )
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--variants",
        nargs="+",
        choices=list(VARIANTS),
        default=list(VARIANTS),
    )
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--timeout", type=float, default=900)
    args = parser.parse_args()

    questions = load_questions(args.questions)
    if args.limit > 0:
        questions = questions[: args.limit]
    args.output_dir.mkdir(parents=True, exist_ok=True)

    summary = {
        "endpoint": args.endpoint,
        "question_file": str(args.questions),
        "question_count": len(questions),
        "variants": args.variants,
        "results": {},
    }
    for variant_id in args.variants:
        variant_dir = args.output_dir / variant_id
        variant_dir.mkdir(parents=True, exist_ok=True)
        successes = 0
        failures = 0
        for index, question in enumerate(questions, 1):
            target = variant_dir / f"{question['question_id']}.json"
            if target.exists():
                try:
                    existing = json.loads(target.read_text(encoding="utf-8"))
                    if existing.get("success"):
                        successes += 1
                        continue
                except Exception:
                    pass
            result = run_one(
                args.endpoint,
                question,
                variant_id,
                args.timeout,
            )
            target.write_text(
                json.dumps(result, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            if result.get("success"):
                successes += 1
            else:
                failures += 1
            print(
                f"[{variant_id}] {index}/{len(questions)} "
                f"{question['question_id']} success={result.get('success')}",
                flush=True,
            )
        summary["results"][variant_id] = {
            "successes": successes,
            "failures": failures,
        }
        (args.output_dir / "run_summary.json").write_text(
            json.dumps(summary, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )


if __name__ == "__main__":
    main()