| """PFB-MAX router: one LLM classification call + channel dispatch. |
| |
| Contract: |
| |
| def solve(query: str, client, llm: LLM, inserted_before: str | None) -> Submission |
| |
| Routes by LLM classification (semantic | specific | metadata), NEVER by query id. |
| |
| Design (IRIS-proven; ported from iris_asta/iris_asta/solvers/pfb.py): |
| * ONE json_mode LLM call classifies the query into |
| {"route": "semantic" | "specific" | "metadata"}; any parse failure or |
| invalid label -> "semantic" (the judged channel degrades gracefully; a |
| wrong deterministic plan returns a confidently wrong set). |
| * Deterministic guard (IRIS v2 forensics; one such misroute alone cost |
| 0.250 -> 0.022): a "metadata" label on relevance-phrased queries |
| ("find papers relevant/related to ...") is a CONTENT ask and is forced |
| back to "semantic", date constraints or not. |
| * Dispatch: semantic -> semantic_retrieval.build_pool + |
| semantic_judge.rank_and_emit; specific -> specific_solver.solve_specific; |
| metadata -> metadata_solver.solve_metadata. |
| * No-empty-fallback (IRIS lesson): any exception OR empty result from a |
| deterministic channel falls back to the semantic channel; total failure |
| returns [] (never raises). |
| * Graceful degradation: sibling modules may not have landed yet. Missing |
| ones are reported ONCE on stderr and replaced by minimal built-in stubs |
| (semantic: one snippet_search + one paper_search fusion with verbatim |
| evidence; specific: single title-resolve emit; metadata: empty -> the |
| semantic fallback fires), so this file runs today and improves as |
| siblings land. |
| |
| Integrity: logic keys on query text + general rules only; no |
| gold corpus ids, no per-query branching; all emitted evidence is verbatim |
| corpus text (title/snippet/abstract windows). Stdlib only. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import importlib |
| import os |
| import re |
| import sys |
|
|
| Submission = list[tuple[str, str]] |
|
|
| |
| MAX_RESULTS = 250 |
|
|
| |
| |
| _PFBMAX_DIR = os.path.dirname(os.path.abspath(__file__)) |
| if _PFBMAX_DIR not in sys.path: |
| sys.path.insert(0, _PFBMAX_DIR) |
|
|
| |
| |
| |
| ROUTE_MARKER = "[pfbmax:route]" |
| CRITERIA_MARKER = "[pfbmax:criteria]" |
|
|
| _VALID_ROUTES = ("semantic", "specific", "metadata") |
|
|
| |
| |
| |
| |
| _RELEVANCE_PHRASING_RE = re.compile( |
| r"\b(?:papers?|articles?|works?|studies|publications?|literature)\s+" |
| r"(?:relevant|related)\s+to\b", |
| re.IGNORECASE, |
| ) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| _DISCRIMINATING_META_RE = re.compile( |
| r"(?:\bauthors?\b|\bauthored\b|\bco-?authored\b|\bby\s+(?-i:[A-Z])|\bvenue\b|\bconference\b|\bjournal\b|\bproceedings\b|\bworkshop\b|(?-i:\b[A-Z]{2,}\b)|\bcit(?:e|es|ed|ing|ation|ations)\b|\breferences?\b|\bmore\s+than\s+\d+|\bat\s+least\s+\d+|\bfewer\s+than\s+\d+|\bover\s+\d+)", |
| re.IGNORECASE, |
| ) |
|
|
|
|
| |
| |
| |
| _DATE_CUE_RE = re.compile( |
| r"(?:\b(?:since|after|before|during|between|from)\s+(?:19|20)\d{2}\b|\b(?:19|20)\d{2}\b|\bpublished\s+(?:on\s+or\s+)?(?:after|before|since)\b|\blast\s+\d+\s+years?\b|\brecent\s+(?:years?|work|papers?)\b)", |
| re.IGNORECASE, |
| ) |
|
|
|
|
| |
| |
| |
|
|
| ROUTE_SYSTEM = f"""You route queries for a scientific paper-finding system. {ROUTE_MARKER} |
| Classify the query into exactly one category: |
| - "specific": the query asks for ONE particular known paper, referred to by \ |
| (approximate) name or an unmistakable description of a single paper. \ |
| Examples: "the BERT paper"; "find 'Attention Is All You Need'"; \ |
| "the paper that introduced residual networks". |
| - "metadata": the wanted set is defined by structured constraints -- venue, \ |
| publication year(s), author names, citation count, or citation relations \ |
| (papers citing X / cited by X / NOT citing Y), possibly nested or negated. \ |
| Examples: "NeurIPS 2021 papers by Kaiming He"; "papers since 2019 citing \ |
| 'Attention Is All You Need' but not 'BERT' with over 100 citations". |
| - "semantic": the wanted papers are described by their CONTENT/topic/findings \ |
| in natural language. Examples: "papers showing retrieval augmentation reduces \ |
| hallucination in QA"; "work on protein folding with diffusion models". |
| If constraints and content BOTH appear, choose "metadata" only when the \ |
| structured constraints alone pin down the answer set; otherwise "semantic". |
| Respond with JSON only: {{"route": "specific" | "metadata" | "semantic"}}""" |
|
|
| CRITERIA_SYSTEM = f"""You decompose a scientific paper-finding query into relevance criteria. {CRITERIA_MARKER} |
| List the INDEPENDENT requirements a paper must EACH satisfy to be PERFECTLY |
| relevant to the query -- topic, method, finding, setting, population, etc. |
| Each criterion must be a single, independently checkable statement about the |
| paper. Use between 1 and 6 criteria; fewer, sharper criteria are better. |
| Respond with JSON only: {{"criteria": ["...", "..."]}}""" |
|
|
|
|
| |
| |
| |
|
|
| def classify(query: str, llm) -> str: |
| """One LLM json call -> "semantic" | "specific" | "metadata". |
| |
| Safe default "semantic" on: llm exception, non-dict reply, missing or |
| invalid label. Deterministic guard: label "metadata" + relevance |
| phrasing in the query -> "semantic" (applied to the RAW label, before |
| validity, exactly as IRIS does). |
| """ |
| try: |
| obj = llm.json(ROUTE_SYSTEM, query, max_tokens=64) or {} |
| except Exception: |
| obj = {} |
| if not isinstance(obj, dict): |
| obj = {} |
| label = str(obj.get("route") or obj.get("category") or "").strip().lower() |
| if label == "metadata" and _RELEVANCE_PHRASING_RE.search(query or ""): |
| return "semantic" |
| |
| |
| |
| if (label == "metadata" |
| and _DATE_CUE_RE.search(query or "") |
| and not _DISCRIMINATING_META_RE.search(query or "")): |
| return "semantic" |
| if label in _VALID_ROUTES: |
| return label |
| return "semantic" |
|
|
|
|
| def derive_criteria(query: str, llm) -> list[str]: |
| """Per-query relevance criteria (1-6 strings; fallback: the query itself). |
| |
| The semantic channel contract (build_pool / rank_and_emit) consumes |
| criteria; the router derives them once from the query text (IRIS |
| `_derive_criteria` port) and shares them across both stages. |
| """ |
| try: |
| obj = llm.json(CRITERIA_SYSTEM, query, max_tokens=400) or {} |
| except Exception: |
| obj = {} |
| raw = obj.get("criteria") if isinstance(obj, dict) else None |
| criteria: list[str] = [] |
| if isinstance(raw, list): |
| for item in raw: |
| if isinstance(item, str) and item.strip(): |
| criteria.append(item.strip()) |
| if len(criteria) >= 6: |
| break |
| return criteria or [(query or "").strip() or "relevant to the query"] |
|
|
|
|
| |
| |
| |
|
|
| _SIB_CACHE: dict[str, object] = {} |
| _REPORTED: set[str] = set() |
|
|
|
|
| def _report_once(key: str, msg: str) -> None: |
| if key not in _REPORTED: |
| _REPORTED.add(key) |
| print(msg, file=sys.stderr, flush=True) |
|
|
|
|
| def _sib(name: str): |
| """Import sibling pfbmax module ``name`` (flat layout); None if absent. |
| |
| A sibling that exists but fails to import is reported once and treated |
| as absent (the stub keeps the pipeline alive). Tests monkeypatch this |
| function to control dispatch. |
| """ |
| if name in _SIB_CACHE: |
| return _SIB_CACHE[name] |
| try: |
| mod = importlib.import_module(name) |
| except ImportError: |
| mod = None |
| except Exception as exc: |
| _report_once( |
| f"broken:{name}", |
| f"[pfbmax.router] sibling {name}.py failed to import ({exc!r}) " |
| "-- treating as absent (stub behavior)", |
| ) |
| mod = None |
| _SIB_CACHE[name] = mod |
| return mod |
|
|
|
|
| |
| |
| |
|
|
| def _norm_cid(cid) -> str | None: |
| """Digits-only corpus id, or None (non-normalizable ids can never match |
| gold and would only waste ranked slots -- IRIS `_cid_of` lesson).""" |
| s = str(cid if cid is not None else "").strip() |
| if s.lower().startswith("corpusid:"): |
| s = s[len("corpusid:"):].strip() |
| return s if s.isdigit() else None |
|
|
|
|
| def _finalize(results) -> Submission: |
| """Dedupe by corpus id (first occurrence wins = best rank), normalize, |
| coerce evidence to str, cap at the scorer's 250.""" |
| out: Submission = [] |
| seen: set[str] = set() |
| for item in results or []: |
| try: |
| cid, ev = item[0], item[1] |
| except Exception: |
| continue |
| c = _norm_cid(cid) |
| if c is None or c in seen: |
| continue |
| seen.add(c) |
| out.append((c, str(ev) if ev is not None else "")) |
| if len(out) >= MAX_RESULTS: |
| break |
| return out |
|
|
|
|
| def _violates_cutoff(paper, inserted_before) -> bool: |
| """True when the paper is provably outside the corpus snapshot (IRIS |
| `_violates_cutoff` port): year strictly greater than the cutoff year. |
| Same-year papers are kept (the snapshot may cut mid-year).""" |
| if not inserted_before: |
| return False |
| try: |
| cutoff_year = int(str(inserted_before)[:4]) |
| except (TypeError, ValueError): |
| return False |
| year = getattr(paper, "year", None) |
| return isinstance(year, int) and year > cutoff_year |
|
|
|
|
| def _paper_evidence(paper) -> str: |
| """Markdown evidence for a paper: title/year + verbatim abstract lead-in.""" |
| title = getattr(paper, "title", "") or "(untitled)" |
| year = getattr(paper, "year", None) |
| abstract = (getattr(paper, "abstract", "") or "")[:400] |
| head = f"**{title}**" + (f" ({year})" if isinstance(year, int) else "") |
| return (head + ("\n" + abstract if abstract else "")).strip() |
|
|
|
|
| |
| |
| |
|
|
| def _fallback_semantic(query: str, client, llm, inserted_before) -> Submission: |
| """Degraded semantic channel: one snippet_search + one paper_search, |
| snippet hits first (relevance-ranked), verbatim evidence windows.""" |
| out: Submission = [] |
| seen: set[str] = set() |
| try: |
| snippets = client.snippet_search( |
| query, limit=100, inserted_before=inserted_before |
| ) |
| except Exception: |
| snippets = [] |
| for sn in snippets or []: |
| cid = _norm_cid(getattr(sn, "corpus_id", None)) |
| if cid is None or cid in seen: |
| continue |
| seen.add(cid) |
| title = getattr(sn, "title", "") or "" |
| text = (getattr(sn, "text", "") or "")[:600] |
| ev = (f"**{title}**\n{text}").strip() if title else text |
| out.append((cid, ev)) |
| try: |
| papers = client.paper_search( |
| query, limit=100, publication_date_before=inserted_before |
| ) |
| except Exception: |
| papers = [] |
| for p in papers or []: |
| cid = _norm_cid(getattr(p, "corpus_id", None)) |
| if cid is None or cid in seen or _violates_cutoff(p, inserted_before): |
| continue |
| seen.add(cid) |
| out.append((cid, _paper_evidence(p))) |
| return out[:MAX_RESULTS] |
|
|
|
|
| def _fallback_specific(query: str, client, llm, inserted_before) -> Submission: |
| """Degraded specific channel: single title-resolve emit (every extra |
| result halves precision on 1-gold navigational queries -- IRIS lesson). |
| Empty on a miss; the caller's semantic fallback takes over.""" |
| try: |
| paper = client.search_paper_by_title(query) |
| except Exception: |
| paper = None |
| if paper is None or _violates_cutoff(paper, inserted_before): |
| return [] |
| cid = _norm_cid(getattr(paper, "corpus_id", None)) |
| if cid is None: |
| return [] |
| return [(cid, _paper_evidence(paper))] |
|
|
|
|
| |
| |
| |
|
|
| def _semantic_channel(query, client, llm, inserted_before, trace=None) -> Submission: |
| sr = _sib("semantic_retrieval") |
| sj = _sib("semantic_judge") |
| if sr is not None and sj is not None: |
| try: |
| criteria = derive_criteria(query, llm) |
| pool = sr.build_pool(query, criteria, client, llm, inserted_before) |
| if trace is not None: |
| try: |
| trace["pool_size"] = pool.size() |
| except Exception: |
| pass |
| trace["pool_calls"] = getattr(pool, "calls_used", None) |
| |
| |
| |
| ce = _sib("ce_rerank") |
| if ce is not None: |
| try: |
| pool = ce.ce_rerank_pool( |
| query, criteria, pool, |
| trace=trace if trace is not None else None) |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"ce:{exc!r}") |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| if os.environ.get("PFBMAX_CJ_POOL", "").strip() in ("1", "true", "yes"): |
| cjp = _sib("criterion_judge") |
| if cjp is not None and hasattr(cjp, "rerank_pool"): |
| try: |
| pool = cjp.rerank_pool( |
| query, criteria, pool, |
| meter=getattr(llm, "meter", None), |
| trace=trace if trace is not None else None) |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"cj_pool:{exc!r}") |
| if os.environ.get("PFBMAX_PERCRIT_POOL", "").strip() in ("1", "true", "yes"): |
| pcm = _sib("percrit_rank") |
| if pcm is not None and hasattr(pcm, "rerank_pool"): |
| try: |
| pool = pcm.rerank_pool(query, criteria, pool, |
| trace=trace if trace is not None else None) |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"percrit_pool:{exc!r}") |
| sub = sj.rank_and_emit(query, criteria, pool, llm) |
| |
| |
| |
| if os.environ.get("PFBMAX_DILIGENT", "").strip() in ("1", "true", "yes"): |
| dmod, cjm = _sib("diligent"), _sib("criterion_judge") |
| if dmod is not None and cjm is not None: |
| try: |
| extra = dmod.expand(query, criteria, pool, client, |
| inserted_before, cjm, |
| trace=trace if trace is not None else None) |
| if extra: |
| |
| |
| |
| |
| |
| have = {c for c, _e in sub} |
| fresh = [(c, ev, sc) for c, ev, sc in extra |
| if c not in have] |
| fresh.sort(key=lambda r: -r[2]) |
| merged, k = [], 0 |
| for i, row in enumerate(sub): |
| merged.append(row) |
| |
| if k < len(fresh) and (i + 1) % 4 == 0: |
| merged.append((fresh[k][0], fresh[k][1])) |
| k += 1 |
| merged += [(c, ev) for c, ev, _s in fresh[k:]] |
| sub = merged[:MAX_RESULTS] |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"diligent:{exc!r}") |
| |
| |
| |
| |
| |
| |
| if os.environ.get("PFBMAX_IRIS_CHANNEL", "").strip() in ("1", "true", "yes"): |
| icmod = _sib("iris_channel") |
| if icmod is not None: |
| try: |
| iris_sub = icmod.retrieve(query, client, inserted_before, |
| trace=trace) |
| if iris_sub: |
| have = {c for c, _ in sub} |
| merged = list(iris_sub) |
| merged += [(c, e) for c, e in sub |
| if c not in {x for x, _ in iris_sub}] |
| sub = merged[:MAX_RESULTS] |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"iris_channel:{exc!r}") |
|
|
| |
| |
| |
| |
| |
| if sub and os.environ.get("PFBMAX_LISTWISE", "").strip() in ("1", "true", "yes"): |
| lwmod = _sib("listwise") |
| if lwmod is not None: |
| try: |
| texts = {} |
| for cid, ev in sub: |
| texts[cid] = ev |
| order = lwmod.rerank( |
| query, criteria, [(c, texts.get(c, "")) for c, _ in sub], |
| meter=getattr(llm, "meter", None), |
| trace=trace if trace is not None else None) |
| pos = {c: i for i, c in enumerate(order)} |
| sub = sorted(sub, key=lambda r: pos.get(r[0], 10**6)) |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"listwise:{exc!r}") |
| |
| |
| |
| |
| |
| if sub and os.environ.get("PFBMAX_PERCRIT_URL", "").strip(): |
| pcmod = _sib("percrit_rank") |
| if pcmod is not None: |
| try: |
| ptexts = {cid: ev for cid, ev in sub} |
| porder = pcmod.rerank( |
| query, criteria, |
| [(c, ptexts.get(c, "")) for c, _ in sub], |
| trace=trace if trace is not None else None) |
| ppos = {c: i for i, c in enumerate(porder)} |
| sub = sorted(sub, key=lambda r: ppos.get(r[0], 10**6)) |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"percrit:{exc!r}") |
| |
| |
| |
| |
| |
| |
| if sub and os.environ.get("PFBMAX_CRITJUDGE", "").strip() in ("1", "true", "yes"): |
| cjmod = _sib("criterion_judge") |
| if cjmod is not None: |
| try: |
| |
| |
| |
| |
| |
| |
| def _judge_text(cid: str, ev: str) -> str: |
| try: |
| title = pool.title(cid) or "" |
| parts = pool.texts(cid) or ([ev] if ev else []) |
| except Exception: |
| title, parts = "", ([ev] if ev else []) |
| body = " ".join(p for p in parts if p) |
| return (title + ". " + body).strip() if title else body |
|
|
| order = cjmod.rerank( |
| query, criteria, |
| [(c, _judge_text(c, ev)) for c, ev in sub], |
| meter=getattr(llm, "meter", None), |
| trace=trace if trace is not None else None) |
| pos = {c: i for i, c in enumerate(order)} |
| sub = sorted(sub, key=lambda r: pos.get(r[0], 10**6)) |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"critjudge:{exc!r}") |
| |
| |
| |
| |
| |
| if sub: |
| return sub |
| except Exception as exc: |
| if trace is not None: |
| trace.setdefault("errors", []).append(f"semantic:{exc!r}") |
| |
| |
| else: |
| _report_once( |
| "stub:semantic", |
| "[pfbmax.router] semantic_retrieval/semantic_judge not present -- " |
| "using built-in degraded snippet+paper fallback (STUB)", |
| ) |
| return _fallback_semantic(query, client, llm, inserted_before) |
|
|
|
|
| def _specific_channel(query, client, llm, inserted_before, trace=None) -> Submission: |
| m = _sib("specific_solver") |
| if m is not None: |
| return m.solve_specific(query, client, llm, inserted_before) |
| _report_once( |
| "stub:specific", |
| "[pfbmax.router] specific_solver not present -- using built-in " |
| "title-resolve STUB (single emit)", |
| ) |
| return _fallback_specific(query, client, llm, inserted_before) |
|
|
|
|
| def _metadata_channel(query, client, llm, inserted_before, trace=None) -> Submission: |
| m = _sib("metadata_solver") |
| if m is not None: |
| return m.solve_metadata(query, client, llm, inserted_before) |
| _report_once( |
| "stub:metadata", |
| "[pfbmax.router] metadata_solver not present -- STUB returns [] " |
| "(the semantic fallback will fire)", |
| ) |
| return [] |
|
|
|
|
| |
| |
| |
|
|
| def solve(query: str, client, llm, inserted_before: str | None, |
| *, trace: dict | None = None) -> Submission: |
| """Route + solve one PFB query. Never raises. |
| |
| route via :func:`classify` (one LLM json call + deterministic guard), |
| then dispatch; any exception falls back to the semantic channel, as does |
| an EMPTY result from a deterministic channel (no-empty-fallback: an |
| empty submission scores exactly 0, so the judged channel can only help); |
| total failure returns []. ``trace`` (optional, keyword-only) collects |
| route / fallback / pool diagnostics for the iteration engine. |
| """ |
| route = "semantic" |
| results: Submission = [] |
| try: |
| route = classify(query, llm) |
| if trace is not None: |
| trace["route"] = route |
| if route == "specific": |
| results = _finalize( |
| _specific_channel(query, client, llm, inserted_before, trace)) |
| elif route == "metadata": |
| results = _finalize( |
| _metadata_channel(query, client, llm, inserted_before, trace)) |
| else: |
| results = _finalize( |
| _semantic_channel(query, client, llm, inserted_before, trace)) |
| except Exception as exc: |
| if trace is not None: |
| trace["error"] = repr(exc) |
| results = [] |
| if not results and route != "semantic": |
| if trace is not None: |
| trace["fallback"] = "semantic" |
| try: |
| results = _finalize( |
| _semantic_channel(query, client, llm, inserted_before, trace)) |
| except Exception as exc: |
| if trace is not None: |
| trace["fallback_error"] = repr(exc) |
| results = [] |
| return results |
|
|
|
|
| __all__ = [ |
| "MAX_RESULTS", |
| "ROUTE_MARKER", |
| "CRITERIA_MARKER", |
| "ROUTE_SYSTEM", |
| "CRITERIA_SYSTEM", |
| "Submission", |
| "classify", |
| "derive_criteria", |
| "solve", |
| ] |
|
|