"""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.query_set import Query, scale_probes from bench.report import scale_arm_names, write_markdown from bench.run import fingerprint, latency, load_rankings, remeasure_latency 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"]