raglab 0.2.2__tar.gz → 0.2.4__tar.gz

{raglab-0.2.2 → raglab-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: raglab
-Version: 0.2.2
+Version: 0.2.4
 Summary: A medley of tools to make RAG-based applications.
 Project-URL: Homepage, https://github.com/thorwhalen/raglab
 Project-URL: Repository, https://github.com/thorwhalen/raglab
@@ -9,7 +9,7 @@ Author: thorwhalen
 License: mit
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: ir>=0.1.12
+Requires-Dist: ir>=0.1.16
 Provides-Extra: dev
 Requires-Dist: pytest-cov>=4.0; extra == 'dev'
 Requires-Dist: pytest>=7.0; extra == 'dev'

{raglab-0.2.2 → raglab-0.2.4}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "raglab"
-version = "0.2.2"
+version = "0.2.4"
 description = "A medley of tools to make RAG-based applications."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -18,7 +18,7 @@ authors = [
 # strategies (Planner/Formulator/Evaluator) use `oa` lazily — the `llm` extra
 # below — so `import raglab` stays offline by default.
 dependencies = [
-    "ir>=0.1.12",
+    "ir>=0.1.16",
 ]
 
 [project.license]

{raglab-0.2.2 → raglab-0.2.4}/raglab/__init__.py

@@ -42,11 +42,14 @@ from .agent import (
     identity_citer,
     identity_formulator,
     ir_sources,
+    make_rrf_reranker,
     make_search_agent,
     passthrough_evaluator,
+    rrf_reranker,
     score_reranker,
     single_subtask_planner,
 )
+from .llm import EVALUATION_PROMPT, make_llm_evaluator, make_llm_formulator
 
 __all__ = [
     "Query",
@@ -68,5 +71,10 @@ __all__ = [
     "identity_formulator",
     "passthrough_evaluator",
     "score_reranker",
+    "rrf_reranker",
+    "make_rrf_reranker",
     "identity_citer",
+    "make_llm_formulator",
+    "make_llm_evaluator",
+    "EVALUATION_PROMPT",
 ]

{raglab-0.2.2 → raglab-0.2.4}/raglab/agent.py

@@ -6,6 +6,14 @@ This module is the v1 foundation: the immutable value types, the role *Protocols
 is fully parametrized by injected roles. Concrete tools live at the leaves — an
 `ir` corpus becomes one `Retriever` via :func:`ir.as_retriever`.
 
+The agent is **multi-source by default**: the loop stamps each hit's
+provenance (``hit.source``), and the fan-in :class:`Reranker` —
+:func:`rrf_reranker` — merges heterogeneous sources by *rank*, never by raw
+score (scores from different corpora / embedders / modes are incommensurable;
+ir_07/ir_08). Raw magnitudes order and dedup hits *within* one source, and the
+loop's pool always carries them; fused (ordinal) scores appear only at the
+fan-in boundary.
+
 The shape follows ir_09 §3/§6: a small set of named roles —
 ``Planner / Formulator / Retriever / Evaluator / Reranker / Citer`` — and a loop
 whose defining feature is the **back-edge** (evaluator → reformulate) that makes
@@ -27,9 +35,12 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import Any, Protocol, runtime_checkable
 
-# ir owns the retrieval substrate: the Result type and the Retriever leaf
-# contract live there (one-way dependency, ir is the SSOT).
-from ir import Retriever, SearchHit
+# ir owns the retrieval substrate: the Result type, the Retriever leaf
+# contract, and the hit operations (dedup, cross-source fusion) live there
+# (one-way dependency, ir is the SSOT).
+from ir import Retriever, SearchHit, fuse_hits, tag_source
+from ir.base import best_per_artifact
+from ir.retrieve import DFLT_RRF_K, Identity
 
 #: A retrieved item — ir's :class:`~ir.base.SearchHit` (ir_09's ``Result``):
 #: a *pointer + snippet* (``text``) with a ``score`` and ``metadata``.
@@ -55,6 +66,8 @@ __all__ = [
     "identity_formulator",
     "passthrough_evaluator",
     "score_reranker",
+    "rrf_reranker",
+    "make_rrf_reranker",
     "identity_citer",
 ]
 
@@ -188,13 +201,82 @@ def passthrough_evaluator(task: SubTask, results: Sequence[Result]) -> Judgement
 
 
 def score_reranker(results: Sequence[Result]) -> Sequence[Result]:
-    """Final ordering by descending ``score`` (the cross-source merge, v1).
+    """Magnitude merge: one surface per artifact, ordered by descending raw score.
+
+    Delegates to :func:`ir.base.best_per_artifact` (ir is the SSOT for hit
+    operations): an artifact retrieved by several queries / rounds — common once
+    the back-edge re-queries — survives once, at its highest score. Identity is
+    ``(source, artifact_id)``, so two sources' same-id artifacts never collapse.
+
+    A plain score sort compares raw scores **across** sources, which is only
+    sound when every source shares one score scale (same embedder + mode) — it
+    is the explicit homogeneous-sources opt-in. The default fan-in is
+    :func:`rrf_reranker`, which never compares raw scores across sources.
+    """
+    return best_per_artifact(results)
+
+
+def _rrf_rerank(
+    results: Sequence[Result],
+    *,
+    rrf_k: int,
+    weights: Mapping[str, float] | None,
+    identity: Identity,
+) -> Sequence[Result]:
+    """Group by ``hit.source`` and rank-fuse via :func:`ir.fuse_hits`."""
+    # None is preserved as the untagged pseudo-source key: its hits fuse as
+    # one rank group and stay unattributed (never an empty-string stamp).
+    groups: dict[str | None, list[Result]] = {}
+    for h in results:
+        groups.setdefault(h.source, []).append(h)
+    if len(groups) <= 1:
+        # One scale: the magnitude merge, with hits passed through untouched.
+        return best_per_artifact(results)
+    return fuse_hits(groups, rrf_k=rrf_k, weights=weights, identity=identity)
+
+
+def rrf_reranker(results: Sequence[Result]) -> Sequence[Result]:
+    """Cross-source merge (the default fan-in): fuse by rank, never by raw score.
+
+    Groups the accumulated pool by ``hit.source`` and delegates the merge to
+    :func:`ir.fuse_hits` (ir is the SSOT for hit operations): within each
+    source raw scores order and dedup that source's hits — one scale, sound —
+    and across sources only **ranks** interact (Reciprocal Rank Fusion), so
+    heterogeneous embedders / modes can never mis-order the merge, and
+    colliding ``artifact_id``\\ s from different sources stay distinct results
+    (identity is ``(source, artifact_id)``).
+
+    A single-source pool (or an untagged one — hits with no ``source``) keeps
+    its raw scores and exactly :func:`score_reranker`'s ordering; fused,
+    rank-derived scores only appear when there is genuinely something to fuse.
+    Each fused hit keeps its pre-fusion magnitude as
+    ``metadata["source_score"]``. For per-source weights, another ``rrf_k``,
+    or opt-in cross-source duplicate merging, use :func:`make_rrf_reranker`.
+    """
+    return _rrf_rerank(results, rrf_k=DFLT_RRF_K, weights=None, identity=None)
+
 
-    Note: a plain score sort assumes comparable score scales across sources
-    (true when they share an embedder + mode). A rank-based (RRF) cross-source
-    merge for heterogeneous backends is a documented follow-up.
+def make_rrf_reranker(
+    *,
+    rrf_k: int = DFLT_RRF_K,
+    weights: Mapping[str, float] | None = None,
+    identity: Identity = None,
+) -> Reranker:
+    """A parametrized :func:`rrf_reranker` (per-source trust weights, ``rrf_k``).
+
+    Args:
+        rrf_k: the RRF rank constant (standard default 60).
+        weights: optional per-source trust dial, by source name (default 1.0
+            each) — biases the merge without ever comparing raw scores.
+        identity: opt-in cross-source duplicate detection (e.g. ``"pointer"``)
+            — see :data:`ir.retrieve.Identity`. Default: never merge across
+            sources.
     """
-    return sorted(results, key=lambda r: float(getattr(r, "score", 0.0)), reverse=True)
+
+    def reranker(results: Sequence[Result]) -> Sequence[Result]:
+        return _rrf_rerank(results, rrf_k=rrf_k, weights=weights, identity=identity)
+
+    return reranker
 
 
 def identity_citer(results: Sequence[Result]) -> Sequence[Result]:
@@ -221,7 +303,7 @@ class SingleContextAgent:
     planner: Planner = single_subtask_planner
     formulator: Formulator = identity_formulator
     evaluator: Evaluator = passthrough_evaluator
-    reranker: Reranker = score_reranker
+    reranker: Reranker = rrf_reranker
     citer: Citer = identity_citer
     budget: Budget = field(default_factory=Budget)
 
@@ -243,7 +325,14 @@ class SingleContextAgent:
                 if retriever is None:
                     continue
                 for llq in self.formulator(current, source):
-                    found.extend(retriever(llq.query, **dict(llq.params)))
+                    # Stamp the registry key on hits the retriever did not
+                    # self-attribute (ir-backed retrievers stamp the corpus
+                    # name themselves, and their tags win), so any custom
+                    # Retriever still yields attributable hits — the fan-in
+                    # reranker merges by source.
+                    found.extend(
+                        tag_source(retriever(llq.query, **dict(llq.params)), source)
+                    )
             judged = self.evaluator(current, found[: self.budget.max_results_per_task])
             found = list(judged.relevant)
             if judged.sufficient or judged.refinement is None:
@@ -265,17 +354,22 @@ def make_search_agent(
     """Build a :class:`SingleContextAgent` over *sources* with smart defaults.
 
     ``sources`` is a ``Mapping[name, Retriever]`` — e.g.
-    ``{"skills": ir.as_retriever("skills")}``. Every role defaults to its no-LLM
-    thin-slice implementation, so ``make_search_agent(sources)("query")`` just
-    works; inject an LLM ``formulator`` / ``evaluator`` to turn on rewriting and
-    the back-edge.
+    ``{"skills": ir.as_retriever("skills")}``, :func:`ir_sources`, or the lazy
+    ``ir.retrievers()`` view. Every role defaults to its no-LLM thin-slice
+    implementation, so ``make_search_agent(sources)("query")`` just works;
+    inject an LLM ``formulator`` / ``evaluator`` to turn on rewriting and the
+    back-edge.
+
+    A custom ``Retriever`` must return :class:`ir.SearchHit` instances (the
+    ``Result`` alias): the loop stamps provenance (``hit.source``) on its
+    output, so duck-typed hit objects raise at the tagging step.
     """
     return SingleContextAgent(
         sources=sources,
         planner=planner or single_subtask_planner,
         formulator=formulator or identity_formulator,
         evaluator=evaluator or passthrough_evaluator,
-        reranker=reranker or score_reranker,
+        reranker=reranker or rrf_reranker,
         citer=citer or identity_citer,
         budget=budget or Budget(),
     )
@@ -284,9 +378,11 @@ def make_search_agent(
 def ir_sources(*names: str, **search_defaults: Any) -> dict[str, Retriever]:
     """A source registry ``{name: Retriever}`` backed by named ``ir`` corpora.
 
-    Each name is bound to ``ir.as_retriever(name, **search_defaults)``. A thin
-    convenience; once ir ships ``registry.retrievers()`` (a lazy view), prefer
-    that. ``search_defaults`` (e.g. ``mode="hybrid"``) apply to every source.
+    Each name is bound to ``ir.as_retriever(name, **search_defaults)``, opened
+    eagerly. For the *lazy* live view over everything registered (a corpus
+    opens only when its key is first used), use ``ir.retrievers()`` instead —
+    the agent accepts either, or any ``Mapping[name, Retriever]``.
+    ``search_defaults`` (e.g. ``mode="hybrid"``) apply to every source.
     """
     import ir
 

raglab-0.2.4/raglab/llm.py

@@ -0,0 +1,297 @@
+"""LLM-backed roles — the Formulator and Evaluator (ir_09 §8 steps 1 & 3).
+
+`raglab`'s no-LLM thin slice (:mod:`raglab.agent`) runs offline: an identity
+formulator and a pass-through evaluator. This module supplies the two LLM roles
+that turn that slice into a real agent — query understanding and the back-edge:
+
+- :func:`make_llm_formulator` — query rewrite / expand / HyDE. A thin adapter
+  over ir's :func:`ir.make_llm_formulator` (the ``formulate=`` seam, ir_09 §3):
+  ir owns the lazy-:mod:`oa` rewriter and the identity fallback; raglab only
+  adapts the shape ``str -> str | [str]`` to the role contract
+  ``(SubTask, source) -> [LowLevelQuery]``.
+- :func:`make_llm_evaluator` — sufficiency + refinement. ``ir.select`` owns the
+  *relevance* decision (the calibrated committed subset, ir_09 §3); the LLM owns
+  *sufficiency* — informed by ir's model-free :attr:`ir.Selection.sufficient`
+  hint, it judges whether the committed subset actually satisfies the goal and,
+  when it does not, emits a ``refinement`` SubTask. That refinement is the
+  **back-edge** that makes the loop an agent rather than a DAG.
+
+Two load-bearing boundaries (ir ↔ raglab), guarded here:
+
+1. A Formulator returns **queries, never SubTasks** — decomposition is the
+   Planner's job. :func:`make_llm_formulator` only ever yields
+   :class:`~raglab.agent.LowLevelQuery`\\ s.
+2. The **back-edge lives in raglab**, never in ir: ir derives a ``sufficient``
+   *signal* from its own selection; raglab's Evaluator is what reads it, decides,
+   and re-queries.
+
+Both builders mirror :func:`ir.select.make_llm_selector`: an injectable callable
+(a test double, or your own router), built lazily on :mod:`oa` only when omitted
+(so ``import raglab`` stays offline), with a **safe fallback** — a formulator must
+never make retrieval worse than the raw query, and an evaluator must never
+fabricate an endless loop (on any failure it returns no refinement, which is the
+loop's break condition).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any, Callable
+
+import ir
+from ir.base import best_per_artifact
+
+from .agent import (
+    Evaluator,
+    Formulator,
+    Judgement,
+    LowLevelQuery,
+    Reranker,
+    Result,
+    SubTask,
+    rrf_reranker,
+)
+
+__all__ = ["make_llm_formulator", "make_llm_evaluator", "EVALUATION_PROMPT"]
+
+#: An ir-style query formulator: a query string -> one query or several.
+QueryFormulator = Callable[[str], "str | Sequence[str]"]
+
+#: An evaluator judge: ``(goal, results) -> (sufficient, refinement)``.
+#: ``sufficient`` ends the loop; a non-empty ``refinement`` query becomes the
+#: next sub-goal (the back-edge). A raw text reply is also accepted and parsed.
+Judge = Callable[..., "tuple[bool, str | None] | str"]
+
+#: Truncate each rendered result's text to this many chars in the judge prompt.
+DFLT_MAX_RESULT_CHARS = 500
+
+
+# --------------------------------------------------------------------------- #
+# LLM Formulator — adapt ir's query-level formulator to the role contract
+# --------------------------------------------------------------------------- #
+
+
+def make_llm_formulator(
+    *,
+    formulate: QueryFormulator | None = None,
+    params: Mapping[str, Any] | None = None,
+    **make_kwargs: Any,
+) -> Formulator:
+    """An LLM-backed raglab :class:`~raglab.agent.Formulator`.
+
+    Adapts an ir-style query formulator (``str -> str | [str]``) to the role
+    contract ``(SubTask, source) -> [LowLevelQuery]``: the sub-goal text is
+    rewritten / expanded into one or more search queries, each wrapped as a
+    :class:`~raglab.agent.LowLevelQuery` against ``source``.
+
+    Args:
+        formulate: an injectable ir-style formulator (a test double, or one you
+            built). When omitted it is built once via
+            :func:`ir.make_llm_formulator` (``**make_kwargs`` are forwarded —
+            e.g. ``n``, ``prompt``, ``rewriter``); that call is offline (``oa`` is
+            imported lazily only when the formulator is *invoked*), so
+            ``import raglab`` stays offline.
+        params: per-call retriever overrides (e.g. ``{"mode": "hybrid", "k": 5}``)
+            attached to every emitted :class:`~raglab.agent.LowLevelQuery`.
+
+    The boundary: this returns **queries, never SubTasks**. ir's formulator
+    already falls back to identity on any failure, so the emitted list is never
+    empty — a formulator must never make retrieval worse than the raw sub-goal.
+    """
+    extra = dict(params or {})
+    fn: QueryFormulator = (
+        formulate if formulate is not None else ir.make_llm_formulator(**make_kwargs)
+    )
+
+    def formulator(task: SubTask, source: str) -> list[LowLevelQuery]:
+        try:
+            out = fn(task.goal)
+        except Exception:
+            out = None  # a formulator must never make retrieval worse: fall back
+        queries = [out] if isinstance(out, str) else list(out or [])
+        queries = [q for q in queries if isinstance(q, str) and q.strip()] or [
+            task.goal
+        ]
+        return [
+            LowLevelQuery(source=source, query=q, params=dict(extra)) for q in queries
+        ]
+
+    return formulator
+
+
+# --------------------------------------------------------------------------- #
+# LLM Evaluator — ir.select owns relevance; the LLM owns sufficiency + the
+# back-edge (ir_09 §3/§4)
+# --------------------------------------------------------------------------- #
+
+#: Default prompt for :func:`make_llm_evaluator` — judge sufficiency, else emit a
+#: single improved query (the refinement that drives the back-edge).
+EVALUATION_PROMPT = """\
+A search agent is pursuing this goal:
+
+{goal}
+
+A calibrated selector reviewed the retrieved candidates and committed to the
+results below (it abstained if this list is empty):
+
+{results}
+
+Decide whether these results are SUFFICIENT to satisfy the goal.
+- If they are, reply with exactly: SUFFICIENT
+- If they are not, reply with: INSUFFICIENT
+  then, on the next line, a single improved search query that would retrieve what
+  is still missing. Keep it a terse search phrase — no prose, no numbering.
+"""
+
+
+def make_llm_evaluator(
+    *,
+    judge: Judge | None = None,
+    select_strategy: str = "conservative",
+    select_kwargs: Mapping[str, Any] | None = None,
+    prerank: Reranker | None = None,
+    prompt: str = EVALUATION_PROMPT,
+    max_result_chars: int = DFLT_MAX_RESULT_CHARS,
+    **prompt_function_kwargs: Any,
+) -> Evaluator:
+    """An LLM-backed raglab :class:`~raglab.agent.Evaluator` (turns on the back-edge).
+
+    Relevance is ir's: each round the accumulated results are passed through
+    :func:`ir.select` and the committed subset becomes the ``Judgement.relevant``
+    (so the LLM stays in its lane — LLM relevance is known-fragile, ir_01 §3).
+    Sufficiency is the LLM's: it reads the committed subset (informed by ir's
+    model-free :attr:`~ir.Selection.sufficient` hint via abstention) and decides
+    whether the goal is satisfied. When it is not, the judge's improved query
+    becomes a ``refinement`` :class:`~raglab.agent.SubTask` over the same sources
+    — the **back-edge**.
+
+    Args:
+        judge: an injectable ``(goal, results) -> (sufficient, refinement)``
+            callable (a test double, or your own router); a raw text reply is also
+            accepted and parsed. When omitted it is built lazily on :mod:`oa`.
+        select_strategy: ir selection strategy for the relevance decision
+            (default ``"conservative"`` — distractor-robust).
+        select_kwargs: extra args forwarded to :func:`ir.select`
+            (e.g. ``max_k``, ``rel``). ``min_score`` is allowed only while the
+            round's pool keeps raw scores (single-source, or
+            ``prerank=score_reranker``): an absolute floor is per-(corpus,
+            mode, embedder), so a round whose pool was rank-fused across
+            sources **raises** rather than silently mass-abstaining on
+            ordinal RRF scores.
+        prerank: the per-round merge that puts the accumulated pool best-first
+            (and deduped) before ``ir.select``. Defaults to
+            :func:`~raglab.agent.rrf_reranker` — a round's pool can already mix
+            sources (a SubTask spans several), so the same rank-based,
+            scale-safe merge as the fan-in applies; single-source rounds keep
+            raw scores. Inject :func:`~raglab.agent.score_reranker` for
+            sources known to share one score scale.
+        max_result_chars: per-result text truncation in the judge prompt.
+
+    Safe fallback: any judge error returns ``refinement=None`` — the control
+    loop's break condition — so a failing judge can never fabricate an endless
+    loop. Sufficiency without a refinement query is likewise treated as a stop.
+    """
+    sel_kw = dict(select_kwargs or {})
+    # The floor is handled here, not blindly forwarded: it must only ever meet
+    # raw (per-source-scale) scores — see the guard in the evaluator below.
+    floor = sel_kw.pop("min_score", None)
+    prerank = prerank or rrf_reranker
+
+    def _ask_judge(goal: str, rendered: str) -> tuple[bool, str | None]:
+        fn = (
+            judge
+            if judge is not None
+            else _default_llm_judge(prompt, **prompt_function_kwargs)
+        )
+        return _normalize_verdict(fn(goal=goal, results=rendered))
+
+    def evaluator(task: SubTask, results: Sequence[Result]) -> Judgement:
+        # ``ir.select`` documents a best-first precondition; the loop accumulates
+        # hits across rounds/sources in arbitrary order, so merge them first —
+        # the same (scale-safe) merge the final fan-in reranker uses.
+        ranked = list(prerank(results))
+        if floor is not None and any("source_rank" in h.metadata for h in ranked):
+            # fuse_hits stamps source_rank exactly when it fused: the pool
+            # spans sources, so the scores below are ordinal RRF values — an
+            # absolute floor against them is the mis-scaled comparison ir
+            # refuses loudly (ir_07). Never silently mass-abstain.
+            raise ValueError(
+                "min_score with a multi-source round: the pool was rank-fused, "
+                "so an absolute floor would compare against ordinal RRF scores. "
+                "Floors are per-(corpus, mode, embedder) — drop min_score, or "
+                "inject prerank=score_reranker for sources known to share one "
+                "score scale."
+            )
+        selection = ir.select(
+            ranked, strategy=select_strategy, min_score=floor, **sel_kw
+        )
+        # Map the committed subset back to the PRE-fusion originals: fused
+        # (ordinal) scores are local to this selection. ``Judgement.relevant``
+        # re-enters the loop's pool and the final fan-in re-fuses it, so it
+        # must carry raw per-source magnitudes — otherwise round N+1 compares
+        # round N's fused scores against raw ones inside one source group,
+        # and ``source_score`` gets overwritten by an already-fused value.
+        raw = {(h.source, h.artifact_id): h for h in best_per_artifact(results)}
+        relevant = [raw.get((h.source, h.artifact_id), h) for h in selection.selected]
+        rendered = (
+            _render_results(relevant, max_result_chars)
+            if relevant
+            else "(none — the selector abstained)"
+        )
+        try:
+            sufficient, refinement = _ask_judge(task.goal, rendered)
+        except Exception:
+            # Trust ir's model-free signal for the report, but never re-query on a
+            # judge failure: refinement=None is the loop's break condition.
+            return Judgement(
+                relevant=relevant, sufficient=selection.sufficient, refinement=None
+            )
+        if sufficient or not refinement:
+            return Judgement(relevant=relevant, sufficient=True, refinement=None)
+        return Judgement(
+            relevant=relevant,
+            sufficient=False,
+            refinement=SubTask(goal=refinement, sources=task.sources),
+        )
+
+    return evaluator
+
+
+def _render_results(results: Sequence[Result], max_chars: int) -> str:
+    """Render committed hits as ``- id: text`` lines for the judge prompt."""
+    return "\n".join(f"- {h.artifact_id}: {str(h.text)[:max_chars]}" for h in results)
+
+
+def _normalize_verdict(out: "tuple[bool, str | None] | str") -> tuple[bool, str | None]:
+    """Coerce a judge reply (a ``(sufficient, refinement)`` tuple, or raw text)."""
+    if isinstance(out, tuple):
+        sufficient, refinement = out
+        refinement = str(refinement).strip() if refinement else None
+        return bool(sufficient), (refinement or None)
+    return _parse_verdict(str(out or ""))
+
+
+def _parse_verdict(text: str) -> tuple[bool, str | None]:
+    """Parse a SUFFICIENT / INSUFFICIENT + refinement-query text reply."""
+    lines = [line.strip() for line in text.splitlines() if line.strip()]
+    if not lines or lines[0].upper().startswith("SUFFICIENT"):
+        return True, None
+    refinement = lines[1] if len(lines) > 1 else None
+    return False, (refinement or None)
+
+
+def _default_llm_judge(prompt: str, **prompt_function_kwargs: Any) -> Judge:
+    """Build the default sufficiency judge on :mod:`oa` (lazy import)."""
+    import oa
+
+    fn = oa.prompt_function(
+        prompt,
+        egress=_parse_verdict,
+        name="evaluate_sufficiency",
+        **prompt_function_kwargs,
+    )
+
+    def judge(*, goal: str, results: str) -> tuple[bool, str | None]:
+        return fn(goal=goal, results=results)
+
+    return judge

{raglab-0.2.2 → raglab-0.2.4}/tests/test_agent.py

@@ -39,6 +39,8 @@ def test_defaults_satisfy_role_protocols():
     assert isinstance(raglab.identity_formulator, raglab.Formulator)
     assert isinstance(raglab.passthrough_evaluator, raglab.Evaluator)
     assert isinstance(raglab.score_reranker, raglab.Reranker)
+    assert isinstance(raglab.rrf_reranker, raglab.Reranker)
+    assert isinstance(raglab.make_rrf_reranker(), raglab.Reranker)
     assert isinstance(raglab.identity_citer, raglab.Citer)
 
 
@@ -57,13 +59,41 @@ def test_query_string_or_object_equivalent():
     assert agent("q") == agent(Query(text="q"))
 
 
-def test_cross_source_merge_reranks_by_score():
+def test_cross_source_merge_is_rank_based_not_score_based():
+    # Two sources' raw scores live on different scales, so the default fan-in
+    # (rrf_reranker) fuses by per-source rank: both hits are rank 1 in their own
+    # source, and the tie breaks by source order — never by the (incomparable)
+    # raw scores. For a raw-score merge, inject score_reranker explicitly.
     sources = {
         "s1": _fake_retriever(_hits(("a", 0.3))),
         "s2": _fake_retriever(_hits(("b", 0.9))),
     }
     results = make_search_agent(sources)("q")
-    assert [r.artifact_id for r in results] == ["b", "a"]  # by score desc
+    assert [r.artifact_id for r in results] == ["a", "b"]  # source order, not score
+    by_score = make_search_agent(sources, reranker=raglab.score_reranker)("q")
+    assert [r.artifact_id for r in by_score] == ["b", "a"]  # the explicit opt-in
+
+
+def test_final_results_have_no_duplicate_artifacts():
+    # An artifact re-retrieved across back-edge rounds collapses to its best score.
+    retr = _fake_retriever(_hits(("a", 0.9), ("b", 0.4)))
+    rounds = {"n": 0}
+
+    def refining(task, results):
+        rounds["n"] += 1
+        if rounds["n"] < 2:
+            return Judgement(
+                list(results),
+                sufficient=False,
+                refinement=SubTask(task.goal, task.sources),
+            )
+        return Judgement(list(results), sufficient=True)
+
+    ids = [
+        r.artifact_id for r in make_search_agent({"s": retr}, evaluator=refining)("q")
+    ]
+    assert ids == ["a", "b"]  # best score per artifact, descending
+    assert len(ids) == len(set(ids))  # the re-query did not duplicate artifacts
 
 
 def test_passthrough_evaluator_does_not_loop():
@@ -133,6 +163,26 @@ def test_budget_bounds_a_never_sufficient_loop():
     assert len(retr.calls) == 3  # exactly max_rounds — the safety net holds
 
 
+def test_budget_caps_results_per_task_seen_by_evaluator():
+    retr = _fake_retriever(_hits(*[(f"a{i}", 1.0 - i / 100) for i in range(20)]))
+    seen = {}
+
+    def evaluator(task, results):
+        seen["n"] = len(results)
+        return Judgement(list(results), sufficient=True)
+
+    make_search_agent(
+        {"s": retr}, evaluator=evaluator, budget=Budget(max_results_per_task=5)
+    )("q")
+    assert seen["n"] == 5  # the evaluator sees at most max_results_per_task
+
+
+def test_budget_caps_sources_per_task():
+    retrs = {f"s{i}": _fake_retriever(_hits((f"a{i}", 0.5))) for i in range(6)}
+    make_search_agent(retrs, budget=Budget(max_sources_per_task=2))("q")
+    assert sum(1 for r in retrs.values() if r.calls) == 2  # only first 2 sources hit
+
+
 # ----- end-to-end over a REAL ir corpus (hermetic: light embedder) ---------- #
 
 

raglab-0.2.4/tests/test_fanin.py

@@ -0,0 +1,288 @@
+"""Tests for the Reranker at fan-in — rank-based cross-source merge (ir_09 §3).
+
+The property under test: raw scores never cross a source boundary. Within one
+source they order and dedup that source's hits; across sources only ranks
+interact (RRF via ``ir.fuse_hits``). Hermetic: fake retrievers with canned
+hits; one end-to-end test wires two REAL ir corpora (light embedder, in-memory
+stores) whose artifact ids deliberately collide.
+"""
+
+import pytest
+
+import raglab
+from ir import SearchHit
+from raglab import (
+    Judgement,
+    SubTask,
+    make_rrf_reranker,
+    make_search_agent,
+    rrf_reranker,
+)
+
+
+def _hits(*specs, source=None):
+    """``(artifact_id, score)`` pairs -> ir.SearchHits (optionally source-tagged)."""
+    return [
+        SearchHit(aid, "k", score, f"text {aid}", {}, source) for aid, score in specs
+    ]
+
+
+def _fake_retriever(hits):
+    """A Retriever that records its calls and returns canned hits."""
+    calls = []
+
+    def retrieve(query, **kw):
+        calls.append((query, kw))
+        return list(hits)
+
+    retrieve.calls = calls
+    return retrieve
+
+
+# ----- rrf_reranker: the role in isolation ---------------------------------- #
+
+
+def test_heterogeneous_scales_interleave_by_rank():
+    # A cosine-scale source (~[0,1]) and a BM25-scale source (~[0,50]): a raw
+    # score sort would bury the cosine source entirely; rank fusion interleaves.
+    pool = _hits(("c1", 0.92), ("c2", 0.85), source="cos") + _hits(
+        ("b1", 31.0), ("b2", 24.0), source="bm25"
+    )
+    fused = rrf_reranker(pool)
+    assert {h.artifact_id for h in fused[:2]} == {"c1", "b1"}  # both rank-1s lead
+    assert {h.artifact_id for h in fused[2:]} == {"c2", "b2"}
+
+
+def test_colliding_ids_across_sources_stay_distinct():
+    pool = _hits(("dol", 0.9), source="skills") + _hits(("dol", 28.0), source="pkgs")
+    fused = rrf_reranker(pool)
+    assert len(fused) == 2  # same id, different corpus = different artifact
+    assert {h.source for h in fused} == {"skills", "pkgs"}
+
+
+def test_single_source_pool_keeps_raw_scores():
+    pool = _hits(("a", 0.9), ("b", 0.5), source="s")
+    fused = rrf_reranker(pool)
+    assert [h.score for h in fused] == [0.9, 0.5]  # = score_reranker's ordering
+    assert list(fused) == list(raglab.score_reranker(pool))
+
+
+def test_untagged_pool_keeps_raw_scores_and_no_stamp():
+    pool = _hits(("a", 0.9), ("b", 0.5))  # no source anywhere
+    fused = rrf_reranker(pool)
+    assert [h.score for h in fused] == [0.9, 0.5]
+    assert all(h.source is None for h in fused)  # passthrough, no "" stamping
+
+
+def test_cross_round_duplicates_collapse_per_source():
+    # The same (source, artifact) re-retrieved across back-edge rounds counts
+    # once, at its best rank — no duplicate ids, no double RRF mass.
+    pool = (
+        _hits(("a", 0.7), ("a", 0.9), source="s1")
+        + _hits(("z", 5.0), source="s2")
+        + _hits(("a", 0.8), source="s1")
+    )
+    fused = rrf_reranker(pool)
+    keyed = [(h.source, h.artifact_id) for h in fused]
+    assert len(keyed) == len(set(keyed)) == 2
+    a = next(h for h in fused if h.artifact_id == "a")
+    assert a.metadata["source_score"] == 0.9  # the best raw magnitude survives
+    assert a.metadata["source_rank"] == 1
+
+
+def test_make_rrf_reranker_weights_bias_the_merge():
+    pool = _hits(("a", 0.9), source="s1") + _hits(("b", 0.9), source="s2")
+    fused = make_rrf_reranker(weights={"s2": 2.0})(pool)
+    assert fused[0].artifact_id == "b"  # trust dial, no score comparability needed
+
+
+# ----- the agent loop: tagging + fan-in ------------------------------------- #
+
+
+def test_loop_stamps_registry_key_on_untagged_hits():
+    sources = {
+        "s1": _fake_retriever(_hits(("a", 0.9))),
+        "s2": _fake_retriever(_hits(("b", 7.0))),
+    }
+    results = make_search_agent(sources)("q")
+    assert {h.source for h in results} == {"s1", "s2"}
+
+
+def test_retriever_self_attribution_wins_over_registry_key():
+    # An ir-backed retriever stamps the corpus name itself; the loop must not
+    # overwrite it with the (possibly different) registry key.
+    sources = {"alias": _fake_retriever(_hits(("a", 0.9), source="corpus_x"))}
+    results = make_search_agent(sources)("q")
+    assert [h.source for h in results] == ["corpus_x"]
+
+
+def test_agent_keeps_colliding_ids_from_two_sources():
+    sources = {
+        "skills": _fake_retriever(_hits(("dol", 0.9))),
+        "pkgs": _fake_retriever(_hits(("dol", 28.0))),
+    }
+    results = make_search_agent(sources)("q")
+    assert sorted((h.source, h.artifact_id) for h in results) == [
+        ("pkgs", "dol"),
+        ("skills", "dol"),
+    ]
+
+
+def test_back_edge_rounds_do_not_duplicate_across_sources():
+    sources = {
+        "s1": _fake_retriever(_hits(("a", 0.9))),
+        "s2": _fake_retriever(_hits(("a", 6.0))),  # same id, other source
+    }
+    rounds = {"n": 0}
+
+    def refining(task, results):
+        rounds["n"] += 1
+        if rounds["n"] < 3:
+            return Judgement(
+                list(results),
+                sufficient=False,
+                refinement=SubTask(task.goal, task.sources),
+            )
+        return Judgement(list(results), sufficient=True)
+
+    results = make_search_agent(sources, evaluator=refining)("q")
+    keyed = [(h.source, h.artifact_id) for h in results]
+    assert sorted(keyed) == [("s1", "a"), ("s2", "a")]  # one per (source, artifact)
+
+
+# ----- the LLM evaluator's per-round prerank --------------------------------- #
+
+
+def test_llm_evaluator_prerank_is_scale_safe_by_default():
+    from raglab import make_llm_evaluator
+
+    # A round's pool already mixes sources: the per-round merge must keep both
+    # same-id artifacts (distinct sources) visible to ir.select and the judge.
+    pool = _hits(("readme", 0.9), source="s1") + _hits(("readme", 30.0), source="s2")
+    evaluator = make_llm_evaluator(judge=lambda *, goal, results: (True, None))
+    judgement = evaluator(SubTask("g", ("s1", "s2")), pool)
+    assert sorted((h.source, h.artifact_id) for h in judgement.relevant) == [
+        ("s1", "readme"),
+        ("s2", "readme"),
+    ]
+
+
+def test_llm_evaluator_prerank_is_injectable():
+    from raglab import make_llm_evaluator, score_reranker
+
+    pool = _hits(("readme", 0.9), source="s1") + _hits(("readme", 30.0), source="s2")
+    evaluator = make_llm_evaluator(
+        judge=lambda *, goal, results: (True, None), prerank=score_reranker
+    )
+    judgement = evaluator(SubTask("g", ("s1", "s2")), pool)
+    # The explicit magnitude opt-in keeps both too (identity is per source)…
+    # but ranks them by raw score: the BM25-scale hit wins the top slot.
+    assert judgement.relevant[0].source == "s2"
+
+
+def test_llm_evaluator_relevant_carries_prefusion_magnitudes():
+    from raglab import make_llm_evaluator
+
+    # Judgement.relevant re-enters the loop's pool, so it must carry RAW
+    # per-source scores — the fused (ordinal) view is local to selection.
+    pool = _hits(("a", 0.9), source="s1") + _hits(("b", 30.0), source="s2")
+    evaluator = make_llm_evaluator(judge=lambda *, goal, results: (True, None))
+    judgement = evaluator(SubTask("g", ("s1", "s2")), pool)
+    assert {h.score for h in judgement.relevant} == {0.9, 30.0}  # not ~1/61
+    assert all("source_rank" not in h.metadata for h in judgement.relevant)
+
+
+def test_agent_with_llm_evaluator_never_mixes_fused_and_raw_scores():
+    from raglab import make_llm_evaluator
+
+    # Regression for the double-fusion feedback: round-1 fused scores must not
+    # re-enter the pool, or round 2 compares ordinal RRF values against raw
+    # magnitudes INSIDE one source group (and source_score gets overwritten by
+    # an already-fused value).
+    def varying_retriever(rounds_hits):
+        state = {"round": 0}
+
+        def retrieve(query, **kw):
+            hits = rounds_hits[min(state["round"], len(rounds_hits) - 1)]
+            state["round"] += 1
+            return list(hits)
+
+        return retrieve
+
+    sources = {
+        "s1": varying_retriever([_hits(("a1", 0.9)), _hits(("a2", 0.5))]),
+        "s2": varying_retriever([_hits(("b1", 30.0)), _hits(("b2", 24.0))]),
+    }
+    verdicts = iter([(False, "refined q"), (True, None)])
+    evaluator = make_llm_evaluator(judge=lambda *, goal, results: next(verdicts))
+    results = make_search_agent(
+        sources, evaluator=evaluator, formulator=raglab.identity_formulator
+    )("q")
+    raw = {("s1", "a1"): 0.9, ("s1", "a2"): 0.5, ("s2", "b1"): 30.0, ("s2", "b2"): 24.0}
+    for h in results:
+        assert h.metadata["source_score"] == raw[(h.source, h.artifact_id)]
+    # Within one source, the final order follows raw magnitudes.
+    s1 = [h.artifact_id for h in results if h.source == "s1"]
+    assert s1 == sorted(s1, key=lambda a: -raw[("s1", a)])
+
+
+def test_llm_evaluator_min_score_works_single_source_raises_multi():
+    import pytest
+
+    from raglab import make_llm_evaluator
+
+    judge = lambda *, goal, results: (True, None)  # noqa: E731
+    single = _hits(("a", 0.9), ("b", 0.1), source="s1")
+    evaluator = make_llm_evaluator(judge=judge, select_kwargs={"min_score": 0.5})
+    judgement = evaluator(SubTask("g", ("s1",)), single)
+    assert [h.artifact_id for h in judgement.relevant] == ["a"]  # floor met raw scores
+
+    multi = _hits(("a", 0.9), source="s1") + _hits(("b", 30.0), source="s2")
+    with pytest.raises(ValueError, match="rank-fused"):
+        evaluator(SubTask("g", ("s1", "s2")), multi)  # never silently mass-abstain
+
+
+def test_mixed_tagged_untagged_pool_keeps_none_provenance():
+    pool = _hits(("u", 0.5)) + _hits(("t", 9.0), source="s1")
+    fused = rrf_reranker(pool)
+    by_id = {h.artifact_id: h for h in fused}
+    assert by_id["u"].source is None  # untagged pseudo-source, never ""
+    assert by_id["t"].source == "s1"
+    assert by_id["u"].to_dict()["source"] is None
+
+
+# ----- end-to-end over two REAL ir corpora (hermetic: light embedder) -------- #
+
+
+def test_agent_federates_two_real_corpora_with_colliding_ids():
+    import ir
+    from ir.store import CorpusStore
+
+    def corpus(docs, name):
+        return ir.build(
+            ir.CorpusSource.from_mapping(docs, name=name, strategy=ir.WholeText()),
+            store=CorpusStore.memory(),
+            embedder="light",
+        )
+
+    docs_a = {"shared": "zebra zephyr zucchini", "alpha": "alpha apple avocado"}
+    docs_b = {"shared": "zebra zephyr zucchini", "beta": "beta banana blueberry"}
+    agent = make_search_agent(
+        {
+            "one": ir.as_retriever(corpus(docs_a, "one"), k=3),
+            "two": ir.as_retriever(corpus(docs_b, "two"), k=3),
+        }
+    )
+    results = agent("zebra zephyr zucchini")
+    keyed = {(h.source, h.artifact_id) for h in results}
+    # The colliding "shared" artifact survives from BOTH corpora, attributed.
+    assert {("one", "shared"), ("two", "shared")} <= keyed
+    assert results[0].to_dict()["source"] in {"one", "two"}  # serialization-clean
+    fused_scores = [h.score for h in results]
+    assert fused_scores == sorted(fused_scores, reverse=True)  # best-first
+
+
+def test_score_reranker_remains_the_homogeneous_opt_in():
+    pytest.importorskip("ir")
+    pool = _hits(("a", 0.3), source="s1") + _hits(("b", 0.9), source="s2")
+    assert [h.artifact_id for h in raglab.score_reranker(pool)] == ["b", "a"]

raglab-0.2.4/tests/test_llm_roles.py

@@ -0,0 +1,250 @@
+"""Tests for raglab's LLM-backed roles (the Formulator and Evaluator).
+
+Hermetic and deterministic: every "LLM" is an injected test double (no model, no
+network). The Formulator adapts an ir-style ``str -> [str]`` rewriter; the
+Evaluator delegates relevance to ``ir.select`` and sufficiency to the injected
+judge. One end-to-end demo wires a real ``ir`` corpus (the light, numpy-only
+embedder, in-memory store) and shows the **back-edge** recovering a gold document
+that single-shot retrieval misses.
+"""
+
+import ir
+from ir import SearchHit
+from ir.store import CorpusStore
+
+from raglab import (
+    Budget,
+    LowLevelQuery,
+    SubTask,
+    make_llm_evaluator,
+    make_llm_formulator,
+    make_search_agent,
+)
+
+
+def _hits(*specs):
+    """``(artifact_id, score)`` pairs -> ir.SearchHits (no corpus needed)."""
+    return [SearchHit(aid, "k", score, f"text {aid}", {}) for aid, score in specs]
+
+
+def _fake_retriever(hits):
+    """A Retriever that records its calls and returns canned hits."""
+    calls = []
+
+    def retrieve(query, **kw):
+        calls.append((query, kw))
+        return list(hits)
+
+    retrieve.calls = calls
+    return retrieve
+
+
+# ----- LLM Formulator: adapt str -> [str] to (SubTask, source) -> [LLQ] ------ #
+
+
+def test_llm_formulator_fans_out_one_llq_per_query():
+    formulator = make_llm_formulator(formulate=lambda q: [q, q + " alt"])
+    llqs = formulator(SubTask(goal="deploy", sources=("s",)), "s")
+    assert [q.query for q in llqs] == ["deploy", "deploy alt"]
+    assert all(isinstance(q, LowLevelQuery) and q.source == "s" for q in llqs)
+
+
+def test_llm_formulator_accepts_a_bare_string():
+    formulator = make_llm_formulator(formulate=lambda q: q + "!")
+    llqs = formulator(SubTask(goal="x", sources=("s",)), "s")
+    assert [q.query for q in llqs] == ["x!"]
+
+
+def test_llm_formulator_attaches_params_to_every_query():
+    formulator = make_llm_formulator(
+        formulate=lambda q: [q, q + " b"], params={"mode": "hybrid", "k": 5}
+    )
+    llqs = formulator(SubTask(goal="g", sources=("s",)), "s")
+    assert all(q.params == {"mode": "hybrid", "k": 5} for q in llqs)
+
+
+def test_llm_formulator_empty_output_falls_back_to_the_goal():
+    # A formulator must never make retrieval worse than the raw sub-goal.
+    formulator = make_llm_formulator(formulate=lambda q: [])
+    llqs = formulator(SubTask(goal="the goal", sources=("s",)), "s")
+    assert [q.query for q in llqs] == ["the goal"]
+
+
+def test_llm_formulator_swallows_a_raising_callable():
+    # A failing custom formulator must fall back to the goal, never propagate.
+    def boom(_q):
+        raise RuntimeError("rewriter down")
+
+    formulator = make_llm_formulator(formulate=boom)
+    llqs = formulator(SubTask(goal="the goal", sources=("s",)), "s")
+    assert [q.query for q in llqs] == ["the goal"]
+
+
+def test_llm_formulator_drives_the_agent_loop():
+    retr = _fake_retriever(_hits(("a", 0.9)))
+    formulator = make_llm_formulator(formulate=lambda q: [q, q + " expanded"])
+    make_search_agent({"s": retr}, formulator=formulator)("q")
+    assert [c[0] for c in retr.calls] == ["q", "q expanded"]
+
+
+# ----- LLM Evaluator: ir.select owns relevance, the LLM owns sufficiency ----- #
+
+
+def test_evaluator_relevance_comes_from_ir_select():
+    # Conservative selection keeps only the near-top hit; "b" is a distractor.
+    evaluator = make_llm_evaluator(judge=lambda **kw: (True, None))
+    judged = evaluator(SubTask("g", ("s",)), _hits(("a", 0.9), ("b", 0.1)))
+    assert [h.artifact_id for h in judged.relevant] == ["a"]
+    assert judged.sufficient and judged.refinement is None
+
+
+def test_evaluator_emits_a_refinement_when_insufficient():
+    evaluator = make_llm_evaluator(judge=lambda **kw: (False, "better query"))
+    judged = evaluator(SubTask("g", ("s1", "s2")), _hits(("a", 0.9)))
+    assert judged.sufficient is False
+    assert judged.refinement == SubTask(goal="better query", sources=("s1", "s2"))
+
+
+def test_evaluator_insufficient_without_a_query_stops_the_loop():
+    # Insufficient but no refinement query -> nothing better to try -> stop.
+    evaluator = make_llm_evaluator(judge=lambda **kw: (False, None))
+    judged = evaluator(SubTask("g", ("s",)), _hits(("a", 0.9)))
+    assert judged.sufficient is True and judged.refinement is None
+
+
+def test_evaluator_parses_a_raw_text_reply():
+    evaluator = make_llm_evaluator(
+        judge=lambda **kw: "INSUFFICIENT\nvector database filtering"
+    )
+    judged = evaluator(SubTask("g", ("s",)), _hits(("a", 0.9)))
+    assert judged.refinement.goal == "vector database filtering"
+
+
+def test_evaluator_judge_error_falls_back_to_signal_no_loop():
+    def boom(**kw):
+        raise RuntimeError("model down")
+
+    evaluator = make_llm_evaluator(judge=boom)
+    judged = evaluator(SubTask("g", ("s",)), _hits(("a", 0.9)))
+    # refinement=None is the loop's break condition: a judge error never spins.
+    assert judged.refinement is None
+    assert judged.sufficient is True  # ir.select committed to "a" -> sufficient
+
+
+def test_evaluator_renders_abstention_to_the_judge():
+    seen = {}
+
+    def judge(*, goal, results):
+        seen["results"] = results
+        return (True, None)
+
+    evaluator = make_llm_evaluator(judge=judge)
+    evaluator(SubTask("g", ("s",)), [])  # no results -> ir.select abstains
+    assert "abstained" in seen["results"]
+
+
+def test_evaluator_forwards_select_kwargs_to_ir_select():
+    # Two near-tied hits: conservative keeps only "a" by default, but a loose
+    # rel threshold (forwarded via select_kwargs) admits "b" too.
+    hits = _hits(("a", 0.9), ("b", 0.6))
+    strict = make_llm_evaluator(judge=lambda **kw: (True, None))
+    loose = make_llm_evaluator(
+        judge=lambda **kw: (True, None), select_kwargs={"rel": 0.5}
+    )
+    assert [h.artifact_id for h in strict(SubTask("g", ("s",)), hits).relevant] == ["a"]
+    assert [h.artifact_id for h in loose(SubTask("g", ("s",)), hits).relevant] == [
+        "a",
+        "b",
+    ]
+
+
+def test_evaluator_ranks_heterogeneous_results_before_selecting():
+    # Accumulated cross-source results arrive unordered; the evaluator must rank
+    # best-first before ir.select (which trusts input order).
+    captured = {}
+
+    def judge(*, goal, results):
+        captured["results"] = results
+        return (True, None)
+
+    evaluator = make_llm_evaluator(judge=judge, select_kwargs={"rel": 0.0})
+    unordered = _hits(("lo", 0.1), ("hi", 0.9), ("mid", 0.5))
+    judged = evaluator(SubTask("g", ("s",)), unordered)
+    assert [h.artifact_id for h in judged.relevant] == ["hi", "mid", "lo"]
+
+
+# ----- the back-edge end-to-end, wired through the agent loop ---------------- #
+
+
+def test_evaluator_back_edge_loops_until_sufficient():
+    retr = _fake_retriever(_hits(("a", 0.9)))
+    rounds = {"n": 0}
+
+    def judge(*, goal, results):
+        rounds["n"] += 1
+        if rounds["n"] < 2:
+            return (False, goal + " more")
+        return (True, None)
+
+    make_search_agent({"s": retr}, evaluator=make_llm_evaluator(judge=judge))("q")
+    assert rounds["n"] == 2  # looped once via the back-edge
+    assert [c[0] for c in retr.calls] == ["q", "q more"]  # refinement re-queried
+
+
+def test_evaluator_back_edge_is_bounded_by_budget():
+    retr = _fake_retriever(_hits(("a", 0.9)))
+    evaluator = make_llm_evaluator(judge=lambda **kw: (False, "again"))
+    make_search_agent({"s": retr}, evaluator=evaluator, budget=Budget(max_rounds=3))(
+        "q"
+    )
+    assert len(retr.calls) == 3  # the safety net holds even if never sufficient
+
+
+# ----- end-to-end over a REAL ir corpus (hermetic: light embedder) ---------- #
+
+
+def _light_corpus():
+    docs = {
+        "embed": "embed and cache model vectors",
+        "systemd": "configure systemd units and restart services",
+        "filtering": "narrow similarity search using metadata filters",
+    }
+    return ir.build(
+        ir.CorpusSource.from_mapping(docs, name="t", strategy=ir.WholeText()),
+        store=CorpusStore.memory(),
+        embedder="light",
+    )
+
+
+def test_back_edge_recovers_a_doc_single_shot_misses():
+    """A query that overlaps a distractor misses the gold; the refinement recovers it.
+
+    Deterministic with the light embedder: the round-1 query shares vocabulary
+    with ``embed`` (a positive-score distractor) but none with the gold
+    ``filtering``, so single-shot ranks ``embed`` first; the injected judge
+    declares it insufficient and reformulates to the gold doc's own vocabulary, so
+    round 2 retrieves ``filtering`` to the top via the back-edge.
+    """
+    corpus = _light_corpus()
+    sources = {"t": ir.as_retriever(corpus, k=3)}
+    vague = "cache model results"  # overlaps the `embed` distractor, not the gold
+    gold_query = "narrow similarity search using metadata filters"
+
+    # Baseline: single-shot (no LLM evaluator) surfaces the distractor, not the gold.
+    baseline = make_search_agent(sources)(vague)
+    assert baseline[0].artifact_id == "embed"
+
+    # With the back-edge: reformulate to the gold's vocabulary, then it wins.
+    rounds = {"n": 0}
+
+    def judge(*, goal, results):
+        rounds["n"] += 1
+        if rounds["n"] < 2:
+            return (False, gold_query)
+        return (True, None)
+
+    agent = make_search_agent(sources, evaluator=make_llm_evaluator(judge=judge))
+    results = agent(vague)
+    assert rounds["n"] == 2  # the back-edge fired
+    assert results[0].artifact_id == "filtering"  # gold recovered
+    assert isinstance(results[0], SearchHit)