raglab 0.2.6__tar.gz → 0.2.7__tar.gz

{raglab-0.2.6 → raglab-0.2.7}/.claude/CLAUDE.md

@@ -29,7 +29,10 @@ control loop, with concrete tools injected at the leaves.
   cyclic-graph runner. This back-edge is what makes it an agent vs a DAG.
 - **Budget governor** (`max_rounds` / `max_sources_per_task` / `max_results_per_task`);
   termination as a **separately injected policy** (ir_09 §9), not folded into the
-  evaluator.
+  evaluator. Pool bounding is an injected `pool_cap` policy (default
+  `rank_fair_cap`: per-source best-first, round-robin). It must return hits
+  untouched — never pre-rank the pool with the RRF reranker (fused scores must not
+  re-enter the pool).
 - **Source registry** = a live `Mapping[name, Retriever]` across *heterogeneous*
   backends (ir corpora + web/SQL/graph); cross-source merge + global rerank at
   the fan-in point.

{raglab-0.2.6 → raglab-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: raglab
-Version: 0.2.6
+Version: 0.2.7
 Summary: A medley of tools to make RAG-based applications.
 Project-URL: Homepage, https://thorwhalen.github.io/raglab/
 Project-URL: Repository, https://github.com/thorwhalen/raglab
@@ -68,7 +68,7 @@ import raglab
 # Register named ir corpora as the agent's sources, then search across them.
 sources = raglab.ir_sources("skills", "reports", mode="hybrid")
 agent = raglab.make_search_agent(sources)
-results = agent("how do I deploy the app")     # -> list[ir.SearchHit], ranked
+results = agent("how do I deploy the app")  # -> list[ir.SearchHit], ranked
 ```
 
 A self-contained, runnable example (build a tiny corpus, search it):
@@ -80,18 +80,18 @@ import raglab
 
 docs = {
     "deploy": "deploy the app to the server with systemd units",
-    "embed":  "embed text with a model and cache the vectors",
+    "embed": "embed text with a model and cache the vectors",
     "search": "vector similarity search with metadata filters",
 }
 corpus = ir.build(
     ir.CorpusSource.from_mapping(docs, name="t", strategy=ir.WholeText()),
     store=CorpusStore.memory(),
-    embedder="light",          # numpy-only, no model download
+    embedder="light",  # numpy-only, no model download
 )
 
 agent = raglab.make_search_agent({"t": ir.as_retriever(corpus, k=3)})
 results = agent("deploy the app to the server")
-results[0].artifact_id        # -> "deploy"
+results[0].artifact_id  # -> "deploy"
 ```
 
 `sources` is any `Mapping[name, Retriever]`. Build it with
@@ -147,8 +147,8 @@ import raglab
 
 agent = raglab.make_search_agent(
     sources,
-    formulator=raglab.make_llm_formulator(),   # rewrite/expand the query (HyDE)
-    evaluator=raglab.make_llm_evaluator(),      # sufficiency + the back-edge
+    formulator=raglab.make_llm_formulator(),  # rewrite/expand the query (HyDE)
+    evaluator=raglab.make_llm_evaluator(),  # sufficiency + the back-edge
 )
 results = agent("how do I deploy the app")
 ```
@@ -210,14 +210,19 @@ Every role is just a callable matching a `Protocol`, so you can inject your own:
 ```python
 from raglab import SubTask, LowLevelQuery, Judgement, make_search_agent
 
+
 def my_planner(query, sources):
     # split into two sub-goals over all sources
-    return [SubTask(goal=g, sources=tuple(sources))
-            for g in (query.text, query.text + " examples")]
+    return [
+        SubTask(goal=g, sources=tuple(sources))
+        for g in (query.text, query.text + " examples")
+    ]
+
 
 def my_formulator(task, source):
     return [LowLevelQuery(source=source, query=task.goal, params={"k": 10})]
 
+
 agent = make_search_agent(sources, planner=my_planner, formulator=my_formulator)
 ```
 

{raglab-0.2.6 → raglab-0.2.7}/README.md

@@ -34,7 +34,7 @@ import raglab
 # Register named ir corpora as the agent's sources, then search across them.
 sources = raglab.ir_sources("skills", "reports", mode="hybrid")
 agent = raglab.make_search_agent(sources)
-results = agent("how do I deploy the app")     # -> list[ir.SearchHit], ranked
+results = agent("how do I deploy the app")  # -> list[ir.SearchHit], ranked
 ```
 
 A self-contained, runnable example (build a tiny corpus, search it):
@@ -46,18 +46,18 @@ import raglab
 
 docs = {
     "deploy": "deploy the app to the server with systemd units",
-    "embed":  "embed text with a model and cache the vectors",
+    "embed": "embed text with a model and cache the vectors",
     "search": "vector similarity search with metadata filters",
 }
 corpus = ir.build(
     ir.CorpusSource.from_mapping(docs, name="t", strategy=ir.WholeText()),
     store=CorpusStore.memory(),
-    embedder="light",          # numpy-only, no model download
+    embedder="light",  # numpy-only, no model download
 )
 
 agent = raglab.make_search_agent({"t": ir.as_retriever(corpus, k=3)})
 results = agent("deploy the app to the server")
-results[0].artifact_id        # -> "deploy"
+results[0].artifact_id  # -> "deploy"
 ```
 
 `sources` is any `Mapping[name, Retriever]`. Build it with
@@ -113,8 +113,8 @@ import raglab
 
 agent = raglab.make_search_agent(
     sources,
-    formulator=raglab.make_llm_formulator(),   # rewrite/expand the query (HyDE)
-    evaluator=raglab.make_llm_evaluator(),      # sufficiency + the back-edge
+    formulator=raglab.make_llm_formulator(),  # rewrite/expand the query (HyDE)
+    evaluator=raglab.make_llm_evaluator(),  # sufficiency + the back-edge
 )
 results = agent("how do I deploy the app")
 ```
@@ -176,14 +176,19 @@ Every role is just a callable matching a `Protocol`, so you can inject your own:
 ```python
 from raglab import SubTask, LowLevelQuery, Judgement, make_search_agent
 
+
 def my_planner(query, sources):
     # split into two sub-goals over all sources
-    return [SubTask(goal=g, sources=tuple(sources))
-            for g in (query.text, query.text + " examples")]
+    return [
+        SubTask(goal=g, sources=tuple(sources))
+        for g in (query.text, query.text + " examples")
+    ]
+
 
 def my_formulator(task, source):
     return [LowLevelQuery(source=source, query=task.goal, params={"k": 10})]
 
+
 agent = make_search_agent(sources, planner=my_planner, formulator=my_formulator)
 ```
 

{raglab-0.2.6 → raglab-0.2.7}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "raglab"
-version = "0.2.6"
+version = "0.2.7"
 description = "A medley of tools to make RAG-based applications."
 readme = "README.md"
 requires-python = ">=3.10"

{raglab-0.2.6 → raglab-0.2.7}/raglab/__init__.py

@@ -33,18 +33,21 @@ from .agent import (
     Judgement,
     LowLevelQuery,
     Planner,
+    PoolCap,
     Query,
     Reranker,
     Result,
     Retriever,
     SingleContextAgent,
     SubTask,
+    arrival_order_cap,
     identity_citer,
     identity_formulator,
     ir_sources,
     make_rrf_reranker,
     make_search_agent,
     passthrough_evaluator,
+    rank_fair_cap,
     rrf_reranker,
     score_reranker,
     single_subtask_planner,
@@ -64,6 +67,9 @@ __all__ = [
     "Reranker",
     "Citer",
     "Budget",
+    "PoolCap",
+    "rank_fair_cap",
+    "arrival_order_cap",
     "SingleContextAgent",
     "make_search_agent",
     "ir_sources",

{raglab-0.2.6 → raglab-0.2.7}/raglab/agent.py

@@ -31,6 +31,7 @@ Dependency direction is one-way: `raglab` imports `ir`; `ir` never imports
 
 from __future__ import annotations
 
+import math
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import Any, Protocol, runtime_checkable
@@ -59,6 +60,9 @@ __all__ = [
     "Reranker",
     "Citer",
     "Budget",
+    "PoolCap",
+    "rank_fair_cap",
+    "arrival_order_cap",
     "SingleContextAgent",
     "make_search_agent",
     "ir_sources",
@@ -178,6 +182,70 @@ class Budget:
     max_results_per_task: int = 50
 
 
+@runtime_checkable
+class PoolCap(Protocol):
+    """Bound the pool the evaluator sees to at most ``limit`` results.
+
+    A governor decision (which hits survive the budget), injected like every other
+    role so a strategy can be swapped without touching the loop. It must return
+    hits **unchanged** — never re-scored — because the pool keeps raw per-source
+    magnitudes until the fan-in reranker.
+    """
+
+    def __call__(self, results: Sequence[Result], limit: int) -> list[Result]: ...
+
+
+def _rank_key(hit: Result) -> float:
+    """A hit's raw score for ordering *within its source*; a missing or NaN score is last."""
+    score = hit.score
+    if isinstance(score, (int, float)) and not math.isnan(score):
+        return float(score)
+    return -math.inf
+
+
+def rank_fair_cap(results: Sequence[Result], limit: int) -> list[Result]:
+    """Keep at most ``limit`` hits, taking the best of each source in turn (the default).
+
+    Groups the pool by ``hit.source`` (in first-arrival order; untagged hits are one
+    group), collapses each group to one hit per artifact (its best surface, as
+    :func:`ir.base.best_per_artifact` does), orders the group best-first by its own raw
+    score -- sound, since one source is one score scale -- and deals hits round-robin
+    across the groups. So a late source's strong hits are never cut before the evaluator
+    sees them, a duplicate never takes a slot, and no score ever crosses a source
+    boundary or is rewritten (unlike pre-ranking with the RRF reranker, whose fused
+    ordinal scores must never re-enter the pool). The output is in that dealt order
+    whatever the pool's size. A missing or NaN score sorts last in its source.
+    """
+    if limit <= 0:
+        return []
+    groups: dict[str | None, dict[str, Result]] = {}
+    for hit in results:
+        best = groups.setdefault(hit.source, {})
+        cur = best.get(hit.artifact_id)
+        if cur is None or _rank_key(hit) > _rank_key(cur):
+            best[hit.artifact_id] = hit
+    queues = [
+        sorted(best.values(), key=_rank_key, reverse=True) for best in groups.values()
+    ]
+    capped: list[Result] = []
+    for depth in range(max((len(q) for q in queues), default=0)):
+        for queue in queues:
+            if depth < len(queue):
+                capped.append(queue[depth])
+                if len(capped) == limit:
+                    return capped
+    return capped
+
+
+def arrival_order_cap(results: Sequence[Result], limit: int) -> list[Result]:
+    """Keep the first ``limit`` hits in arrival order (source order x query order).
+
+    The pre-governor behaviour: a late source's strong hits can be cut before the
+    evaluator sees them. Kept as an explicit opt-in (``pool_cap=arrival_order_cap``).
+    """
+    return list(results[: max(0, limit)])
+
+
 # --------------------------------------------------------------------------- #
 # Default role implementations (the no-LLM thin slice) — ir_09 §8 step 2
 # --------------------------------------------------------------------------- #
@@ -306,6 +374,7 @@ class SingleContextAgent:
     reranker: Reranker = rrf_reranker
     citer: Citer = identity_citer
     budget: Budget = field(default_factory=Budget)
+    pool_cap: PoolCap = rank_fair_cap
 
     def __call__(self, query: str | Query) -> list[Result]:
         """Run the agent for *query*; returns the final ranked, cited results."""
@@ -317,9 +386,15 @@ class SingleContextAgent:
         return list(self.citer(ranked))
 
     def _run_task(self, task: SubTask) -> list[Result]:
-        found: list[Result] = []
+        # `kept` is what the evaluator judged relevant last round: governor-approved,
+        # carried into the next round untouched. Only a round's *new* hits compete for
+        # the rest of the budget (`pool_cap`), so a re-query can never evict a hit the
+        # evaluator already kept, nor spend a slot re-offering it.
+        kept: list[Result] = []
         current = task
+        limit = self.budget.max_results_per_task
         for _ in range(max(1, self.budget.max_rounds)):
+            found: list[Result] = []
             for source in current.sources[: self.budget.max_sources_per_task]:
                 retriever = self.sources.get(source)
                 if retriever is None:
@@ -333,12 +408,15 @@ class SingleContextAgent:
                     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)
+            seen = {(h.source, h.artifact_id) for h in kept}
+            fresh = [h for h in found if (h.source, h.artifact_id) not in seen]
+            pool = kept + self.pool_cap(fresh, max(0, limit - len(kept)))
+            judged = self.evaluator(current, pool)
+            kept = list(judged.relevant)
             if judged.sufficient or judged.refinement is None:
                 break
             current = judged.refinement  # the back-edge: re-query
-        return found
+        return kept
 
 
 def make_search_agent(
@@ -350,6 +428,7 @@ def make_search_agent(
     reranker: Reranker | None = None,
     citer: Citer | None = None,
     budget: Budget | None = None,
+    pool_cap: PoolCap | None = None,
 ) -> SingleContextAgent:
     """Build a :class:`SingleContextAgent` over *sources* with smart defaults.
 
@@ -372,6 +451,7 @@ def make_search_agent(
         reranker=reranker or rrf_reranker,
         citer=citer or identity_citer,
         budget=budget or Budget(),
+        pool_cap=pool_cap or rank_fair_cap,
     )
 
 

{raglab-0.2.6 → raglab-0.2.7}/tests/test_fanin.py

@@ -286,3 +286,144 @@ 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"]
+
+
+# ----- the budget cap: rank-fair, scale-safe (raglab#5) ---------------------- #
+
+
+def _capturing_evaluator(seen):
+    def evaluator(task, results):
+        seen.append(list(results))
+        return Judgement(list(results), sufficient=True)
+
+    return evaluator
+
+
+def test_small_cap_lets_both_disparate_sources_top_hits_through():
+    # source "big" arrives first with many BM25-scale hits; "cos" arrives last with
+    # cosine-scale hits. Arrival-order truncation would never show "cos" at all.
+    big = _hits(*[(f"b{i}", 40.0 - i) for i in range(10)], source="big")
+    cos = _hits(("c_best", 0.95), ("c2", 0.4), source="cos")
+    seen = []
+    make_search_agent(
+        {"big": _fake_retriever(big), "cos": _fake_retriever(cos)},
+        evaluator=_capturing_evaluator(seen),
+        budget=raglab.Budget(max_results_per_task=4),
+    )("q")
+    pool = seen[0]
+    assert len(pool) == 4
+    ids = [h.artifact_id for h in pool]
+    assert ids[:2] == ["b0", "c_best"]  # each source's best, dealt round-robin
+    assert set(ids) == {"b0", "b1", "c_best", "c2"}
+
+
+def test_capped_pool_keeps_raw_per_source_magnitudes():
+    big = _hits(*[(f"b{i}", 40.0 - i) for i in range(10)], source="big")
+    cos = _hits(("c_best", 0.95), ("c2", 0.4), source="cos")
+    originals = {(h.source, h.artifact_id): h for h in big + cos}
+    seen = []
+    make_search_agent(
+        {"big": _fake_retriever(big), "cos": _fake_retriever(cos)},
+        evaluator=_capturing_evaluator(seen),
+        budget=raglab.Budget(max_results_per_task=3),
+    )("q")
+    for hit in seen[0]:
+        assert hit == originals[(hit.source, hit.artifact_id)]  # untouched
+        assert "source_rank" not in hit.metadata
+        assert "source_score" not in hit.metadata  # no fused-score stamp either
+
+
+def test_rank_fair_cap_orders_within_source_by_its_own_score():
+    # one source, scores out of arrival order (e.g. two formulator queries)
+    pool = _hits(("lo", 0.1), ("hi", 0.9), ("mid", 0.5), source="s")
+    assert [h.artifact_id for h in raglab.rank_fair_cap(pool, 2)] == ["hi", "mid"]
+    # untagged hits are one group of their own
+    mixed = _hits(("u1", 3.0), ("u2", 2.0)) + _hits(("t1", 0.5), source="t")
+    assert [h.artifact_id for h in raglab.rank_fair_cap(mixed, 2)] == ["u1", "t1"]
+
+
+def test_rank_fair_cap_within_limit_and_degenerate_limits():
+    pool = _hits(("a", 0.1), ("b", 0.9), source="s")
+    # within budget nothing is cut, and the order is the same dealt order as above it
+    assert [h.artifact_id for h in raglab.rank_fair_cap(pool, 5)] == ["b", "a"]
+    assert raglab.rank_fair_cap(pool, 0) == []
+    assert raglab.rank_fair_cap(pool, -1) == []
+    uneven = _hits(*[(f"x{i}", 1.0 - i / 10) for i in range(5)], source="x") + _hits(
+        ("y0", 9.0), source="y"
+    )
+    # a short source runs dry; the long one fills the rest
+    assert [h.artifact_id for h in raglab.rank_fair_cap(uneven, 4)] == [
+        "x0",
+        "y0",
+        "x1",
+        "x2",
+    ]
+
+
+def test_pool_cap_is_injectable_and_arrival_order_is_the_opt_out():
+    big = _hits(*[(f"b{i}", 40.0 - i) for i in range(10)], source="big")
+    cos = _hits(("c_best", 0.95), source="cos")
+    seen = []
+    make_search_agent(
+        {"big": _fake_retriever(big), "cos": _fake_retriever(cos)},
+        evaluator=_capturing_evaluator(seen),
+        budget=raglab.Budget(max_results_per_task=3),
+        pool_cap=raglab.arrival_order_cap,
+    )("q")
+    assert [h.artifact_id for h in seen[0]] == ["b0", "b1", "b2"]
+    assert isinstance(raglab.rank_fair_cap, raglab.PoolCap)
+
+
+def test_cap_collapses_duplicates_and_tolerates_missing_scores():
+    dup = _hits(("a", 0.9), ("a", 0.5), ("b", 0.8), ("c", 0.7), source="s")
+    assert [h.artifact_id for h in raglab.rank_fair_cap(dup, 3)] == ["a", "b", "c"]
+    assert raglab.rank_fair_cap(dup, 3)[0].score == 0.9  # the best surface survives
+    odd = _hits(("nan", float("nan")), ("none", None), ("ok", 0.1), source="s")
+    assert raglab.rank_fair_cap(odd, 1)[0].artifact_id == "ok"
+
+
+def _refine_once(seen, keep):
+    """An evaluator that keeps only `keep` in round one and asks for a re-query."""
+
+    def evaluator(task, results):
+        seen.append(list(results))
+        if len(seen) == 1:
+            kept = [h for h in results if h.artifact_id in keep]
+            return Judgement(
+                kept, sufficient=False, refinement=SubTask("again", task.sources)
+            )
+        return Judgement(list(results), sufficient=True)
+
+    return evaluator
+
+
+def _rounds(*per_round):
+    """A retriever returning a different canned hit list on each call."""
+    calls = iter(per_round)
+    return lambda query, **kw: list(next(calls))
+
+
+def test_back_edge_never_evicts_a_hit_the_evaluator_kept():
+    round1 = _hits(
+        ("weak_but_relevant", 0.1), ("junk1", 0.9), ("junk2", 0.8), source="s"
+    )
+    round2 = _hits(("new1", 0.95), ("new2", 0.85), ("new3", 0.7), source="s")
+    seen = []
+    out = make_search_agent(
+        {"s": _rounds(round1, round2)},
+        evaluator=_refine_once(seen, {"weak_but_relevant"}),
+        budget=raglab.Budget(max_results_per_task=3),
+    )("q")
+    assert [h.artifact_id for h in seen[1]] == ["weak_but_relevant", "new1", "new2"]
+    assert "weak_but_relevant" in {h.artifact_id for h in out}
+
+
+def test_back_edge_does_not_spend_slots_on_re_retrieved_hits():
+    same = _hits(("a", 0.9), ("b", 0.8), ("c", 0.7), source="s")
+    seen = []
+    make_search_agent(
+        {"s": _rounds(same, same)},
+        evaluator=_refine_once(seen, {"a"}),
+        budget=raglab.Budget(max_results_per_task=3),
+    )("q")
+    assert [h.artifact_id for h in seen[1]] == ["a", "b", "c"]  # no second "a"