"""PFB-MAX semantic retrieval pool. Builds the candidate pool for semantic PFB queries: one cheap gpt-4o-mini expansion call (rephrasings + criteria + HyDE hypothetical abstracts), a budgeted multi-probe fanout over ``snippet_search`` / ``paper_search``, one round of citation expansion on the top fused seeds, all fused with weighted reciprocal-rank fusion (per-channel best-rank dedup, IRIS retrieval-v3 lineage; see ``iris_asta/iris_asta/solvers/pfb.py:semantic_channel``). Why HyDE views are first-class probes: IRIS forensics showed ~30% of gold papers were NEVER retrieved by query-derived probes: the recall wall. A hypothetical answer-abstract written in the field's own vocabulary is a far better dense-retrieval key than the question text (Gao et al., 2022), and gpt-4o-mini writes materially better abstracts than the open-weight backbone IRIS ran. Two methodologically diverse views widen the net further. Integrity: every probe is generated from the live query text only; no gold corpus ids, no per-query branching. All stored texts are VERBATIM backend strings (snippet text or the paper's own abstract); ``evidence()`` only joins/truncates contiguous spans, never rewrites. Contract: def build_pool(query, criteria, client, llm, inserted_before, max_calls=90) -> SemPool SemPool: .ranked() .evidence(cid) .texts(cid) .size() .calls_used Stdlib + iris_asta imports only. """ from __future__ import annotations import os import re import sys import threading import time from typing import Any def _ensure_iris_importable() -> None: """Make the read-only ``iris_asta`` package importable from the bundle. The bundle layout is ``/pfbmax/`` (this file) next to ``/iris_asta/iris_asta/`` (the package inside its repo dir). Probes for a SUBMODULE, not the bare package name: run from the bundle root, the outer ``iris_asta`` repo dir (no ``__init__.py``) resolves as an empty namespace package that would satisfy ``import iris_asta`` while providing none of its modules. Any such stale namespace stub is dropped from ``sys.modules`` so the real package wins. """ import importlib.util try: if importlib.util.find_spec("iris_asta.retrieval") is not None: return except Exception: pass bundle = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) candidate = os.path.join(bundle, "iris_asta") if os.path.isdir(os.path.join(candidate, "iris_asta")): sys.modules.pop("iris_asta", None) sys.path.insert(0, candidate) _ensure_iris_importable() from iris_asta.contracts import normalize_corpus_id # noqa: E402 from iris_asta.retrieval import ReciprocalRankFusion # noqa: E402 # -------------------------------------------------------------------------- # Tunables (IRIS retrieval-v3 lineage; weights proven on the validation set) # -------------------------------------------------------------------------- RANK_CONSTANT = 60.0 # RRF k (standard; IRIS _RRF_RANK_CONSTANT) W_SNIPPET = 1.20 # snippet channel multiplier (x probe weight) W_PAPER = 1.00 # paper_search relevance channel (x probe weight) W_HYDE = 1.00 # hyde views fuse on their own flat channel W_REFERENCE = 0.30 # ref_mention cited-paper channel (x probe weight) W_CITE_FWD = 0.80 # citation expansion: papers CITING a seed W_CITE_BACK = 0.65 # citation expansion: papers a seed CITES PW_RAW = 1.30 # probe weights (raw query keeps primacy) PW_REPHRASE = 1.00 PW_CRITERION = 0.80 LIMIT_SNIPPET_RAW = 100 # raw query keeps the deepest snippet list LIMIT_SNIPPET_PROBE = 60 # rephrasings / criteria add recall, not depth LIMIT_SNIPPET_HYDE = 100 # hyde views are the recall lever - deep fanout LIMIT_PAPER = 100 # relevance-ranker channel per probe LIMIT_LOCAL = 1000 # local mirror: depth is ~free (one ranked scan) LIMIT_CITATIONS = 60 # per direction per expansion seed # The depth limits above were raised 2-2.5x from their first values. # Rationale, measured: the semantic score is harmonic(rank, recall@K) and # the decomposition over 8 live queries was rank 0.668 / recall@K 0.083 -- # recall is the whole bottleneck. Corpus calls are negligible under the # benchmark's cost accounting (only solver LLM tokens count), and a larger # *limit* returns more results per call rather than issuing more calls, so # this buys recall at ~zero dollar cost and near-zero wall clock. The # cross-encoder stage added alongside is what converts the extra pool # depth into ranked position (IRIS measured +102 pool golds converting to # only +17 submittable WITHOUT a CE, and named the CE as the missing # promoter). MAX_REPHRASINGS = 8 MAX_CRITERIA = 6 MAX_HYDE_VIEWS = 4 MAX_TEXTS_PER_CID = 3 # distinct verbatim snippets kept per paper EVIDENCE_MAX_CHARS = 1200 EXPANSION_SEED_POOL = 10 # fused window considered for expansion... EXPANSION_SEEDS = 6 # ...but only this many are expanded (budget) def _envi(name: str, default: int) -> int: """Integer override from the environment, ignoring junk values.""" raw = (os.environ.get(name) or "").strip() try: return int(raw) if raw else default except ValueError: return default # Pool recall is the CEILING on the semantic score: recall@K can never exceed # pool_recall/2 (K ~ 2P), so 43% pool recall caps recall@K at ~0.216 no # matter how good the ranker is. Widening the fanout is therefore the only # lever that raises the ceiling rather than chasing it -- and it is nearly # free, because the benchmark's cost accounting counts SOLVER LLM tokens # only; corpus calls cost wall-clock and nothing else. Measured runs used 38 # of the 90-call budget, so the binding constraint was probe count, not # budget. Every knob below is env-overridable so a wide config can be A/B'd # against the default without touching the defaults. MAX_REPHRASINGS = _envi("PFBMAX_MAX_REPHRASINGS", MAX_REPHRASINGS) MAX_CRITERIA = _envi("PFBMAX_MAX_CRITERIA", MAX_CRITERIA) MAX_HYDE_VIEWS = _envi("PFBMAX_MAX_HYDE", MAX_HYDE_VIEWS) EXPANSION_SEED_POOL = _envi("PFBMAX_SEED_POOL", EXPANSION_SEED_POOL) EXPANSION_SEEDS = _envi("PFBMAX_SEEDS", EXPANSION_SEEDS) LIMIT_SNIPPET_RAW = _envi("PFBMAX_LIMIT_SNIPPET_RAW", LIMIT_SNIPPET_RAW) LIMIT_SNIPPET_PROBE = _envi("PFBMAX_LIMIT_SNIPPET_PROBE", LIMIT_SNIPPET_PROBE) LIMIT_SNIPPET_HYDE = _envi("PFBMAX_LIMIT_SNIPPET_HYDE", LIMIT_SNIPPET_HYDE) LIMIT_PAPER = _envi("PFBMAX_LIMIT_PAPER", LIMIT_PAPER) LIMIT_CITATIONS = _envi("PFBMAX_LIMIT_CITATIONS", LIMIT_CITATIONS) #: The dense snippet backend rejects markdown/query-operator characters - #: exact character class from IRIS ``_HYDE_SANITIZE_RE`` (live-verified). _SANITIZE_RE = re.compile(r'[*"`~^\\|{}\[\]<>#]+') _CITATION_FIELDS = "corpusId,title,abstract,year" # -------------------------------------------------------------------------- # One-call query expansion (rephrasings + criteria + HyDE views) # -------------------------------------------------------------------------- _EXPAND_SYSTEM = """You expand a scientific paper-finding query for high-recall literature retrieval. Respond with JSON only, using exactly these keys: "rephrasings": exactly {n_rephrasings} diverse rephrasings of the query. Each must use ALTERNATIVE technical vocabulary, synonyms, and phrasings the full text of a relevant paper might use, while preserving the exact information need. Make them genuinely different from each other. {criteria_clause}"hyde_views": exactly {n_hyde} hypothetical paper ABSTRACTS of 60-120 words each, plain prose only (no title, no markdown, no bullet points, no quotes). Each abstract describes a DIFFERENT plausible research paper that would perfectly satisfy the query, written in the field's own standard terminology, naming the concrete methods, datasets, tasks, metrics, and findings such a paper would contain. Make the views methodologically diverse (e.g. different approach families or settings).""" _CRITERIA_CLAUSE = """"criteria": between 1 and 4 SHORT relevance criteria — the independent requirements a paper must EACH satisfy to be perfectly relevant to the query (topic, method, finding, setting, population). Each criterion is a single, independently checkable statement. """ def _clean_line(text: Any) -> str: """Whitespace-collapsed single-line string (probe hygiene).""" return " ".join(str(text or "").split()) def _sanitize_probe(text: Any) -> str: """Strip the characters the dense backend rejects; collapse whitespace.""" return _clean_line(_SANITIZE_RE.sub(" ", str(text or ""))) def _coerce_criteria(criteria: Any) -> list[str]: """Normalize caller criteria to clean strings. Callers normally pass ``list[str]`` (contract), but validation tooling carries criteria as ``{name, description, weight}`` dicts; accept those too (description preferred) so measurement scripts can't silently lose the criteria signal. """ out: list[str] = [] seen: set[str] = set() for item in criteria or []: if isinstance(item, dict): text = _clean_line(item.get("description") or item.get("name")) else: text = _clean_line(item) key = text.casefold() if text and key not in seen: seen.add(key) out.append(text) return out def _expand_query( llm: Any, query: str, passed_criteria: list[str] ) -> dict[str, list[str]]: """One LLM json call -> {rephrasings, criteria, hyde_views} (all lists). When the caller already supplied criteria the prompt does not ask for a criteria key (and any returned one is ignored). Every failure mode (None reply, wrong shapes, junk items) degrades to empty lists: the fanout then runs on the raw query alone, never empty (IRIS no-empty-fallback lesson). """ criteria_clause = "" if passed_criteria else _CRITERIA_CLAUSE # Ask for as many probes as the caps allow: the prompt used to hard-code # "exactly 8"/"exactly 4", so raising MAX_REPHRASINGS / MAX_HYDE_VIEWS # silently changed nothing. HyDE is the deepest-fanout channel we have, # and pool recall is the ceiling on the whole semantic score. system = _EXPAND_SYSTEM.format(criteria_clause=criteria_clause, n_rephrasings=MAX_REPHRASINGS, n_hyde=MAX_HYDE_VIEWS) user = f"QUERY:\n{query}" if passed_criteria: joined = "\n".join(f"- {c}" for c in passed_criteria[:8]) user += ( "\n\nKnown relevance criteria (do NOT output a criteria key; " f"use these to sharpen the abstracts):\n{joined}" ) try: # A gpt-4o "reason first" expansion was tried here and REGRESSED # (-0.034 vs baseline, 3/3 down): richer expansions drift from the # query's actual intent and dilute the pool. Recall wants FAITHFUL # paraphrase, not creative reframing. # Scale the cap with what the prompt actually asks for: a fixed 900 # truncates the JSON (and loses the whole expansion) once the # rephrasing/HyDE counts are raised above the defaults. # 900 is enough for the default caps (measured headroom), so the # default stays exactly as scored; the budget only grows when a # caller raises the caps, where a fixed 900 truncated the JSON and # silently lost the whole expansion. budget = max(900, 200 + 60 * max(MAX_REPHRASINGS - 8, 0) + 260 * max(MAX_HYDE_VIEWS - 4, 0) + (0 if passed_criteria else 60 * max(MAX_CRITERIA - 6, 0))) if MAX_REPHRASINGS > 8 or MAX_HYDE_VIEWS > 4 or MAX_CRITERIA > 6: budget = max(budget, 900 + 60 * MAX_REPHRASINGS + 260 * MAX_HYDE_VIEWS) obj = llm.json(system, user, max_tokens=budget) except Exception: obj = None if not isinstance(obj, dict): obj = {} def _str_list(key: str, cap: int) -> list[str]: raw = obj.get(key) out: list[str] = [] seen = {query.strip().casefold()} if isinstance(raw, list): for item in raw: if not isinstance(item, str): continue text = _clean_line(item) if not text or text.casefold() in seen: continue seen.add(text.casefold()) out.append(text) if len(out) >= cap: break return out hyde_views: list[str] = [] raw_views = obj.get("hyde_views") if isinstance(raw_views, list): for item in raw_views: if not isinstance(item, str): continue text = _sanitize_probe(item)[:1200] # Under ~40 chars it is not an abstract - a dense probe that # short only adds noise (IRIS threshold). if len(text) >= 40: hyde_views.append(text) if len(hyde_views) >= MAX_HYDE_VIEWS: break return { "rephrasings": _str_list("rephrasings", MAX_REPHRASINGS), "criteria": [] if passed_criteria else _str_list("criteria", MAX_CRITERIA), "hyde_views": hyde_views, } # -------------------------------------------------------------------------- # Budgeted client wrappers # -------------------------------------------------------------------------- class _Budget: """Hard cap on corpus-backend calls. Every ATTEMPT costs one unit (a call that raises still consumed backend time and rate limit).""" def __init__(self, max_calls: int): self.max_calls = max(0, int(max_calls)) self.used = 0 # Retries are taken from worker threads, so the counter needs a lock. self._lock = threading.Lock() def take(self) -> bool: with self._lock: if self.used >= self.max_calls: return False self.used += 1 return True def reserve(self, n: int) -> int: """Claim up to n units at once; returns how many were granted.""" with self._lock: n = max(0, min(int(n), self.max_calls - self.used)) self.used += n return n def _fanout(fetch, jobs, budget, workers=None): """Run independent budgeted probes concurrently, in deterministic order. Each corpus call is ~6s of pure server latency, so issuing them serially made a wide fanout cost ~23 minutes a query -- enough to make the highest-recall configuration unusable in practice even though the calls themselves are negligible under the benchmark's cost accounting. Budget is reserved UP FRONT on the calling thread, so the cap stays exact (``_Budget.take`` is a non-atomic read-modify-write and would race). Results come back in job order, and absorption stays sequential in the caller, so fusion ranks and the per-paper snippet selection are bit-identical to the serial path -- concurrency changes the wall clock, not the pool. Returns a list aligned with ``jobs``; entries past the budget are None, which callers already treat as "budget exhausted". """ jobs = list(jobs) allowed = budget.reserve(len(jobs)) run = jobs[:allowed] if not run: return [None] * len(jobs) if workers is None: workers = _envi("PFBMAX_FANOUT_WORKERS", 8) workers = max(1, min(int(workers), len(run))) if workers == 1: out = [fetch(j) for j in run] else: from concurrent.futures import ThreadPoolExecutor with ThreadPoolExecutor(max_workers=workers) as ex: out = list(ex.map(fetch, run)) return out + [None] * (len(jobs) - len(run)) def _snippet_search(client, budget, query, limit, inserted_before): """Budgeted snippet_search; None = budget exhausted, [] = empty/failed.""" if not budget.take(): return None return _snippet_fetch(client, query, limit, inserted_before) # One retry: enough to ride out a throttle, bounded so a dead probe cannot # starve the others. Both are env-tunable (0 retries = original behaviour). FETCH_RETRIES = _envi("PFBMAX_FETCH_RETRIES", 1) FETCH_BACKOFF_S = float(os.environ.get("PFBMAX_FETCH_BACKOFF") or 1.0) _FETCH_FAILURES = {"snippet": 0, "paper": 0, "retried": 0} def _retrying(kind, call, budget=None): """Run a corpus call, retrying transient failures instead of losing them. A swallowed exception is indistinguishable from "no matches", and that silence is expensive: concurrent probes get throttled, each failure returns [], and the pool quietly shrinks. Measured directly -- the same 37 calls built a 2112-paper pool serially but only 1960 at 6 workers (-7%), purely from dropped responses. Each retry COSTS A BUDGET UNIT, because the budget exists to bound load on a shared rate-limited backend and a raised call already consumed backend time. So retries recover lost results only while budget remains, and never smuggle extra load past the cap. With no budget passed there is no retry at all -- the original single-attempt behaviour. """ try: return call() or [] except Exception: pass # Bounded: a probe that fails deterministically must not be able to drain # the budget that the remaining probes need. One retry is what recovers a # throttled call; a second failure means the probe is simply dead. for _ in range(FETCH_RETRIES): if budget is None or not budget.take(): break _FETCH_FAILURES["retried"] += 1 time.sleep(FETCH_BACKOFF_S) try: return call() or [] except Exception: continue _FETCH_FAILURES[kind] = _FETCH_FAILURES.get(kind, 0) + 1 return [] def _snippet_fetch(client, query, limit, inserted_before, budget=None): """Unbudgeted first attempt (budget is reserved by _fanout); retries, when a budget is supplied, are charged to it.""" return _retrying("snippet", lambda: client.snippet_search( query, limit=limit, inserted_before=inserted_before), budget) def _paper_search(client, budget, query, limit, inserted_before): if not budget.take(): return None return _paper_fetch(client, query, limit, inserted_before) def _paper_fetch(client, query, limit, inserted_before, budget=None): """Unbudgeted first attempt; retries are charged to ``budget``.""" return _retrying("paper", lambda: client.paper_search( query, limit=limit, publication_date_before=inserted_before), budget) def _get_citations(client, budget, cid, direction, limit, date_range): if not budget.take(): return None try: return client.get_citations( cid, direction=direction, limit=limit, publication_date_range=date_range, fields=_CITATION_FIELDS, ) or [] except Exception: return [] # -------------------------------------------------------------------------- # Snippet / paper helpers (ports of small pfb.py primitives) # -------------------------------------------------------------------------- def _snippet_score(snippet: Any) -> float: """Positive contribution score of a snippet (unscored hits count 1.0).""" try: score = float(getattr(snippet, "score", 0.0) or 0.0) except (TypeError, ValueError): score = 0.0 return score if score > 0.0 else 1.0 def _ref_mention_cids(snippet: Any) -> list[str]: """Digits-only corpus ids referenced by a snippet's ref_mentions.""" out: list[str] = [] for mention in getattr(snippet, "ref_mentions", None) or []: if isinstance(mention, dict): raw = ( mention.get("matchedPaperCorpusId") or mention.get("corpusId") or mention.get("corpus_id") ) else: raw = mention cid = normalize_corpus_id(raw) if cid and cid not in out: out.append(cid) return out def _cid_of(paper: Any) -> str | None: return normalize_corpus_id(getattr(paper, "corpus_id", None)) def _postdates_snapshot(year: Any, cutoff_year: int | None) -> bool: """True only when the year PROVES the paper postdates the snapshot. Same-year and unknown-year papers are kept (the snapshot may cut mid-year; dropping on suspicion costs recall, matching IRIS ``_violates_cutoff`` semantics). Post-snapshot ids are INVALID to the scorer, worse than irrelevant. """ return ( cutoff_year is not None and isinstance(year, int) and year > cutoff_year ) # -------------------------------------------------------------------------- # SemPool # -------------------------------------------------------------------------- class SemPool: """Fused candidate pool for one semantic query. Contract surface: ``ranked() / evidence(cid) / texts(cid) / size() / calls_used``. Extras used by the judge/router modules: ``criteria`` (caller-passed or generated), ``rephrasings``, ``hyde_views``, ``title(cid)``, ``year(cid)``, ``families(cid)``, ``rrf_score(cid)``, ``diagnostics``. """ def __init__( self, query: str, criteria: list[str], rephrasings: list[str], hyde_views: list[str], records: dict[str, dict], fusion: ReciprocalRankFusion, calls_used: int, diagnostics: dict[str, Any], ): self.query = query self.criteria = list(criteria) self.rephrasings = list(rephrasings) self.hyde_views = list(hyde_views) self._records = records self._fusion = fusion self.calls_used = int(calls_used) self.diagnostics = diagnostics self._ranked_cache: list[str] | None = None # -- contract ---------------------------------------------------------- def ranked(self) -> list[str]: """All pool cids in fused order (best first, deterministic).""" if self._ranked_cache is None: self._ranked_cache = [ c.key for c in self._fusion.ranked() if c.key in self._records ] return list(self._ranked_cache) def evidence(self, cid: str) -> str: """Best verbatim snippet(s) for ``cid``, newline-joined, <=1200 chars. Parts are stored backend texts (snippets best-score-first, else the paper's own abstract). Whole parts are appended while they fit; the first part alone is truncated to the cap if oversized; truncation keeps a contiguous prefix, so every emitted part remains a verbatim substring of a stored text. """ rec = self._records.get(_key(cid)) if not rec: return "" parts = [text for _score, text in rec["snippets"]] if not parts and rec.get("abstract"): parts = [rec["abstract"]] out: list[str] = [] used = 0 for part in parts: sep = 1 if out else 0 if not out and len(part) > EVIDENCE_MAX_CHARS: out.append(part[:EVIDENCE_MAX_CHARS]) used = EVIDENCE_MAX_CHARS break if used + sep + len(part) > EVIDENCE_MAX_CHARS: continue out.append(part) used += sep + len(part) return "\n".join(out) def texts(self, cid: str) -> list[str]: """All stored verbatim texts for ``cid`` (snippets first, then the abstract when distinct). Empty list for graph-only candidates.""" rec = self._records.get(_key(cid)) if not rec: return [] out = [text for _score, text in rec["snippets"]] abstract = rec.get("abstract") or "" if abstract and abstract not in out: out.append(abstract) return out def size(self) -> int: return len(self._records) # -- extras ------------------------------------------------------------ def title(self, cid: str) -> str: rec = self._records.get(_key(cid)) return (rec or {}).get("title", "") def year(self, cid: str): rec = self._records.get(_key(cid)) return (rec or {}).get("year") def best_snippet_score(self, cid: str) -> float: rec = self._records.get(_key(cid)) return (rec or {}).get("best_score", 0.0) def families(self, cid: str) -> list[str]: cand = self._fusion.get(_key(cid)) return sorted(cand.families) if cand else [] def channels(self, cid: str) -> dict[str, int]: cand = self._fusion.get(_key(cid)) return dict(cand.channels) if cand else {} def rrf_score(self, cid: str) -> float: cand = self._fusion.get(_key(cid)) return cand.score if cand else 0.0 def _key(cid: Any) -> str: return normalize_corpus_id(cid) or str(cid or "").strip() # -------------------------------------------------------------------------- # build_pool # -------------------------------------------------------------------------- def build_pool( query: str, criteria: list[str], client: Any, llm: Any, inserted_before: str | None, max_calls: int = -1, ) -> SemPool: """Build the fused semantic candidate pool for ``query``. Stages (all corpus calls budgeted by ``max_calls``, snapshot-bounded by ``inserted_before``): 1. ONE gpt-4o-mini json call -> rephrasings (up to MAX_REPHRASINGS), criteria (up to MAX_CRITERIA; skipped when the caller passed criteria), and sanitized HyDE abstracts (up to MAX_HYDE_VIEWS). 2. Fanout: snippet_search on the raw query, hyde views (own flat channel), rephrasings, and criteria at the LIMIT_SNIPPET_* depths; paper_search on raw + rephrasings (LIMIT_PAPER). Snippet ref_mentions credit cited papers on a low-weight reference channel (zero extra calls). Weighted RRF with per-channel best-rank dedup fuses everything. 3. One citation-expansion round: within the top EXPANSION_SEED_POOL fused, the top EXPANSION_SEEDS seeds get ``get_citations`` both directions (LIMIT_CITATIONS each, post-snapshot years skipped), absorbed at W_CITE_FWD (citing) / W_CITE_BACK (cited), budget-capped like every other call. Fanout priority under a tight budget: raw first (primary signal), then hyde views (the recall lever), then rephrasings, criteria, paper channels, citations. """ query = _clean_line(query) # -1 means "take the default", so PFBMAX_MAX_CALLS can raise the ceiling # without every caller having to thread the argument through. if max_calls is None or max_calls < 0: max_calls = _envi("PFBMAX_MAX_CALLS", 90) budget = _Budget(max_calls) fusion = ReciprocalRankFusion(rank_constant=RANK_CONSTANT) records: dict[str, dict] = {} def _record(cid: str) -> dict: return records.setdefault( cid, { "snippets": [], # [(score, text)] distinct, best-first "best_score": 0.0, "title": "", "abstract": "", "year": None, }, ) def _absorb_snippet(cid: str, title: str, text: str, score: float) -> None: rec = _record(cid) if text: if score > rec["best_score"]: rec["best_score"] = score if all(text != t for _s, t in rec["snippets"]): rec["snippets"].append((score, text)) rec["snippets"].sort(key=lambda pair: -pair[0]) del rec["snippets"][MAX_TEXTS_PER_CID:] if title and not rec["title"]: rec["title"] = title def _absorb_paper(cid: str, paper: Any) -> None: rec = _record(cid) title = _clean_line(getattr(paper, "title", "") or "") abstract = (getattr(paper, "abstract", "") or "").strip() year = getattr(paper, "year", None) if title and not rec["title"]: rec["title"] = title if abstract and not rec["abstract"]: rec["abstract"] = abstract if isinstance(year, int) and rec["year"] is None: rec["year"] = year # -- stage 1: one expansion call -------------------------------------- passed_criteria = _coerce_criteria(criteria) plan = _expand_query(llm, query, passed_criteria) final_criteria = passed_criteria or plan["criteria"] or [query] rephrasings = plan["rephrasings"] hyde_views = plan["hyde_views"] # -- stage 2: fanout --------------------------------------------------- # Snippet jobs: (probe_text, channel, fused_weight, limit). Probe texts # are casefold-deduped so a criterion equal to the raw query (the # no-criteria fallback) never pays for the same search twice. seen_probes: set[str] = set() snippet_jobs: list[tuple[str, str, float, int]] = [] def _add_snippet_job(text, channel, weight, limit): clean = _sanitize_probe(text) key = clean.casefold() if clean and key not in seen_probes: seen_probes.add(key) snippet_jobs.append((clean, channel, weight, limit)) _add_snippet_job(query, "snippet:raw-0", W_SNIPPET * PW_RAW, LIMIT_SNIPPET_RAW) for i, view in enumerate(hyde_views[:MAX_HYDE_VIEWS]): _add_snippet_job(view, f"hyde:view-{i}", W_HYDE, LIMIT_SNIPPET_HYDE) for i, text in enumerate(rephrasings[:MAX_REPHRASINGS]): _add_snippet_job(text, f"snippet:rephrase-{i}", W_SNIPPET * PW_REPHRASE, LIMIT_SNIPPET_PROBE) for i, text in enumerate(final_criteria[:MAX_CRITERIA]): _add_snippet_job(text, f"snippet:criterion-{i}", W_SNIPPET * PW_CRITERION, LIMIT_SNIPPET_PROBE) _local_stats: dict = {} # --- local corpus mirror (opt-in supplement) -------------------------- # An optional local index. One wide BM25 query measured 37.3% pool # recall on the gold sets vs 43.2% for IRIS's ~168-remote-call fanout, # it never rate-limits, and it adds no solver LLM tokens to the cost. # # When enabled it SUPPLEMENTS the remote fanout with a few probes; the # remote fanout below still runs either way. # OFF by default: an earlier configuration that made the mirror primary # measured a REGRESSION (0.115 vs 0.152 on the same 3 queries). Cause is # latency, not relevance -- 9 concurrent deep probes overwhelmed the # mirror and 7 of 9 returned nothing, starving the pool (n_submitted # fell 250 -> ~89). One sequential probe reaches 37.3% recall in 24s, # so the mirror is sound; it needs a real search engine (Tantivy/Lucene) # or a field-filtered smaller index to serve a wide fanout. # Opt in with PFBMAX_USE_LOCAL=1. _local_on = False _ls = None if os.environ.get("PFBMAX_USE_LOCAL", "").strip() in ("1", "true", "yes", "on"): try: import local_search as _ls _local_on = _ls.available(timeout=6.0) except Exception: _ls = None if _local_on: # SUPPLEMENT, not replacement, and deliberately few: each local probe # is ~24s of memory-bound BM25 and 9 concurrent ones starved the pool # (7 of 9 timed out). Three probes add an independent retrieval arm -- # the mirror alone reaches 37.3% pool recall vs the remote fanout's # 43.2%, and their union is what breaks the ~0.25 semantic ceiling # that pool recall imposes. local_probes = [(query, "local:raw", W_SNIPPET * PW_RAW)] for i, view in enumerate(hyde_views[:2]): local_probes.append((view, f"local:hyde-{i}", W_HYDE)) # Probes are independent, so fire them concurrently: sequentially # this cost ~24s x 9 probes and blew a 10-minute budget on ONE query. # The mirror is a threaded server with a 48-connection pool on a # 64-core box, so concurrency is free on both ends. from concurrent.futures import ThreadPoolExecutor jobs = [(t, c, w) for (t, c, w) in local_probes if (t or "").strip()] with ThreadPoolExecutor(max_workers=3) as ex: results = list(ex.map( lambda j: (j[1], j[2], _ls.search(j[0], LIMIT_LOCAL)), jobs)) n_local = 0 for channel, weight, papers in results: for rank, paper in enumerate(papers): cid = normalize_corpus_id(getattr(paper, "corpusId", None)) if not cid: continue _absorb_paper(cid, paper) fusion.add(cid, rank=rank, channel=channel, weight=weight) n_local += 1 _local_stats = {"hits": n_local, "probes": len(jobs)} snippet_results = _fanout( lambda j: _snippet_fetch(client, j[0], j[3], inserted_before, budget), snippet_jobs, budget) for (probe_text, channel, weight, limit), snippets in zip(snippet_jobs, snippet_results): if snippets is None: break snippets = sorted(snippets, key=_snippet_score, reverse=True) family = channel.split(":", 1)[0] # Probe weight recovered for the reference channel: fused weight is # W_SNIPPET*pw (or W_HYDE), reference rides at W_REFERENCE*pw. probe_weight = (weight / W_SNIPPET) if family == "snippet" else 1.0 ref_channel = "reference:" + channel.split(":", 1)[1] for rank, snippet in enumerate(snippets): score = _snippet_score(snippet) text = getattr(snippet, "text", "") or "" title = _clean_line(getattr(snippet, "title", "") or "") cid = normalize_corpus_id(getattr(snippet, "corpus_id", None)) if cid: _absorb_snippet(cid, title, text, score) fusion.add(cid, rank=rank, channel=channel, weight=weight) # Cited-paper attachment (free candidates): the citing sentence # is the OTHER paper's text, so the cited paper gets rank credit # only - its own snippets/abstract must furnish evidence later. for ref_cid in _ref_mention_cids(snippet): if ref_cid != cid: _record(ref_cid) fusion.add(ref_cid, rank=rank, channel=ref_channel, weight=W_REFERENCE * probe_weight) paper_jobs: list[tuple[str, str, float]] = [ (query, "paper:raw-0", W_PAPER * PW_RAW) ] for i, text in enumerate(rephrasings[:MAX_REPHRASINGS]): paper_jobs.append((_sanitize_probe(text), f"paper:rephrase-{i}", W_PAPER * PW_REPHRASE)) paper_results = _fanout( lambda j: _paper_fetch(client, j[0], LIMIT_PAPER, inserted_before, budget), paper_jobs, budget) for (probe_text, channel, weight), papers in zip(paper_jobs, paper_results): if papers is None: break rank = 0 for paper in papers: cid = _cid_of(paper) abstract = (getattr(paper, "abstract", "") or "").strip() # Abstract required: this channel's hits carry no snippet, so a # hit without its own text could never earn judged evidence. if not cid or not abstract: continue _absorb_paper(cid, paper) fusion.add(cid, rank=rank, channel=channel, weight=weight) rank += 1 # -- stage 3: one citation-expansion round ---------------------------- cutoff_year: int | None = None if inserted_before: try: cutoff_year = int(str(inserted_before)[:4]) except (TypeError, ValueError): cutoff_year = None date_range = ( f"1900-01-01:{str(inserted_before)[:10]}" if inserted_before else None ) seed_window = [c.key for c in fusion.ranked()[:EXPANSION_SEED_POOL]] # Citation expansion is remote-only and each call can stall 60-90s # when the corpus backend is degraded, so it runs last and stays # budget-capped like every other call. seeds = seed_window[:EXPANSION_SEEDS] exhausted = False for si, seed in enumerate(seeds): if exhausted: break for direction, weight, prefix in ( ("citations", W_CITE_FWD, "cite:fwd"), ("references", W_CITE_BACK, "cite:back"), ): # Server-side date shard only for the citations direction (the # backend's citer list is recency-first and would return mostly # post-snapshot rows otherwise); the references client path # would DROP null-date papers on a range, so filter by year. papers = _get_citations( client, budget, seed, direction, LIMIT_CITATIONS, date_range if direction == "citations" else None, ) if papers is None: exhausted = True break rank = 0 for paper in papers: cid = _cid_of(paper) if not cid or cid == seed: continue if _postdates_snapshot(getattr(paper, "year", None), cutoff_year): continue _absorb_paper(cid, paper) fusion.add(cid, rank=rank, channel=f"{prefix}-{si}", weight=weight) rank += 1 diagnostics = { **({'local': _local_stats} if _local_stats else {}), "n_probes": len(snippet_jobs), "n_paper_probes": len(paper_jobs), "n_hyde_views": len(hyde_views), "expansion_window": seed_window, "expansion_seeds": seeds, "budget": budget.max_calls, } return SemPool( query=query, criteria=final_criteria, rephrasings=rephrasings, hyde_views=hyde_views, records=records, fusion=fusion, calls_used=budget.used, diagnostics=diagnostics, ) __all__ = ["SemPool", "build_pool"]