inputs/repro_agentselect/audit_release.py

#!/usr/bin/env python3
"""Audit exact AgentSelect release counts, integrity, and partition topology."""

from __future__ import annotations

import argparse
import hashlib
import json
from collections import Counter
from pathlib import Path
from typing import Any


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


def gini(values: list[int]) -> float:
    if not values or not sum(values):
        return 0.0
    values = sorted(values)
    n = len(values)
    weighted = sum((index + 1) * value for index, value in enumerate(values))
    return (2 * weighted) / (n * sum(values)) - (n + 1) / n


def top_share(counts: Counter[str], fraction: float) -> float:
    if not counts:
        return 0.0
    total = sum(counts.values())
    take = max(1, round(len(counts) * fraction))
    return sum(value for _, value in counts.most_common(take)) / total


def load(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    report: dict[str, Any] = {"parts": {}, "source_manifest": []}
    all_queries = all_agents = all_protocol_positive_pairs = 0
    for part in ("PartI", "PartII", "PartIII"):
        root = args.data_root / part
        paths = {kind: root / f"{kind}.json" for kind in ("agents", "questions", "rankings")}
        agents = load(paths["agents"])
        questions = load(paths["questions"])
        payload = load(paths["rankings"])
        rankings: dict[str, list[str]] = payload["rankings"]
        raw_pairs = sum(map(len, rankings.values()))
        # The paper's claimed 251,103 total decomposes as Part I top-5,
        # Part II's sole listed positive, and all listed Part III positives.
        # That arithmetic differs from the later code's Part I top-10 / Part
        # III top-5 training defaults, so both conventions are reported.
        claim_cutoff = {"PartI": 5, "PartII": 1, "PartIII": None}[part]
        code_cutoff = {"PartI": 10, "PartII": 1, "PartIII": 5}[part]
        claim_lists = [ranked if claim_cutoff is None else ranked[:claim_cutoff] for ranked in rankings.values()]
        claim_pairs = sum(map(len, claim_lists))
        code_pairs = sum(len(ranked[:code_cutoff]) for ranked in rankings.values())
        reuse = Counter(aid for ranked in claim_lists for aid in ranked)
        all_queries += len(questions)
        all_agents += len(agents)
        all_protocol_positive_pairs += claim_pairs
        report["parts"][part] = {
            "task": payload.get("task"),
            "questions_catalog": len(questions),
            "questions_with_rankings": len(rankings),
            "agents_catalog": len(agents),
            "raw_ranked_pairs": raw_pairs,
            "claim_counting_cutoff": "all listed" if claim_cutoff is None else claim_cutoff,
            "claim_counting_positive_pairs": claim_pairs,
            "official_code_training_cutoff": code_cutoff,
            "official_code_training_positive_pairs": code_pairs,
            "ranking_length_min": min(map(len, rankings.values())),
            "ranking_length_max": max(map(len, rankings.values())),
            "unique_agents_in_protocol_positives": len(reuse),
            "agent_reuse_mean_for_seen_agents": claim_pairs / len(reuse),
            "agent_reuse_median_for_seen_agents": sorted(reuse.values())[len(reuse) // 2],
            "agent_reuse_gini": gini(list(reuse.values())),
            "top_1pct_positive_share": top_share(reuse, 0.01),
            "top_5pct_positive_share": top_share(reuse, 0.05),
            "positive_density": claim_pairs / (len(questions) * len(agents)),
            "referential_integrity": {
                "ranked_questions_missing_from_question_catalog": len(set(rankings) - set(questions)),
                "ranked_agents_missing_from_agent_catalog": len(
                    {aid for ranked in rankings.values() for aid in ranked if aid not in agents}
                ),
            },
        }
        for kind, path in paths.items():
            report["source_manifest"].append({"path": str(path.relative_to(args.data_root)), "sha256": sha256(path)})
    tools_path = args.data_root / "Tools" / "tools.json"
    report["tools_catalog"] = len(load(tools_path))
    report["source_manifest"].append({"path": str(tools_path.relative_to(args.data_root)), "sha256": sha256(tools_path)})
    report["current_release_totals_with_paper_claim_counting_convention"] = {
        "queries": all_queries,
        "agents": all_agents,
        "positive_pairs": all_protocol_positive_pairs,
    }
    (args.output_dir / "release_audit.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
    with (args.output_dir / "release_manifest.tsv").open("w", encoding="utf-8") as f:
        f.write("path\tsha256\n")
        for entry in report["source_manifest"]:
            f.write(f"{entry['path']}\t{entry['sha256']}\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()