ely-eye/bench/tests/test_joint.py

"""JointRunner must rank exactly like ArmRunner, and the rankings cache must resume.

The ablation switched to a single pass that computes each signal once. If that pass
drifted from the per-arm path, the baseline would silently measure a different system.
"""

from __future__ import annotations

import hashlib
import json
from collections.abc import Sequence
from pathlib import Path
from typing import Any

import numpy as np
import pytest

from bench import systems
from bench.run import (
    Query,
    fingerprint,
    latency,
    load_rankings,
    remeasure_latency,
    scale_arm_names,
    scale_probes,
    write_markdown,
)
from bench.systems import ARMS, ArmRunner, Atom, JointRunner
from bench.tests.test_systems import ATOMS

MODEL_ARMS = ["bm25_only", "graph_only", "dense_only", "no_rerank", "fusion_rerank"]


def _queries() -> list[str]:
    return ["alpha report", "beta invoice", "gamma", ""]


class _TinyDense:
    """Stand-in embedding model: these tests compare two harness paths, not embedding quality."""

    def encode(self, texts: Sequence[str], **_: Any) -> np.ndarray:
        return np.stack([_unit_vector(str(text)) for text in texts]).astype(np.float32)


class _TinyReranker:
    """Stand-in cross-encoder with the same scores(query, texts) shape as Reranker."""

    def __init__(self, model_dir: Path) -> None:
        self.model_dir = model_dir

    def scores(self, query: str, texts: Sequence[str]) -> np.ndarray:
        return np.array([_score(query, text) for text in texts], dtype=np.float64)


def _unit_vector(text: str) -> np.ndarray:
    seed = int(hashlib.sha256(text.encode()).hexdigest()[:8], 16)
    vector = np.random.default_rng(seed).standard_normal(systems.DENSE_DIM)
    return vector / np.linalg.norm(vector)


def _score(query: str, text: str) -> float:
    digest = hashlib.sha256(f"{query}\n{text}".encode()).hexdigest()
    return int(digest[:8], 16) / 0xFFFFFFFF


def _stub_models(monkeypatch: pytest.MonkeyPatch) -> None:
    dense_dir = systems.local_model_dir(systems.MODELS, systems.DENSE_MODEL)
    monkeypatch.setattr(systems, "_MODELS", {str(dense_dir): _TinyDense()})
    monkeypatch.setattr(systems, "Reranker", _TinyReranker)


