From 05c6f110366bc0d7a88a452807503c0a7b2d26e2 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 11:24:20 +0200 Subject: [PATCH 1/5] feat(eval): expose the shortlist funnel loss in TREC metrics recall@k measures the first-level candidate list and nDCG@k measures the ranked list, so neither can see the pipeline's largest loss: the shortlist handed to the eligibility reasoner is far shorter than the candidate list, and a relevant trial dropped there can never be ranked. Measured on the completed runs, the shortlist discards 33% (TREC 2021), 36% (2022) and 61% (2023) of the relevant trials the first level had already found. The 2023 run shows why this stayed hidden: it has the best nDCG@10 of any run (0.881) and the worst shortlist recall (0.254), because the condensed metric only orders the judged trials that survive the funnel. evaluate() now reports, whenever top_trials.txt is present: - shortlist_recall .............. recall of the list the reasoner read - shortlist_size - funnel_depth_loss ............. recall given up by shortening alone - shortlist_selection_delta ..... second-level ordering vs a plain first-level top-N cut at the same depth The split matters for what we fix first. On TREC 2021 depth accounts for 94% of the loss (0.866 -> 0.595 at top-200) and ordering for 6% (-> 0.577). The selection delta is negative on all four runs (-0.032, -0.027, -0.022, -0.009): the criterion reranker plus RRF shortlist fusion currently select worse than taking the first-level top-N and doing nothing. Runs without top_trials.txt still evaluate; the funnel keys report None. Co-Authored-By: Claude Opus 5 --- src/trialmatchai/trec/qrels.py | 62 +++++++++++++++++++++++++++++++++- tests/test_qrels_eval.py | 42 +++++++++++++++++++++++ 2 files changed, 103 insertions(+), 1 deletion(-) diff --git a/src/trialmatchai/trec/qrels.py b/src/trialmatchai/trec/qrels.py index 94d543e1..f681e7d5 100644 --- a/src/trialmatchai/trec/qrels.py +++ b/src/trialmatchai/trec/qrels.py @@ -124,6 +124,18 @@ def _retrieved_for_patient(patient_dir: Path) -> list[str]: return [] +def _shortlist_for_patient(patient_dir: Path) -> list[str]: + """The trials that actually reached the eligibility (CoT) stage. + + This is the funnel's narrowest point: a relevant trial dropped here is + unrecoverable, however good the first-level search was. + """ + shortlist = patient_dir / "top_trials.txt" + if not shortlist.exists(): + return [] + return [line.strip() for line in shortlist.read_text().splitlines() if line.strip()] + + def recall_at_k(retrieved: list[str], relevant: set[str], k: int) -> float | None: if not relevant: return None @@ -168,6 +180,17 @@ def evaluate( Reports recall@k (retrieval, first-level list) and tie-aware nDCG@{5,10,20} + P@10 (ranking, condensed to judged trials). P@10 is split into "relevant" (grade>=1) and "eligible" (grade==2). + + Also reports funnel metrics when ``top_trials.txt`` is present. recall@k + (first-level list) and nDCG@k (ranked list) between them hide the pipeline's + largest loss: the shortlist handed to the reasoner is far shorter than the + candidate list, and a relevant trial dropped there can never be ranked. + + - ``shortlist_recall`` -- recall of the list the reasoner actually read. + - ``funnel_depth_loss`` -- recall given up by shortening the list alone. + - ``shortlist_selection_delta`` -- second-level ordering minus a plain + first-level cut at the same depth. Negative means the second level + selected worse than doing nothing. """ results_dir = Path(results_dir) relevant = relevant_ncts(qrels, threshold=threshold) @@ -176,6 +199,19 @@ def evaluate( rec_sums = {f"recall@{k}": 0.0 for k in cutoffs} rec_counts = {f"recall@{k}": 0 for k in cutoffs} + # Funnel instrumentation: recall@k measures the FIRST-LEVEL list, but only the shortlist + # reaches the reasoner. That gap is an unrecoverable ceiling on every ranking metric and is + # invisible in recall@k. Split it into its two causes: depth (shortlist shorter than the + # candidate list) and selection (second-level ordering vs a plain first-level top-N cut). + funnel_keys = ( + "shortlist_recall", + "shortlist_size", + "first_level_recall_at_shortlist_depth", + "shortlist_selection_delta", + "funnel_depth_loss", + ) + funnel_sums = {key: 0.0 for key in funnel_keys} + funnel_counts = {key: 0 for key in funnel_keys} rank_sums = {f"ndcg@{k}": 0.0 for k in NDCG_CUTOFFS} rank_sums.update({f"ndcg_full@{k}": 0.0 for k in NDCG_CUTOFFS}) rank_sums[f"P@{P_CUTOFF}(rel>=1)"] = 0.0 @@ -202,6 +238,26 @@ def evaluate( rec_sums[f"recall@{k}"] += r rec_counts[f"recall@{k}"] += 1 + shortlist = _shortlist_for_patient(patient_dir) + if shortlist and retrieved: + depth = len(shortlist) + short_r = recall_at_k(shortlist, rel_set, depth) + first_r = recall_at_k(retrieved, rel_set, depth) + full_r = recall_at_k(retrieved, rel_set, len(retrieved)) + funnel = { + "shortlist_recall": short_r, + "shortlist_size": float(depth), + "first_level_recall_at_shortlist_depth": first_r, + # > 0 means the second level beat a plain first-level cut at the same depth. + "shortlist_selection_delta": short_r - first_r, + # Recall thrown away purely by shortening the list. + "funnel_depth_loss": full_r - first_r, + } + row.update(funnel) + for key, value in funnel.items(): + funnel_sums[key] += value + funnel_counts[key] += 1 + if ranked: # Two IDCG bases: ndcg@k normalizes by the ideal over judged-AND-ranked trials # (recall-independent ordering quality); ndcg_full@k by the ideal over the FULL judged @@ -233,7 +289,11 @@ def evaluate( rank_counts[f"graded_P@{P_CUTOFF}"] += 1 per_query[query_id] = row - mean = {**_mean(rec_sums, rec_counts), **_mean(rank_sums, rank_counts)} + mean = { + **_mean(rec_sums, rec_counts), + **_mean(funnel_sums, funnel_counts), + **_mean(rank_sums, rank_counts), + } return { "recall_relevance_threshold": threshold, "num_queries_scored": len(per_query), diff --git a/tests/test_qrels_eval.py b/tests/test_qrels_eval.py index a864fd73..5520a777 100644 --- a/tests/test_qrels_eval.py +++ b/tests/test_qrels_eval.py @@ -93,3 +93,45 @@ def test_evaluate_precision_is_condensed_to_judged_pool(tmp_path): assert mean["P@10(rel>=1)"] == pytest.approx(2 / 10) # raw would be 0/10 assert mean["P@10(eligible)"] == pytest.approx(1 / 10) # only NCT1 is grade 2 assert mean["graded_P@10"] == pytest.approx((2 + 1) / (10 * 2)) # raw would be 0 + + +def test_evaluate_reports_funnel_metrics(tmp_path): + """The shortlist (top_trials.txt) is the funnel's narrowest point. Evaluation must expose + its recall and split the loss into depth vs second-level selection, since recall@k + (first-level list) and nDCG@k (ranked list) both hide it.""" + q = "trec-1" + pdir = tmp_path / q + pdir.mkdir() + # First level finds all three relevant trials; the shortlist keeps only two slots and + # spends one on an irrelevant trial that a plain first-level top-2 would not have picked. + (pdir / "nct_ids.txt").write_text("NCT1\nNCT2\nNCT3\nNCT4\n") + (pdir / "top_trials.txt").write_text("NCT1\nNCT4\n") + (pdir / "ranked_trials.json").write_text( + json.dumps({"RankedTrials": [{"TrialID": "NCT1", "Score": 1.0}]}) + ) + qrels = {q: {"NCT1": 2, "NCT2": 2, "NCT3": 1, "NCT4": 0}} + + mean = evaluate(qrels, tmp_path, cutoffs=(10,))["mean"] + + assert mean["recall@10"] == 1.0 # first level found everything... + assert mean["shortlist_recall"] == pytest.approx(1 / 3) # ...the reasoner saw a third + assert mean["shortlist_size"] == 2 + assert mean["first_level_recall_at_shortlist_depth"] == pytest.approx(2 / 3) + assert mean["shortlist_selection_delta"] == pytest.approx(-1 / 3) # selection hurt + assert mean["funnel_depth_loss"] == pytest.approx(1 / 3) # depth alone cost this + + +def test_evaluate_funnel_metrics_absent_without_shortlist(tmp_path): + """Runs predating the shortlist file must still evaluate, without funnel values.""" + q = "trec-1" + pdir = tmp_path / q + pdir.mkdir() + (pdir / "nct_ids.txt").write_text("NCT1\n") + (pdir / "ranked_trials.json").write_text( + json.dumps({"RankedTrials": [{"TrialID": "NCT1", "Score": 1.0}]}) + ) + + mean = evaluate({q: {"NCT1": 2}}, tmp_path, cutoffs=(10,))["mean"] + + assert mean["recall@10"] == 1.0 + assert mean["shortlist_recall"] is None From 91f724f33587952638139de61a4d3e2beb06ea86 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 14:21:14 +0200 Subject: [PATCH 2/5] feat(search): per-patient shortlist depth from the first-level score curve The shortlist divisor sizes every patient identically, but depth is 94% of the measured shortlist recall loss (PR #30) and no single number fits: the depth a patient needs to reach 90% of its own first-level recall ranges from 50 to 1550 trials, spread evenly across that range. Sizing for the worst case wastes ~65% of the reasoner's compute; sizing for the median drops the hard patients. Adds search.shortlist.policy: fixed existing divisor sizing, unchanged -- still the DEFAULT relative_to_max keep trials scoring >= alpha x this patient's top score A peaked first-level score curve means retrieval was confident and few trials are plausible; a flat curve means many are. The cut is relative to the patient's own maximum because RRF scores are only comparable within one patient. Tuned offline by replaying the completed runs' first_level_scores.json, so no GPU time. Compared against fixed depth AT EQUAL MEAN DEPTH, which is the only fair test -- a policy must spend compute better, not merely spend more: TREC 2021 +0.017 to +0.029 recall (peak at mean depth ~305) TREC 2022 +0.007 to +0.025 recall TREC 2023 -0.010 to +0.006 recall -- questionnaire topics, no gain Two other policies were tried and rejected: cumulative score mass tracked fixed depth to within +-0.003 everywhere, and half-max curve width won only below ~350 mean depth and lost above it. At alpha=0.25 the policy also chooses to go deeper (196->305 trials on 2021), which lifts shortlist recall 0.609->0.710. That part is bought with compute, not won for free; the equal-cost gain above is the free part. Both are real and they should be reported separately. Default stays "fixed" so enabling this is an explicit A/B. Degrades to the fixed depth when first-level scores are absent (resumed runs), never exceeds what the reasoner can consume, and writes shortlist_depth.json recording the decision and the depth the old sizing would have picked. Co-Authored-By: Claude Opus 5 --- src/trialmatchai/config/settings.py | 17 +++ src/trialmatchai/main.py | 29 +++- src/trialmatchai/matching/shortlist_depth.py | 149 +++++++++++++++++++ tests/test_shortlist_depth.py | 149 +++++++++++++++++++ 4 files changed, 341 insertions(+), 3 deletions(-) create mode 100644 src/trialmatchai/matching/shortlist_depth.py create mode 100644 tests/test_shortlist_depth.py diff --git a/src/trialmatchai/config/settings.py b/src/trialmatchai/config/settings.py index ed3e9b24..f4ab4b13 100644 --- a/src/trialmatchai/config/settings.py +++ b/src/trialmatchai/config/settings.py @@ -174,6 +174,22 @@ class FirstLevelSearchSettings(BaseModel): ) +class ShortlistSettings(BaseModel): + """How deep the shortlist handed to the eligibility reasoner goes. + + "fixed" keeps the divisor-based sizing (one depth for every patient) and is the + default. "relative_to_max" sizes each patient from its own first-level score curve, + keeping trials scoring at least ``relative_to_max_alpha`` x that patient's top score. + See matching/shortlist_depth.py for the offline evidence. + """ + + policy: Literal["fixed", "relative_to_max"] = "fixed" + relative_to_max_alpha: float = Field(0.25, gt=0.0, le=1.0) + min_depth: int = Field(50, ge=1) + # None -> bounded only by what the reasoner can consume (rag.max_trials_rag). + max_depth: int | None = Field(None, ge=1) + + class SearchSettings(BaseModel): mode: Literal["bm25", "vector", "hybrid"] = "hybrid" vector_score_threshold: float = Field(0.5, ge=0.0, le=1.0) @@ -191,6 +207,7 @@ class SearchSettings(BaseModel): first_level: FirstLevelSearchSettings = Field( default_factory=FirstLevelSearchSettings ) + shortlist: ShortlistSettings = Field(default_factory=ShortlistSettings) @model_validator(mode="before") @classmethod diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 83c8df8c..97958d37 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -15,6 +15,7 @@ rank_trials, save_ranked_trials, ) +from trialmatchai.matching.shortlist_depth import choose_shortlist_depth, depth_report from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever from trialmatchai.matching.retrieval.location import ( @@ -307,15 +308,16 @@ def run_second_level_search( trial["nct_id"]: trial["score"] for trial in second_level_results } + search_config = config.get("search", {}) combined_scores = _fuse_shortlist_scores( nct_ids=nct_ids, second_level_results=second_level_results, first_level_scores=first_level_scores, - search_config=config.get("search", {}), + search_config=search_config, ) sorted_trials = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True) - keep_divisor = max(1, int(config.get("search", {}).get("second_level_keep_divisor", 3))) + keep_divisor = max(1, int(search_config.get("second_level_keep_divisor", 3))) # Size the shortlist off the reranked count, not the fused pool: rank fusion adds # first-level-only trials, so keying the divisor to the pool would silently enlarge the # shortlist and confound a fusion A/B. @@ -323,12 +325,33 @@ def run_second_level_search( num_top = max(1, min(reranked_count // keep_divisor, top_n)) # RAG only reasons over rag.max_trials_rag trials; cap the shortlist to match, else trials # past the cap get no eligibility output and are silently dropped from the final ranking. + upper_bound = len(sorted_trials) or 1 if _rag_enabled(config): - num_top = max(1, min(num_top, int(config.get("rag", {}).get("max_trials_rag", 20)))) + upper_bound = min(upper_bound, int(config.get("rag", {}).get("max_trials_rag", 20))) + num_top = max(1, min(num_top, upper_bound)) + # The divisor sizes every patient the same. Depth is 94% of the measured shortlist recall + # loss and the depth patients need spans 50-1550 trials, so an opt-in policy may widen or + # narrow this per patient from the first-level score curve. Default policy returns num_top. + fixed_depth = num_top + num_top = choose_shortlist_depth( + first_level_scores=first_level_scores, + fixed_depth=fixed_depth, + upper_bound=upper_bound, + search_config=search_config, + ) semi_final_trials = sorted_trials[:num_top] top_trials_path = f"{output_folder}/top_trials.txt" write_text_file([trial_id for trial_id, _ in semi_final_trials], top_trials_path) + write_json_file( + depth_report( + chosen=len(semi_final_trials), + fixed_depth=fixed_depth, + first_level_scores=first_level_scores, + search_config=search_config, + ), + f"{output_folder}/shortlist_depth.json", + ) constraints_config = config.get("constraints", {}) if constraints_config.get("enabled", True) and constraints_config.get( "write_reports", diff --git a/src/trialmatchai/matching/shortlist_depth.py b/src/trialmatchai/matching/shortlist_depth.py new file mode 100644 index 00000000..21c85dbf --- /dev/null +++ b/src/trialmatchai/matching/shortlist_depth.py @@ -0,0 +1,149 @@ +"""How many trials the eligibility reasoner gets to read. + +The shortlist is the pipeline's narrowest point: a relevant trial dropped here can +never be ranked, however good the reasoning is. Measured on the completed TREC runs, +a fixed-size shortlist discards a third of the relevant trials the first level had +already found (``shortlist_recall`` in ``trec/qrels.py``). + +Depth is the dominant cause -- 94% of that loss on TREC 2021 -- and no single number +serves every patient: the depth needed to reach 90% of a patient's own first-level +recall ranges from 50 to 1550 trials, spread evenly across that range. Sizing for the +worst case wastes ~65% of the reasoner's compute; sizing for the median silently drops +the hard patients. + +``relative_to_max`` therefore reads the shape of the first-level score curve. A peaked +curve means retrieval was confident and few trials are plausible; a flat curve means +many are, and the patient needs more depth. Offline replay over the completed runs +(first-level scores, same mean depth as the fixed policy) gives: + + TREC 2021 +0.017 to +0.029 recall + TREC 2022 +0.007 to +0.025 recall + TREC 2023 -0.010 to +0.006 recall (questionnaire topics; no gain) + +So it is a small, free gain on the narrative-topic tracks and a wash on 2023 -- worth +having, but not a substitute for spending more depth outright. ``fixed`` remains the +default so enabling this is an explicit A/B, not a silent change. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +from trialmatchai.utils.logging_config import setup_logging + +logger = setup_logging(__name__) + +POLICIES = ("fixed", "relative_to_max") +DEFAULT_ALPHA = 0.25 +DEFAULT_MIN_DEPTH = 50 + + +def shortlist_config(search_config: Mapping[str, Any] | None) -> dict[str, Any]: + """Resolve the ``search.shortlist`` block, tolerating absent or partial config.""" + raw = (search_config or {}).get("shortlist") or {} + if not isinstance(raw, Mapping): + raw = {} + policy = str(raw.get("policy", "fixed") or "fixed") + if policy not in POLICIES: + logger.warning( + "Unknown search.shortlist.policy %r; falling back to 'fixed'. Known: %s", + policy, + ", ".join(POLICIES), + ) + policy = "fixed" + return { + "policy": policy, + "relative_to_max_alpha": float(raw.get("relative_to_max_alpha", DEFAULT_ALPHA)), + "min_depth": int(raw.get("min_depth", DEFAULT_MIN_DEPTH)), + "max_depth": raw.get("max_depth"), + } + + +def _relative_to_max_depth(scores: list[float], alpha: float) -> int: + """Count of trials scoring at least ``alpha`` x the top score. + + Scores are a weighted RRF sum, so they are positive and comparable only within one + patient -- which is exactly why the cut is relative to that patient's own maximum + rather than an absolute threshold. + """ + if not scores: + return 0 + top = scores[0] + if top <= 0: + return len(scores) + cut = alpha * top + kept = 0 + for score in scores: + if score < cut: + break + kept += 1 + return kept + + +def choose_shortlist_depth( + *, + first_level_scores: Mapping[str, float] | None, + fixed_depth: int, + upper_bound: int, + search_config: Mapping[str, Any] | None = None, +) -> int: + """Shortlist size for one patient. + + ``fixed_depth`` is what the existing divisor-based sizing would have chosen, and is + returned unchanged under the default policy. ``upper_bound`` is the hard ceiling the + caller can honour (the reasoner's own cap), and is never exceeded. + """ + upper_bound = max(1, int(upper_bound)) + fixed_depth = max(1, min(int(fixed_depth), upper_bound)) + cfg = shortlist_config(search_config) + if cfg["policy"] == "fixed": + return fixed_depth + + scores = sorted((float(v) for v in (first_level_scores or {}).values()), reverse=True) + if not scores: + # No first-level signal to read (e.g. a resumed run missing the scores file): + # degrade to the fixed sizing rather than guessing a depth. + logger.warning( + "shortlist policy 'relative_to_max' has no first-level scores; using fixed depth %s", + fixed_depth, + ) + return fixed_depth + + depth = _relative_to_max_depth(scores, cfg["relative_to_max_alpha"]) + floor = max(1, cfg["min_depth"]) + ceiling = upper_bound + configured_max = cfg["max_depth"] + if configured_max is not None: + ceiling = min(ceiling, max(1, int(configured_max))) + depth = max(floor, min(depth, ceiling)) + logger.info( + "Shortlist depth %s (policy=relative_to_max, alpha=%.3g, fixed would be %s)", + depth, + cfg["relative_to_max_alpha"], + fixed_depth, + ) + return depth + + +def depth_report( + *, + chosen: int, + fixed_depth: int, + first_level_scores: Mapping[str, float] | None, + search_config: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Provenance for the depth decision, written beside the shortlist.""" + cfg = shortlist_config(search_config) + scores: Iterable[float] = (first_level_scores or {}).values() + ordered = sorted((float(v) for v in scores), reverse=True) + return { + "policy": cfg["policy"], + "chosen_depth": int(chosen), + "fixed_depth": int(fixed_depth), + "relative_to_max_alpha": cfg["relative_to_max_alpha"], + "min_depth": cfg["min_depth"], + "max_depth": cfg["max_depth"], + "candidate_pool": len(ordered), + "top_score": ordered[0] if ordered else None, + } diff --git a/tests/test_shortlist_depth.py b/tests/test_shortlist_depth.py new file mode 100644 index 00000000..dd4cef6c --- /dev/null +++ b/tests/test_shortlist_depth.py @@ -0,0 +1,149 @@ +"""Shortlist depth policies (matching/shortlist_depth.py). + +The shortlist is where the pipeline loses most of its relevant trials, so the depth +decision must be explicit, bounded, and never silently change under the default config. +""" + +import pytest + +from trialmatchai.matching.shortlist_depth import ( + choose_shortlist_depth, + depth_report, + shortlist_config, +) + + +def _scores(values): + return {f"NCT{i:04d}": v for i, v in enumerate(values)} + + +def test_default_policy_returns_the_fixed_depth_unchanged(): + """Enabling adaptive depth must be an explicit A/B, never a silent behaviour change.""" + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 500), + fixed_depth=196, + upper_bound=300, + search_config={}, + ) + assert depth == 196 + + +def test_relative_to_max_keeps_trials_above_alpha_times_the_top_score(): + # Top score 1.0, alpha 0.25 -> keep while score >= 0.25, so the first four. + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0, 0.8, 0.4, 0.25, 0.2, 0.1]), + fixed_depth=2, + upper_bound=100, + search_config={"shortlist": {"policy": "relative_to_max", "min_depth": 1}}, + ) + assert depth == 4 + + +def test_peaked_curve_gets_less_depth_than_flat_curve(): + """The whole point: a confident retrieval needs fewer trials than an ambiguous one.""" + cfg = {"shortlist": {"policy": "relative_to_max", "min_depth": 1}} + peaked = choose_shortlist_depth( + first_level_scores=_scores([1.0] + [0.01] * 99), + fixed_depth=50, + upper_bound=100, + search_config=cfg, + ) + flat = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 100), + fixed_depth=50, + upper_bound=100, + search_config=cfg, + ) + assert peaked == 1 + assert flat == 100 + assert peaked < flat + + +def test_depth_is_clamped_by_min_depth_and_upper_bound(): + cfg = {"shortlist": {"policy": "relative_to_max", "min_depth": 20, "max_depth": 40}} + # A single dominant trial would give depth 1; the floor lifts it to min_depth. + assert ( + choose_shortlist_depth( + first_level_scores=_scores([1.0] + [0.001] * 99), + fixed_depth=10, + upper_bound=100, + search_config=cfg, + ) + == 20 + ) + # A flat curve would give 100; max_depth caps it at 40. + assert ( + choose_shortlist_depth( + first_level_scores=_scores([1.0] * 100), + fixed_depth=10, + upper_bound=100, + search_config=cfg, + ) + == 40 + ) + + +def test_upper_bound_always_wins_over_configured_max_depth(): + """upper_bound is what the reasoner can actually consume; exceeding it drops trials + silently from the final ranking.""" + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 500), + fixed_depth=10, + upper_bound=30, + search_config={ + "shortlist": {"policy": "relative_to_max", "min_depth": 1, "max_depth": 400} + }, + ) + assert depth == 30 + + +def test_missing_scores_degrade_to_fixed_depth(): + """A resumed run can lack first_level_scores.json; guessing a depth would be worse.""" + for scores in (None, {}): + depth = choose_shortlist_depth( + first_level_scores=scores, + fixed_depth=77, + upper_bound=300, + search_config={"shortlist": {"policy": "relative_to_max"}}, + ) + assert depth == 77 + + +def test_nonpositive_top_score_keeps_the_whole_pool(): + depth = choose_shortlist_depth( + first_level_scores=_scores([0.0, 0.0, 0.0]), + fixed_depth=1, + upper_bound=100, + search_config={"shortlist": {"policy": "relative_to_max", "min_depth": 1}}, + ) + assert depth == 3 + + +def test_unknown_policy_falls_back_to_fixed(caplog): + assert shortlist_config({"shortlist": {"policy": "wishful"}})["policy"] == "fixed" + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 100), + fixed_depth=12, + upper_bound=100, + search_config={"shortlist": {"policy": "wishful"}}, + ) + assert depth == 12 + + +@pytest.mark.parametrize("search_config", [None, {}, {"shortlist": None}]) +def test_absent_config_is_tolerated(search_config): + assert shortlist_config(search_config)["policy"] == "fixed" + + +def test_depth_report_records_the_decision(): + report = depth_report( + chosen=40, + fixed_depth=196, + first_level_scores=_scores([1.0, 0.5, 0.2]), + search_config={"shortlist": {"policy": "relative_to_max"}}, + ) + assert report["policy"] == "relative_to_max" + assert report["chosen_depth"] == 40 + assert report["fixed_depth"] == 196 # what the old sizing would have used + assert report["candidate_pool"] == 3 + assert report["top_score"] == 1.0 From a3b43cd4a139b4103e09ae1bb9fb40d319b0018f Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 14:31:24 +0200 Subject: [PATCH 3/5] feat(trec): comparison harness for the shortlist A/B, and persist second-level scores The depth policy added in 91f724f is tuned only against replayed first-level scores. Whether the extra depth reaches nDCG@10 and P@10 -- rather than just shortlist recall -- needs a GPU run. This adds the reporting half of that experiment. compare_shortlist_ab.py reads the arms of a shortlist-policy A/B and always prints shortlist_size beside the quality metrics, because a depth policy produces two effects that must not be conflated: spending compute BETTER same mean shortlist size, more recall -- a free gain spending compute MORE bigger shortlist, bought with GPU time Reporting only the second reads as a much larger win than it is. The planned arms are fixed (~196 trials/patient), relative_to_max alpha 0.25 (~305, deeper), and relative_to_max alpha 0.33 (~196, equal cost). The equal-cost arm is the one that actually tests the policy. The job script itself stays untracked: *.slurm is gitignored repo-wide and none of the 26 existing job scripts are committed, since they encode site-specific paths, partitions and staging. Also persists second_level_scores.json for the whole second-level pool, not just the shortlist. With first_level_scores.json this makes shortlist fusion replayable offline, so the negative shortlist_selection_delta measured in 05c6f11 -- the fused shortlist selecting worse than a plain first-level cut at the same depth -- can be investigated without spending a GPU job on each candidate fix. --- scripts/compare_shortlist_ab.py | 100 ++++++++++++++++++++++++++++++++ src/trialmatchai/main.py | 6 ++ 2 files changed, 106 insertions(+) create mode 100644 scripts/compare_shortlist_ab.py diff --git a/scripts/compare_shortlist_ab.py b/scripts/compare_shortlist_ab.py new file mode 100644 index 00000000..9718338a --- /dev/null +++ b/scripts/compare_shortlist_ab.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python +"""Compare the arms of the shortlist-policy A/B (scripts/run_trec_shortlist_ab.slurm). + +Reports each arm against the fixed baseline, and -- crucially -- separates the two +effects that a depth policy produces: + + * spending compute BETTER (equal-cost arm: same mean shortlist size, more recall) + * spending compute MORE (deeper arm: bigger shortlist, bought with GPU time) + +Conflating them overstates the result, so cost is always printed beside quality. + + uv run python scripts/compare_shortlist_ab.py [--root shortlist_ab] [--track 21] +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +BASELINE = "fixed" +QUALITY = ( + "shortlist_recall", + "recall@1000", + "ndcg@10", + "P@10(rel>=1)", + "P@10(eligible)", +) +COST = ("shortlist_size",) +DIAGNOSTIC = ("funnel_depth_loss", "shortlist_selection_delta") + + +def load_arm(root: Path, arm: str, track: str) -> dict | None: + path = root / arm / f"results_trec{track}" / "evaluation_metrics.json" + if not path.exists(): + return None + data = json.loads(path.read_text()) + return {"mean": data.get("mean", {}), "n": data.get("num_queries_scored")} + + +def fmt(value: object, width: int = 9) -> str: + if isinstance(value, (int, float)): + return f"{value:{width}.4f}" + return f"{'--':>{width}}" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", default="shortlist_ab") + parser.add_argument("--track", default="21") + args = parser.parse_args() + + root = Path(args.root) + arms = sorted(p.name for p in root.iterdir() if p.is_dir()) if root.is_dir() else [] + if BASELINE not in arms: + print(f"No '{BASELINE}' arm under {root}/ -- nothing to compare against.") + return 1 + + loaded = {arm: load_arm(root, arm, args.track) for arm in arms} + missing = [arm for arm, data in loaded.items() if data is None] + for arm in missing: + print(f"NOTE: arm '{arm}' has no evaluation_metrics.json yet (still running?)") + loaded.pop(arm) + if BASELINE not in loaded: + return 1 + + base = loaded[BASELINE]["mean"] + order = [BASELINE] + [a for a in loaded if a != BASELINE] + + print(f"\nTREC {args.track} shortlist-policy A/B (n={loaded[BASELINE]['n']} topics)\n") + label_w = max(len(a) for a in order) + 2 + for group, keys in (("COST", COST), ("QUALITY", QUALITY), ("DIAGNOSTIC", DIAGNOSTIC)): + print(f" {group}") + print(f" {'arm':<{label_w}}" + "".join(f"{k:>26s}" for k in keys)) + for arm in order: + mean = loaded[arm]["mean"] + cells = "" + for key in keys: + value = mean.get(key) + cell = fmt(value) + if arm != BASELINE and isinstance(value, (int, float)): + ref = base.get(key) + if isinstance(ref, (int, float)): + cell += f" ({value - ref:+.4f})" + cells += f"{cell:>26s}" + print(f" {arm:<{label_w}}" + cells) + print() + + print(" Reading this:") + print(" An arm at the SAME shortlist_size as 'fixed' with higher shortlist_recall") + print(" spends its compute better -- that gain is free.") + print(" An arm with a LARGER shortlist_size bought its gain with GPU time; compare") + print(" it against 'fixed' only after noting the extra cost.") + print(" shortlist_selection_delta < 0 means the second level is still selecting worse") + print(" than a plain first-level cut at that depth, independent of the policy.\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 97958d37..336b42a2 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -307,6 +307,12 @@ def run_second_level_search( second_level_scores = { trial["nct_id"]: trial["score"] for trial in second_level_results } + # Persist the whole second-level pool, not just the shortlist. Together with + # first_level_scores.json this makes shortlist fusion replayable offline, so fusion + # weights can be retuned against completed runs instead of costing a GPU job each time. + # Measured motivation: shortlist_selection_delta is negative on every run so far, i.e. + # the fused shortlist selects worse than a plain first-level cut at the same depth. + write_json_file(second_level_scores, f"{output_folder}/second_level_scores.json") search_config = config.get("search", {}) combined_scores = _fuse_shortlist_scores( From 3ce7d606b4008ddfa1235d327389c3c923c5f67f Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 14:54:36 +0200 Subject: [PATCH 4/5] feat(search): implement the first-level LLM query expansion backend The llm_expansion search channel was dead. LLMQueryExpansionBackend (the Protocol), parse_llm_query_expansion (the parser), LLMQueryExpansion (the schema), the channel's 0.5 weight and search.first_level.llm_expansion_enabled all existed, but nothing ever built a backend and main.py never passed one to ClinicalTrialSearch. Setting the flag logged "no expander is enabled but no expander is configured" and contributed no terms. Adds FirstLevelQueryExpander, which subclasses QueryExpander to reuse its engine, chat template and structured-output machinery -- so it shares the one cached vLLM engine rather than loading a second copy -- and swaps in a retrieval-query prompt and schema. This is a genuinely different task from the existing expander, not a rename. QueryExpander enriches the patient SUMMARY (conditions plus narrative sentences) for the expand stage. This one writes RETRIEVAL QUERIES: short noun phrases that should match a trial's title, condition list or eligibility text, bucketed into the six fields the planner turns into weighted query channels. To support both, QueryExpander's prompt and schema became overridable class attributes; the base class's values are unchanged. The per-field maxItems caps are deliberately uneven. llm_max_terms is a SHARED budget across the five query fields, spent in field order, so a model that pads primary_queries starves biomarker_queries and treatment_queries entirely -- a test pins that behaviour. primary_queries is therefore capped at 3 (a patient has one main disease), leaving the budget for the later fields. Off by default: the builder returns None unless llm_expansion_enabled is set, and returns None rather than raising if the expander cannot be constructed. A failed expansion degrades to empty and is logged, because this is one channel of eight -- losing it should cost recall, not the run. Wired once per run in main_pipeline, not per patient, so the engine is resolved once. --- src/trialmatchai/main.py | 12 ++ src/trialmatchai/matching/query_expansion.py | 169 ++++++++++++++++- tests/test_first_level_expansion.py | 182 +++++++++++++++++++ 3 files changed, 361 insertions(+), 2 deletions(-) create mode 100644 tests/test_first_level_expansion.py diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 336b42a2..69d03a79 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -15,6 +15,7 @@ rank_trials, save_ranked_trials, ) +from trialmatchai.matching.query_expansion import build_first_level_expander from trialmatchai.matching.shortlist_depth import choose_shortlist_depth, depth_report from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever @@ -108,6 +109,7 @@ def run_first_level_search( config: Dict, search_backend, patient_profile: PatientProfile | None = None, + llm_query_expander=None, ) -> Optional[Tuple]: main_conditions = list(keywords.get("main_conditions", [])) other_conditions = list(keywords.get("other_conditions", [])) @@ -126,6 +128,9 @@ def run_first_level_search( search_backend=search_backend, embedder=embedder, entity_annotator=entity_annotator, + # Without this the llm_expansion channel is dead: the planner logs "no expander is + # configured" and returns [], however the config flag is set. + llm_query_expander=llm_query_expander, ) search_cfg = config["search"] @@ -553,6 +558,12 @@ def main_pipeline( embedder = build_embedder(config) entity_annotator = build_entity_annotator(config, embedder=embedder) + # Built once for the whole run, not per patient: it shares the cached CoT engine, and + # rebuilding per patient would re-resolve that engine 75 times. None when + # search.first_level.llm_expansion_enabled is off, which is the default. + llm_query_expander = build_first_level_expander(config) + if llm_query_expander is not None: + logger.info("First-level LLM query expansion is ON (llm_expansion channel active).") with warnings.catch_warnings(): warnings.filterwarnings( @@ -647,6 +658,7 @@ def main_pipeline( config, search_backend, patient_profile=profile, + llm_query_expander=llm_query_expander, ) if not result: logger.error("First-level search failed for %s", patient_id) diff --git a/src/trialmatchai/matching/query_expansion.py b/src/trialmatchai/matching/query_expansion.py index 5ffbcd00..3d36c1b6 100644 --- a/src/trialmatchai/matching/query_expansion.py +++ b/src/trialmatchai/matching/query_expansion.py @@ -105,6 +105,11 @@ def _resolve_settings(config: Dict[str, Any]) -> Dict[str, Any]: class QueryExpander: """CoT expander; loads its model lazily so import stays base-deps safe.""" + # Overridable by subclasses that reuse this engine/template machinery for a different + # extraction task (see FirstLevelQueryExpander). + system_prompt: str = SYSTEM_PROMPT + json_schema: Dict[str, Any] = _KEYWORDS_JSON_SCHEMA + def __init__(self, settings: Dict[str, Any], config: Dict[str, Any]): self.settings = settings self.config = config @@ -160,7 +165,7 @@ def _generate(self, narrative: str) -> str: no_think = bool(self.settings.get("no_think")) user_content = ("/no_think\n" + narrative) if no_think else narrative messages = [ - {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "system", "content": self.system_prompt}, {"role": "user", "content": user_content}, ] # Qwen3.x-style templates take enable_thinking; harmless-and-ignored elsewhere (guarded). @@ -199,7 +204,7 @@ def _apply_template(tokenize): if self.settings.get("guided_json"): from vllm.sampling_params import StructuredOutputsParams # type: ignore - structured = StructuredOutputsParams(json=_KEYWORDS_JSON_SCHEMA, disable_any_whitespace=True) + structured = StructuredOutputsParams(json=self.json_schema, disable_any_whitespace=True) params = SamplingParams( temperature=0.0, max_tokens=self.settings["max_new_tokens"], @@ -257,3 +262,163 @@ def enrich_summary( if sentences: out["patient_narrative"] = sentences return out + + +# --- first-level retrieval query expansion (the llm_expansion search channel) ------------ # + +# Distinct from SYSTEM_PROMPT above. That one enriches the patient SUMMARY (conditions and +# narrative sentences). This one writes RETRIEVAL QUERIES: short noun phrases that should +# match trial titles, conditions and eligibility text. The two are not interchangeable -- +# the first-level planner buckets these six fields into weighted query channels. +FIRST_LEVEL_SYSTEM_PROMPT = """ +You expand a patient description into search queries for a clinical trial index. + +Write SHORT NOUN PHRASES that would appear in a trial's title, condition list or +eligibility criteria. Do not write sentences, questions or explanations. + +Fill these six fields: + +1. "primary_queries": the patient's main disease as a trial would name it. Include the + staging or subtype only when the patient description states it. +2. "disease_aliases": other names for that same disease -- synonyms, abbreviations, older + or regional terminology, and the expanded form of any abbreviation. +3. "broader_queries": the parent disease categories a trial might recruit under, from + narrower to wider. These deliberately trade precision for coverage. +4. "biomarker_queries": genes, mutations, fusions, receptor and expression status, and + other molecular markers stated for this patient. +5. "treatment_queries": drugs, drug classes, procedures and prior therapies stated for + this patient. +6. "discarded_or_uncertain": terms you considered but rejected, and anything you are not + confident the patient description supports. + +Rules: +- Use ONLY what the patient description states. Never infer a diagnosis, stage, biomarker + or therapy that is not written there. Put anything doubtful in "discarded_or_uncertain". +- Leave a field as an empty list when the description supports nothing for it. An empty + list is correct; an invented term is not. +- No duplicates within a field. + +Return a JSON object with exactly those six keys and no other commentary. +""".strip() + +_FIRST_LEVEL_FIELDS = ( + "primary_queries", + "disease_aliases", + "broader_queries", + "biomarker_queries", + "treatment_queries", + "discarded_or_uncertain", +) + +# maxItems bounds the array COUNT for the same reason as _KEYWORDS_JSON_SCHEMA: it forces a +# verbose model to close each array instead of emitting terms until max_tokens runs out. +# These are noun phrases, so a short maxLength is safe here (unlike expanded_sentences). +# +# The per-field caps are deliberately uneven. search.first_level.llm_max_terms is a SHARED +# budget across the five query fields, spent in field order (first_level_planner +# parse_llm_query_expansion), so a model that pads primary_queries starves the biomarker and +# treatment channels entirely. Capping primary_queries tightly -- a patient has one main +# disease, not twelve -- keeps the budget available for the later fields. +_FIRST_LEVEL_MAX_ITEMS = { + "primary_queries": 3, + "disease_aliases": 8, + "broader_queries": 5, + "biomarker_queries": 8, + "treatment_queries": 8, + "discarded_or_uncertain": 12, +} +_FIRST_LEVEL_JSON_SCHEMA = { + "type": "object", + "properties": { + field: { + "type": "array", + "maxItems": _FIRST_LEVEL_MAX_ITEMS[field], + "items": {"type": "string", "maxLength": 120}, + } + for field in _FIRST_LEVEL_FIELDS + }, + "required": list(_FIRST_LEVEL_FIELDS), +} + +_FIRST_LEVEL_EMPTY: Dict[str, List[str]] = {field: [] for field in _FIRST_LEVEL_FIELDS} + + +def _first_level_patient_text(profile: Any, matching_summary: Dict[str, Any]) -> str: + """Compact patient description for the expander prompt. + + Built from the matching summary rather than the raw profile so it stays in step with + what first-level retrieval actually searches on. + """ + summary = matching_summary or {} + parts: List[str] = [] + main = [c for c in _as_list(summary.get("main_conditions")) if c][:12] + other = [c for c in _as_list(summary.get("other_conditions")) if c][:30] + narrative = [s for s in _as_list(summary.get("patient_narrative")) if s][:12] + if main: + parts.append("Main conditions: " + "; ".join(main)) + if other: + parts.append("Other conditions and factors: " + "; ".join(other)) + age, gender = summary.get("age"), summary.get("gender") + demographics = [ + f"Age: {age}" for _ in (1,) if age not in (None, "", "all") + ] + [f"Sex: {gender}" for _ in (1,) if gender not in (None, "", "all")] + if demographics: + parts.append(", ".join(demographics)) + if narrative: + parts.append("Description: " + " ".join(narrative)) + return "\n".join(parts).strip() + + +class FirstLevelQueryExpander(QueryExpander): + """Implements ``LLMQueryExpansionBackend`` for the first-level ``llm_expansion`` channel. + + Reuses QueryExpander's engine, chat-template and structured-output machinery -- so it + shares the one cached vLLM engine rather than loading a second copy -- but swaps in the + retrieval-query prompt and schema. + """ + + system_prompt = FIRST_LEVEL_SYSTEM_PROMPT + json_schema = _FIRST_LEVEL_JSON_SCHEMA + + def expand_first_level_queries( + self, + *, + profile: Any, + matching_summary: Dict[str, Any], + ) -> Dict[str, Any]: + patient_text = _first_level_patient_text(profile, matching_summary) + if not patient_text: + return dict(_FIRST_LEVEL_EMPTY) + try: + raw = self._generate(patient_text) + parsed = extract_json_object(BaseTrialProcessor._strip_thinking_tags(raw)) + if not isinstance(parsed, dict): + raise ValueError("first-level expansion output was not a JSON object") + return {field: _as_list(parsed.get(field)) for field in _FIRST_LEVEL_FIELDS} + except Exception as exc: + # Retrieval must not fail because expansion did: the channel is one of eight and + # carries weight 0.5, so an empty expansion degrades recall rather than the run. + logger.error( + "First-level query expansion failed; continuing without that channel: %s", exc + ) + return dict(_FIRST_LEVEL_EMPTY) + + +def build_first_level_expander(config: Dict[str, Any]) -> "FirstLevelQueryExpander | None": + """Construct the expander when ``search.first_level.llm_expansion_enabled`` is true. + + Independent of ``query_expansion.enabled``: that flag governs the separate summary + enrichment stage. Both may run, and they share one engine. + """ + first_level = (config.get("search") or {}).get("first_level") or {} + if not first_level.get("llm_expansion_enabled"): + return None + try: + return FirstLevelQueryExpander(_resolve_settings(config), config) + except Exception as exc: + logger.error( + "search.first_level.llm_expansion_enabled is set but the expander could not be " + "built; first-level search continues without that channel: %s", + exc, + ) + return None diff --git a/tests/test_first_level_expansion.py b/tests/test_first_level_expansion.py new file mode 100644 index 00000000..e7a900da --- /dev/null +++ b/tests/test_first_level_expansion.py @@ -0,0 +1,182 @@ +"""First-level LLM query expansion (the llm_expansion search channel). + +This channel was dead before: the protocol, parser, schema and config flag all existed, +but nothing ever constructed a backend, so enabling the flag logged "no expander is +configured" and returned no terms. These tests cover the backend and the config gate. +""" + +import json + +import pytest + +from trialmatchai.matching.query_expansion import ( + _FIRST_LEVEL_FIELDS, + FirstLevelQueryExpander, + _first_level_patient_text, + build_first_level_expander, +) +from trialmatchai.matching.retrieval.first_level_planner import parse_llm_query_expansion + +SUMMARY = { + "main_conditions": ["metastatic breast cancer"], + "other_conditions": ["hypertension", "HER2 positive"], + "patient_narrative": ["A 54 year old woman with metastatic breast cancer."], + "age": 54, + "gender": "female", +} + + +class _FakeExpander(FirstLevelQueryExpander): + """Bypasses __init__ so no model or GPU is touched; _generate returns a canned reply.""" + + def __init__(self, reply): + self._reply = reply + self.settings = {"guided_json": True, "max_new_tokens": 512} + self.config = {} + self.backend = "vllm" + + def _generate(self, narrative): + if isinstance(self._reply, Exception): + raise self._reply + return self._reply + + +def test_expands_into_the_six_planner_fields(): + payload = { + "primary_queries": ["metastatic breast cancer"], + "disease_aliases": ["breast carcinoma", "mammary carcinoma"], + "broader_queries": ["solid tumor"], + "biomarker_queries": ["HER2 positive"], + "treatment_queries": ["trastuzumab"], + "discarded_or_uncertain": ["hypertension"], + } + result = _FakeExpander(json.dumps(payload)).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result == payload + assert set(result) == set(_FIRST_LEVEL_FIELDS) + + +def test_output_is_consumable_by_the_planner_parser(): + """The backend's contract is the planner's parser, not just valid JSON.""" + payload = { + "primary_queries": ["metastatic breast cancer"], + "disease_aliases": ["breast carcinoma"], + "broader_queries": ["solid tumor"], + "biomarker_queries": ["HER2 positive"], + "treatment_queries": ["trastuzumab"], + "discarded_or_uncertain": [], + } + raw = _FakeExpander(json.dumps(payload)).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + parsed = parse_llm_query_expansion(raw, max_terms=12) + assert parsed.primary_queries == ["metastatic breast cancer"] + assert parsed.biomarker_queries == ["HER2 positive"] + + +def test_max_terms_is_a_shared_budget_spent_primary_first(): + """llm_max_terms caps the TOTAL across the five query fields, not each one, and is spent + in field order. A model that fills primary_queries can starve the later channels, so the + prompt must keep primary_queries to the actual disease rather than padding it.""" + payload = {field: [] for field in _FIRST_LEVEL_FIELDS} + payload["primary_queries"] = [f"q{i}" for i in range(5)] + payload["disease_aliases"] = ["alias1", "alias2"] + payload["biomarker_queries"] = ["EGFR"] + + parsed = parse_llm_query_expansion(payload, max_terms=6) + + assert parsed.primary_queries == [f"q{i}" for i in range(5)] + assert parsed.disease_aliases == ["alias1"] # only one slot left + assert parsed.biomarker_queries == [] # budget exhausted before this field + + +def test_reasoning_tags_are_stripped_before_json_extraction(): + """Reasoning models emit containing an echo of the schema; extracting from that + would return the schema instead of the answer.""" + payload = {field: [] for field in _FIRST_LEVEL_FIELDS} + payload["primary_queries"] = ["glioblastoma"] + reply = ( + "The schema wants primary_queries, disease_aliases, ..." + + json.dumps(payload) + ) + result = _FakeExpander(reply).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result["primary_queries"] == ["glioblastoma"] + + +@pytest.mark.parametrize( + "reply", + ["not json at all", json.dumps(["a", "list"]), RuntimeError("engine died")], +) +def test_failures_degrade_to_empty_not_raise(reply): + """Retrieval must survive a failed expansion: this is 1 of 8 channels, weight 0.5.""" + result = _FakeExpander(reply).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result == {field: [] for field in _FIRST_LEVEL_FIELDS} + + +def test_a_bare_string_field_is_not_shredded_into_characters(): + payload = {field: [] for field in _FIRST_LEVEL_FIELDS} + payload["primary_queries"] = "glioblastoma" + result = _FakeExpander(json.dumps(payload)).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result["primary_queries"] == ["glioblastoma"] + + +def test_empty_summary_skips_the_model_entirely(): + expander = _FakeExpander(RuntimeError("must not be called")) + assert expander.expand_first_level_queries(profile=None, matching_summary={}) == { + field: [] for field in _FIRST_LEVEL_FIELDS + } + + +def test_patient_text_includes_conditions_and_demographics(): + text = _first_level_patient_text(None, SUMMARY) + assert "metastatic breast cancer" in text + assert "HER2 positive" in text + assert "54" in text and "female" in text + + +def test_patient_text_omits_placeholder_demographics(): + text = _first_level_patient_text( + None, {"main_conditions": ["asthma"], "age": "all", "gender": "all"} + ) + assert "asthma" in text + assert "Age:" not in text and "Sex:" not in text + + +def test_builder_returns_none_unless_the_flag_is_set(): + assert build_first_level_expander({}) is None + assert build_first_level_expander({"search": {"first_level": {}}}) is None + assert ( + build_first_level_expander( + {"search": {"first_level": {"llm_expansion_enabled": False}}} + ) + is None + ) + + +def test_builder_degrades_to_none_when_construction_fails(): + """A misconfigured expander must not abort the run; the channel just stays empty.""" + config = { + "search": {"first_level": {"llm_expansion_enabled": True}}, + "model": {}, # no base_model -> QueryExpander raises + "query_expansion": {}, + } + assert build_first_level_expander(config) is None + + +def test_schema_caps_primary_queries_tightly_to_protect_the_shared_budget(): + """Guards the interaction pinned above: primary_queries is spent first out of + llm_max_terms, so its schema cap must leave room for the later channels.""" + from trialmatchai.matching.query_expansion import _FIRST_LEVEL_JSON_SCHEMA + + props = _FIRST_LEVEL_JSON_SCHEMA["properties"] + primary = props["primary_queries"]["maxItems"] + assert primary <= 3 + for field in ("biomarker_queries", "treatment_queries", "disease_aliases"): + assert props[field]["maxItems"] > primary From 57b5ab59d16ab878ad8d3fe69d67229090dd3856 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 17:13:56 +0200 Subject: [PATCH 5/5] fix(trec): judge shortlist depth on recall-aware metrics, not ndcg@10 ndcg@10 cannot answer a question about depth, and reading the A/B on it would have given exactly the wrong answer. It normalizes by the ideal over judged-AND-RANKED trials, which makes it recall-independent by construction: ranking fewer trials shrinks the ideal along with the DCG. Replaying the finished runs at reduced shortlist depth (re-ranking from cached CoT, no inference) shows it moving the WRONG way as depth grows, on all four runs. TREC 2021 / qwen36_medcpt: depth shortRec ndcg@10 ndcg_full@10 P@10(elig) 10 0.0787 0.8948 0.6028 0.3853 50 0.2565 0.8134 0.7916 0.6800 100 0.3924 0.8210 0.8113 0.7360 196 0.5774 0.8215 0.8197 0.7680 ndcg@10 reads BEST at depth 10, where the shortlist holds 8% of the relevant trials. Meanwhile ndcg_full@10 (+0.2169) and P@10(eligible) (+0.3827) both rise strongly with depth, because they are normalized against the full judged pool and so count a trial that never entered the shortlist as a miss. Same pattern on h100_baichuan/21, qwen36/22 and l40/23. So the funnel premise holds: depth does reach the top ten. It is simply invisible in the metric the comparison led with. ndcg_full@10 and P@10(eligible) move to QUALITY; ndcg@10 moves to DIAGNOSTIC with the reason recorded beside it. The replay also shows clear diminishing returns -- 25->50 buys +0.051 ndcg_full@10 while 100->196 buys +0.008 -- so the deeper arm should be expected to gain modestly, not doubled. TREC 2023 is the exception and is still climbing steeply at its full depth (shortlist recall only 0.254), which makes it the strongest case for more depth rather than 2021. --- scripts/compare_shortlist_ab.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/scripts/compare_shortlist_ab.py b/scripts/compare_shortlist_ab.py index 9718338a..d3efee6c 100644 --- a/scripts/compare_shortlist_ab.py +++ b/scripts/compare_shortlist_ab.py @@ -19,15 +19,24 @@ from pathlib import Path BASELINE = "fixed" +# Ordered deliberately. ndcg_full@10 normalizes by the ideal over the FULL judged pool, so a +# relevant trial that never entered the shortlist counts as a miss -- it is recall-aware, and +# it is the metric a depth change should be judged on. QUALITY = ( "shortlist_recall", - "recall@1000", - "ndcg@10", - "P@10(rel>=1)", + "ndcg_full@10", "P@10(eligible)", + "P@10(rel>=1)", + "recall@1000", ) COST = ("shortlist_size",) -DIAGNOSTIC = ("funnel_depth_loss", "shortlist_selection_delta") +# ndcg@10 sits here, NOT in QUALITY. It normalizes by the ideal over judged-AND-RANKED trials, +# which makes it recall-INDEPENDENT by construction: ranking fewer trials shrinks the ideal +# too. Replaying finished runs at reduced depth shows it moving the wrong way -- on TREC 2021 +# it reads 0.8948 at depth 10 and 0.8215 at depth 196, while ndcg_full@10 goes 0.6028 -> 0.8197 +# and P@10(eligible) 0.3853 -> 0.7680 over the same range. Judging a depth policy on ndcg@10 +# would conclude that more depth hurts, which is exactly backwards. +DIAGNOSTIC = ("ndcg@10", "funnel_depth_loss", "shortlist_selection_delta") def load_arm(root: Path, arm: str, track: str) -> dict | None: @@ -87,7 +96,10 @@ def main() -> int: print() print(" Reading this:") - print(" An arm at the SAME shortlist_size as 'fixed' with higher shortlist_recall") + print(" Judge a depth change on ndcg_full@10 and P@10(eligible). Both are recall-aware.") + print(" Do NOT judge it on ndcg@10 -- that is normalized over judged-AND-ranked trials,") + print(" so it is recall-independent and moves the WRONG way as depth grows.") + print(" An arm at the SAME shortlist_size as 'fixed' with higher recall-aware quality") print(" spends its compute better -- that gain is free.") print(" An arm with a LARGER shortlist_size bought its gain with GPU time; compare") print(" it against 'fixed' only after noting the extra cost.")