GEWU/scripts/prepare_answer_batches.py

#!/usr/bin/env python3
"""Prepare bounded answer-agent batches and validate their JSON outputs.

The prepared batches contain only tuning questions, never evaluator answers or
reserved questions. This tool does not call a model or mutate GEWU data.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

MAX_BATCH_SIZE = 4
MAX_SEARCH_CALLS = 10
MAX_SEARCH_BUDGET = 40_000


def load_tuning_questions(path: Path) -> list[dict[str, Any]]:
    raw = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(raw, list) or not raw:
        raise ValueError("Question input must be a nonempty JSON array")
    seen: set[str] = set()
    selected: list[dict[str, Any]] = []
    for row in raw:
        if not isinstance(row, dict) or not isinstance(row.get("id"), str):
            raise ValueError("Each question must be an object with a string id")
        question_id = row["id"]
        if question_id in seen:
            raise ValueError(f"Duplicate question id: {question_id}")
        seen.add(question_id)
        if row.get("partition") == "tuning":
            if not isinstance(row.get("question"), str) or not row["question"].strip():
                raise ValueError(f"Question {question_id} has no question text")
            selected.append({"id": question_id, "question": row["question"]})
    if not selected:
        raise ValueError("No tuning questions found")
    return selected


def build_batches(questions: list[dict[str, Any]], batch_size: int) -> list[list[dict[str, Any]]]:
    if batch_size < 1 or batch_size > MAX_BATCH_SIZE:
        raise ValueError(f"batch-size must be between 1 and {MAX_BATCH_SIZE}")
    return [questions[i : i + batch_size] for i in range(0, len(questions), batch_size)]


def render_prompt(batch_id: str, questions: list[dict[str, Any]]) -> str:
    payload = json.dumps(questions, ensure_ascii=False, indent=2)
    return f"""# Bounded GEWU answer batch {batch_id}

Answer only the listed tuning questions using accepted knowledge in project
`dome-performance` and shared scope as authorized. Do not use evaluator references,
outside knowledge, shell commands, or source-file instructions. Do not write or change
GEWU data.

## Hard run limits

- This batch contains {len(questions)} questions. Complete only these IDs:
  {', '.join(q['id'] for q in questions)}.
- Make at most {MAX_SEARCH_CALLS} total `search` calls for the whole batch. Use one
  focused search per question, plus at most one targeted follow-up if the first result
  is incomplete or reports `BUDGET_TRUNCATED`.
- Never repeat an identical search. Never raise `budget_tokens` above
  {MAX_SEARCH_BUDGET}; the service response limit is 50,000 codepoints.
- Use `limit: 5`. If evidence remains incomplete, state the precise gap and set
  `evidence_gap: true`.
- Cite only evidence returned by GEWU. Preserve attribution, qualifications, units,
  and exact table values. Keep claims separate when evidence differs.
- If the call limit is reached, stop searching and return partial answers for every
  batch question; do not omit IDs.

## Output contract

Return one JSON object and no surrounding prose, with this exact shape:

{{"answers":[{{"id":"Qxx","answer":"...","citations":[],"evidence_gap":false}}]}}

Include exactly one answer per requested ID, in the listed order. Each citation must
be a complete GEWU evidence citation with `node_id`, `claim_id`, `source_revision_id`,
`extraction_id`, `locator` and exact `quote`. Use an empty citations array only when no
relevant accepted evidence was returned. Do not include evaluator scores, expected
points, or reserved questions.

## Questions

{payload}
"""


def prepare(args: argparse.Namespace) -> dict[str, Any]:
    questions = load_tuning_questions(args.questions)
    batches = build_batches(questions, args.batch_size)
    out = args.output.resolve()
    out.mkdir(parents=True, exist_ok=False, mode=0o700)
    index: dict[str, Any] = {
        "schema_version": 1,
        "status": "prepared_offline",
        "source_questions": str(args.questions.resolve()),
        "batch_count": len(batches),
        "question_count": len(questions),
        "max_batch_size": args.batch_size,
        "max_search_calls_per_batch": MAX_SEARCH_CALLS,
        "max_search_budget_tokens": MAX_SEARCH_BUDGET,
        "batches": [],
    }
    for number, batch in enumerate(batches, start=1):
        batch_id = f"batch-{number:02d}"
        (out / f"{batch_id}-questions.json").write_text(
            json.dumps(batch, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        (out / f"{batch_id}-prompt.md").write_text(render_prompt(batch_id, batch), encoding="utf-8")
        index["batches"].append(
            {"id": batch_id, "question_ids": [item["id"] for item in batch], "status": "pending"}
        )
    (out / "batch-index.json").write_text(
        json.dumps(index, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    for path in out.iterdir():
        path.chmod(0o600)
    return index


def validate_answer_output(expected_path: Path, answer_path: Path) -> dict[str, Any]:
    expected = json.loads(expected_path.read_text(encoding="utf-8"))
    actual_doc = json.loads(answer_path.read_text(encoding="utf-8"))
    if not isinstance(expected, list) or not isinstance(actual_doc, dict):
        raise ValueError("Expected question array and answer object")
    if set(actual_doc) != {"answers"} or not isinstance(actual_doc["answers"], list):
        raise ValueError("Answer output must contain only an answers array")
    expected_ids = [row["id"] for row in expected]
    rows = actual_doc["answers"]
    actual_ids = [row.get("id") if isinstance(row, dict) else None for row in rows]
    issues: list[str] = []
    if actual_ids != expected_ids:
        issues.append("answer IDs must exactly match the batch in order")
    for index, row in enumerate(rows):
        if not isinstance(row, dict) or set(row) != {"id", "answer", "citations", "evidence_gap"}:
            issues.append(f"answer row {index} has an invalid shape")
            continue
        if not isinstance(row["answer"], str) or not row["answer"].strip():
            issues.append(f"answer row {index} has empty answer text")
        if not isinstance(row["citations"], list):
            issues.append(f"answer row {index} citations must be an array")
            continue
        if not isinstance(row["evidence_gap"], bool):
            issues.append(f"answer row {index} evidence_gap must be boolean")
        for citation_index, citation in enumerate(row["citations"]):
            required = {
                "node_id",
                "claim_id",
                "source_revision_id",
                "extraction_id",
                "locator",
                "quote",
            }
            if not isinstance(citation, dict) or not required.issubset(citation):
                issues.append(f"answer row {index} citation {citation_index} is incomplete")
    return {
        "schema_version": 1,
        "status": "valid" if not issues else "invalid",
        "expected_question_ids": expected_ids,
        "actual_question_ids": actual_ids,
        "answer_count": len(rows),
        "issues": issues,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    prepare_parser = commands.add_parser("prepare")
    prepare_parser.add_argument("--questions", type=Path, required=True)
    prepare_parser.add_argument("--output", type=Path, required=True)
    prepare_parser.add_argument("--batch-size", type=int, default=MAX_BATCH_SIZE)
    validate_parser = commands.add_parser("validate")
    validate_parser.add_argument("--expected", type=Path, required=True)
    validate_parser.add_argument("--answers", type=Path, required=True)
    args = parser.parse_args()
    if args.command == "prepare":
        result = prepare(args)
    else:
        result = validate_answer_output(args.expected, args.answers)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if result.get("status") == "invalid":
        sys.exit(1)


if __name__ == "__main__":
    main()