data_studio/benchmarks.py

from __future__ import annotations

from typing import Any

from pymongo.database import Database

from data_studio.identity import payload_hash
from data_studio.utils import utc_now


def register_protected_benchmark(
    db: Database,
    benchmark_id: str,
    exact_revision: str,
    task: str,
    exact_payloads: list[dict[str, Any]],
    actor: str,
) -> None:
    fingerprints = sorted({payload_hash(payload) for payload in exact_payloads})
    db["protected_benchmarks"].insert_one(
        {
            "schema_version": "1.0",
            "benchmark_id": benchmark_id,
            "exact_revision": exact_revision,
            "task": task,
            "state": "protected",
            "exact_fingerprints": fingerprints,
            "created_by": actor,
            "created_at": utc_now(),
        }
    )


def exact_benchmark_matches(db: Database, payloads: list[dict[str, Any]]) -> list[dict[str, str]]:
    hashes = {payload_hash(payload) for payload in payloads}
    matches: list[dict[str, str]] = []
    for row in db["protected_benchmarks"].find({"state": "protected"}):
        overlap = hashes.intersection(row.get("exact_fingerprints", []))
        for fingerprint in sorted(overlap):
            matches.append({"benchmark_id": row["benchmark_id"], "fingerprint": fingerprint})
    return matches


def assert_no_exact_benchmark_matches(db: Database, payloads: list[dict[str, Any]]) -> None:
    matches = exact_benchmark_matches(db, payloads)
    if matches:
        raise ValueError(f"Protected benchmark exact-match leakage detected: {matches}")