@pytest.mark.parametrize("arm", MODEL_ARMS)
def test_joint_ranking_matches_the_per_arm_runner(arm: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
    _stub_models(monkeypatch)
    queries = _queries()
    joint = JointRunner([arm], ATOMS, systems.MODELS)
    per_arm = ArmRunner(arm, ATOMS, systems.MODELS)
    for query in queries:
        expected = per_arm.rank(query)
        ranked, seconds = joint.rank_all(query)[arm]
        assert ranked == expected, f"{arm} drifted from the per-arm path on {query!r}"
        assert seconds >= 0.0


def test_joint_runner_computes_the_sparse_signal_once_per_query(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
    _stub_models(monkeypatch)
    joint = JointRunner(["bm25_only", "no_rerank", "fusion_rerank"], ATOMS, systems.MODELS)
    calls = 0
    assert joint._sparse is not None
    original = joint._sparse.scores

    def counting_scores(query: str):
        nonlocal calls
        calls += 1
        return original(query)

    monkeypatch.setattr(joint._sparse, "scores", counting_scores)
    joint.rank_all("alpha report")
    assert calls == 1, "three arms that all need BM25 must still score the corpus once"


def test_joint_runner_rejects_unknown_or_empty_arm_lists() -> None:
    with pytest.raises(SystemExit, match="unknown arm"):
        JointRunner(["vibes"], ATOMS, systems.MODELS)
    with pytest.raises(SystemExit, match="at least one arm"):
        JointRunner([], ATOMS, systems.MODELS)


def test_rankings_cache_resumes_and_refuses_a_mismatched_run(tmp_path: Path) -> None:
    atoms = [Atom(atom_id=f"a{index}", source_id=f"s{index}", text=f"text {index}") for index in range(3)]
    queries = [
        Query(
            query_id=f"q{index}",
            text="t",
            modality="pdf",
            gold=frozenset({"s0"}),
            holder="h",
            answerable=True,
            as_of=None,
        )
        for index in range(2)
    ]
    mark = fingerprint(["bm25_only"], queries, atoms)
    path = tmp_path / "rankings.jsonl"
    assert load_rankings(path, mark) == {}
    rows = [json.dumps({"fingerprint": mark}), json.dumps({"query_id": "q0", "arms": {"bm25_only": {"ranked": ["s0"]}}})]
    path.write_text("\n".join(rows) + "\n", encoding="utf-8")
    assert load_rankings(path, mark) == {"q0": {"bm25_only": {"ranked": ["s0"]}}}
    with pytest.raises(SystemExit, match="different run"):
        load_rankings(path, "other")
    path.write_text("\n".join([rows[0], "{not json"]) + "\n", encoding="utf-8")
    with pytest.raises(SystemExit, match="not JSON"):
        load_rankings(path, mark)


def test_latency_needs_samples() -> None:
    assert latency([0.5, 1.0])["p50"] == pytest.approx(0.75)
    with pytest.raises(SystemExit, match="no latency samples"):
        latency([])


def test_every_arm_declares_the_signals_it_consumes() -> None:
    assert set(systems.ARM_COMPONENTS) == set(ARMS)


def _probe_set(count: int) -> list[Query]:
    return [
        Query(
            query_id=f"q{index}",
            text="t",
            modality="pdf",
            gold=frozenset({f"s{index}"}),
            holder="h",
            answerable=index % 2 == 0,
            as_of=None,
        )
        for index in range(count)
    ]


def test_scale_probes_ignore_gold_so_every_corpus_size_is_measured() -> None:
    queries = _probe_set(100)
    probes = scale_probes(queries, 10)
    assert len(probes) == 10
    assert all(query.answerable for query in probes)
    assert probes == scale_probes(queries, 10), "the curve must be reproducible across runs"
    assert {query.query_id for query in probes} <= {q.query_id for q in queries if q.answerable}
    with pytest.raises(SystemExit, match="at least one probe"):
        scale_probes(queries, 0)
    with pytest.raises(SystemExit, match="needs answerable"):
        scale_probes([], 10)


def _overall() -> dict[str, float]:
    return {
        "recall_at_1": 0.3,
        "recall_at_5": 0.4,
        "recall_at_10": 0.5,
        "ndcg_at_10": 0.42,
        "mrr": 0.4,
        "hit_at_10": 0.52,
        "hit_at_10_wilson_lower": 0.45,
        "abstention_rate_answerable": 0.0,
        "abstention_rate_unanswerable": 0.0,
        "top_hit_from_holder_unanswerable": 0.2,
    }


def _report() -> dict:
    return {
        "meta": {
            "generated_at": "2026-01-01T00:00:00+00:00",
            "git_commit": "a" * 40,
            "corpus_sources": 10,
            "queries": 16,
            "answerable": 8,
            "machine": {
                "system": "Linux",
                "release": "x",
                "cpus": 32,
                "threads_per_query": 4,
                "load_average_1_5_15": [1.0, 2.0, 3.0],
            },
            "models": {"dense": "d", "rerank": "r"},
        },
        "chance": {"recall_at_1": 0.0, "recall_at_5": 0.0, "recall_at_10": 0.0},
        "arms": [
            {
                "arm": "bm25_only",
                "overall": _overall(),
                "by_modality": {"pdf": _overall()},
                "latency_seconds": {"p50": 0.1, "p95": 0.2},
            }
        ],
        "live_product": {
            "queries": 2,
            "top_k": 10,
            "recall_at_10": 0.87,
            "mrr": 0.74,
            "hit_at_10": 0.87,
            "hit_at_10_wilson_lower": 0.82,
            "p50": 1.0,
            "p95": 1.0,
            "mean": 1.0,
            "max": 1.0,
            "top10_overlap_with_bm25_only": 0.0,
            "identical_top10": 0.0,
            "distinct_sources_in_top10": {"mean": 7.5, "min": 1, "max": 10},
        },
        "scale": [
            {"sources": 500, "atoms": 900, "queries": 4, "note": "latency only", "bm25_only": {"p50": 0.001, "p95": 0.002}}
        ],
    }


def test_the_report_writer_survives_a_scale_note(tmp_path: Path) -> None:
    path = tmp_path / "BASELINE.md"
    write_markdown(_report(), path)
    text = path.read_text(encoding="utf-8")
    assert "| 500 | 900 | 0.001/0.002 |" in text
    assert "_latency only_" in text
    assert "load average 1.0 on 32 cores" in text
    assert "## unanswerable queries (50.0% of the suite)" in text
    assert "recall@10 0.870" in text


def test_scale_arm_names_reads_only_latency_dicts() -> None:
    assert scale_arm_names(_report()["scale"]) == ["bm25_only"]
    with pytest.raises(SystemExit, match="no measured arm"):
        scale_arm_names([{"sources": 500, "note": "latency only"}])


def test_remeasure_latency_replaces_cached_seconds_with_a_warm_sample() -> None:
    arms: list[dict[str, Any]] = [{"arm": "bm25_only", "latency_seconds": {"p50": 99.0}}]
    remeasure_latency(arms, JointRunner(["bm25_only"], ATOMS, systems.MODELS), _probe_set(8), 4)
    measured = arms[0]["latency_seconds"]
    assert measured["p50"] < 99.0
    assert measured["queries"] == 4
    assert "warm models" in measured["note"]