rerankgate 0.2.0__py3-none-any.whl

rerankgate/__init__.py

@@ -0,0 +1,30 @@
+"""rerankgate — agentic RAG quality gate.
+
+Drop-in retrieval pipeline with built-in quality gates:
+hybrid search (BM25 + vector) -> rerank -> sufficiency gate ->
+query-rewrite retry or abstain. Local-first; the sufficiency gate
+works without any LLM call by default.
+"""
+
+from rerankgate.config import GateConfig, PipelineConfig, RetryConfig
+from rerankgate.documents import Document, chunk_text
+from rerankgate.gate import GateDecision, SufficiencyGate
+from rerankgate.pipeline import AnswerResult, RAGPipeline, Verdict
+from rerankgate.retrieval import HybridRetriever, ScoredChunk
+
+__all__ = [
+    "AnswerResult",
+    "Document",
+    "GateConfig",
+    "GateDecision",
+    "HybridRetriever",
+    "PipelineConfig",
+    "RAGPipeline",
+    "RetryConfig",
+    "ScoredChunk",
+    "SufficiencyGate",
+    "Verdict",
+    "chunk_text",
+]
+
+__version__ = "0.2.0"

rerankgate/bm25.py

@@ -0,0 +1,70 @@
+"""Pure-stdlib BM25 (Okapi) index over chunks."""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+
+from rerankgate.documents import Chunk
+from rerankgate.tokenize import tokenize
+
+
+class BM25Index:
+    """In-memory BM25 index. Deterministic; rebuild from chunks."""
+
+    def __init__(self, k1: float = 1.5, b: float = 0.75) -> None:
+        self.k1 = k1
+        self.b = b
+        self._chunks: list[Chunk] = []
+        self._doc_freq: Counter[str] = Counter()
+        self._term_freqs: list[Counter[str]] = []
+        self._doc_lens: list[int] = []
+        self._avg_len: float = 0.0
+
+    def __len__(self) -> int:
+        return len(self._chunks)
+
+    def add(self, chunks: list[Chunk]) -> None:
+        for ch in chunks:
+            toks = tokenize(ch.text)
+            tf = Counter(toks)
+            self._chunks.append(ch)
+            self._term_freqs.append(tf)
+            self._doc_lens.append(len(toks))
+            for term in tf:
+                self._doc_freq[term] += 1
+        total = sum(self._doc_lens)
+        self._avg_len = total / len(self._doc_lens) if self._doc_lens else 0.0
+
+    def _idf(self, term: str) -> float:
+        n = len(self._chunks)
+        df = self._doc_freq.get(term, 0)
+        # Standard BM25 IDF with +1 smoothing floor; unseen terms get ~0.
+        return math.log(1 + (n - df + 0.5) / (df + 0.5)) if df else 0.0
+
+    def search(self, query: str, k: int = 100) -> list[tuple[Chunk, float]]:
+        """Return top-k (chunk, score) pairs, highest score first."""
+        q_terms = tokenize(query)
+        if not q_terms or not self._chunks:
+            return []
+        idfs = {t: self._idf(t) for t in set(q_terms)}
+        scored: list[tuple[Chunk, float]] = []
+        for i, ch in enumerate(self._chunks):
+            tf = self._term_freqs[i]
+            dl = self._doc_lens[i]
+            norm = 1 - self.b + self.b * (dl / self._avg_len) if self._avg_len else 1.0
+            score = 0.0
+            for t in q_terms:
+                f = tf.get(t, 0)
+                if f and idfs[t]:
+                    score += idfs[t] * (f * (self.k1 + 1)) / (f + self.k1 * norm)
+            if score > 0:
+                scored.append((ch, score))
+        scored.sort(key=lambda p: p[1], reverse=True)
+        return scored[:k]
+
+    def term_stats(self, chunk_idx: int) -> Counter[str]:
+        return self._term_freqs[chunk_idx]
+
+    def chunks(self) -> list[Chunk]:
+        return list(self._chunks)

rerankgate/cli.py

@@ -0,0 +1,142 @@
+"""rerankgate CLI: ingest | query | eval."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+from rerankgate.config import PipelineConfig
+from rerankgate.documents import Document
+from rerankgate.metrics import EvalItem
+from rerankgate.pipeline import RAGPipeline
+
+
+def _pipeline_for(args: argparse.Namespace, *, for_ingest: bool = False) -> RAGPipeline:
+    cfg = PipelineConfig()
+    if args.gate_preset:
+        cfg.gate_preset = args.gate_preset
+    if args.max_retries is not None:
+        cfg.retry.max_retries = args.max_retries
+    if args.no_vectors:
+        cfg.use_vectors = False
+    if for_ingest:
+        return RAGPipeline(config=cfg)
+    return RAGPipeline.load(args.store, config=cfg)
+
+
+def cmd_ingest(args: argparse.Namespace) -> int:
+    docs: list[Document] = []
+    with open(args.input, encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            text = obj.get("text", "")
+            meta = {k: v for k, v in obj.items() if k not in ("id", "text")}
+            docs.append(Document(id=str(obj.get("id", f"doc-{i}")), text=text, metadata=meta))
+    pipe = _pipeline_for(args, for_ingest=True)
+    n_chunks = pipe.ingest(docs)
+    pipe.save(args.store)
+    print(f"ingested {len(docs)} documents -> {n_chunks} chunks -> {args.store}")
+    if pipe.degraded_reason:
+        print("note: BM25-only mode (no embeddings installed)")
+    return 0
+
+
+def cmd_query(args: argparse.Namespace) -> int:
+    pipe = _pipeline_for(args)
+    result = pipe.ask(args.query)
+    out = result.to_dict()
+    if args.json:
+        print(json.dumps(out, indent=2))
+        return 0
+    print(f"verdict:  {out['verdict']}")
+    print(
+        f"gate:     {'PASS' if out['gate']['passed'] else 'FAIL'} "
+        f"(score {out['gate']['score']} vs threshold {out['gate']['threshold']})"
+    )
+    print(f"reasons:  {', '.join(out['gate']['reasons'])}")
+    print(
+        f"attempts: {out['trace']['attempts']}  "
+        f"modes: {','.join(out['trace']['modes'])}  "
+        f"total: {out['trace']['total_ms']}ms"
+    )
+    if out["trace"]["degraded"]:
+        print("degraded: BM25-only (embeddings unavailable)")
+    print("evidence:")
+    for ev in out["evidence"][:10]:
+        snippet = ev["text"][:120].replace("\n", " ")
+        print(f"  [{ev['score']:.3f}] {ev['chunk_id']}: {snippet}...")
+    return 0
+
+
+def cmd_eval(args: argparse.Namespace) -> int:
+    pipe = _pipeline_for(args)
+    items: list[EvalItem] = []
+    with open(args.qrels, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            items.append(
+                EvalItem(
+                    query=obj["query"],
+                    relevant_ids=list(obj.get("relevant_ids", [])),
+                    answerable=bool(obj.get("answerable", True)),
+                )
+            )
+    report = pipe.evaluate(items)
+    if args.summary:
+        report = {k: v for k, v in report.items() if k != "per_query"}
+    print(json.dumps(report, indent=2))
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    import rerankgate
+
+    p = argparse.ArgumentParser(prog="rerankgate", description="Agentic RAG quality gate")
+    p.add_argument("--version", action="version", version=f"%(prog)s {rerankgate.__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    def _common(sp: argparse.ArgumentParser) -> None:
+        sp.add_argument("--gate-preset", choices=["strict", "balanced", "lenient"], default=None)
+        sp.add_argument("--max-retries", type=int, default=None)
+        sp.add_argument("--no-vectors", action="store_true", help="force BM25-only mode")
+
+    pi = sub.add_parser("ingest", help="ingest a JSONL corpus into a store")
+    pi.add_argument("--input", required=True, help="JSONL with {id, text, ...metadata}")
+    pi.add_argument("--store", required=True, help="output store directory")
+    _common(pi)
+    pi.set_defaults(func=cmd_ingest)
+
+    pq = sub.add_parser("query", help="ask a question against a store")
+    pq.add_argument("--store", required=True)
+    pq.add_argument("query")
+    pq.add_argument("--json", action="store_true")
+    _common(pq)
+    pq.set_defaults(func=cmd_query)
+
+    pe = sub.add_parser("eval", help="evaluate retrieval + gate on labeled qrels")
+    pe.add_argument("--store", required=True)
+    pe.add_argument("--qrels", required=True, help="JSONL with {query, relevant_ids, answerable}")
+    pe.add_argument(
+        "--summary",
+        action="store_true",
+        help="print aggregate metrics only, without the per-query detail",
+    )
+    _common(pe)
+    pe.set_defaults(func=cmd_eval)
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

rerankgate/config.py

@@ -0,0 +1,77 @@
+"""Configuration dataclasses for rerankgate."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class RetryConfig:
+    """Bounds for the query-rewrite retry loop."""
+
+    max_retries: int = 2
+    # Stop early if the best gate score does not improve by at least this much.
+    min_improvement: float = 0.01
+    # How many expansion terms the RM3-style rewriter may add per retry.
+    expansion_terms: int = 5
+
+
+@dataclass
+class GateConfig:
+    """Tunable knobs for the heuristic sufficiency gate.
+
+    The gate score is a fixed weighted sum of interpretable features
+    (see gate.py). Nothing is learned, so there is nothing to overfit;
+    tune the threshold with `rerankgate eval` on your own labeled data.
+    """
+
+    threshold: float = 0.55
+    w_top_score: float = 0.35
+    w_coverage: float = 0.30
+    w_margin: float = 0.20
+    w_support: float = 0.15
+    # Minimum fraction of (non-stopword) query terms that must appear
+    # in the evidence before the gate can pass, regardless of score.
+    min_coverage: float = 0.25
+    # Minimum normalized top-1 rerank score before the gate can pass.
+    min_top_score: float = 0.10
+
+    def preset(self, name: str) -> GateConfig:
+        """Return a copy with a named strictness preset applied."""
+        presets = {
+            # (threshold, min_coverage): stricter gates demand more evidence.
+            "strict": (0.70, 0.35),
+            "balanced": (0.55, 0.25),
+            "lenient": (0.40, 0.15),
+        }
+        if name not in presets:
+            raise ValueError(f"unknown gate preset {name!r}; choose from {sorted(presets)}")
+        import dataclasses
+
+        threshold, min_coverage = presets[name]
+        return dataclasses.replace(self, threshold=threshold, min_coverage=min_coverage)
+
+
+@dataclass
+class PipelineConfig:
+    """Top-level pipeline configuration."""
+
+    # Retrieval
+    retrieve_k: int = 100  # candidates pulled from each retriever before fusion
+    rerank_k: int = 10  # evidence chunks returned after reranking
+    rrf_k: int = 60  # RRF rank constant
+    bm25_weight: float = 1.0  # per-system RRF multipliers
+    vector_weight: float = 1.0
+    use_vectors: bool = True  # degraded to BM25-only if embeddings unavailable
+
+    # Chunking
+    chunk_size: int = 500
+    chunk_overlap: int = 50
+
+    # Gate + retry
+    gate: GateConfig = field(default_factory=GateConfig)
+    retry: RetryConfig = field(default_factory=RetryConfig)
+    gate_preset: str = "balanced"
+
+    def effective_gate(self) -> GateConfig:
+        return self.gate.preset(self.gate_preset)

rerankgate/documents.py

@@ -0,0 +1,97 @@
+"""Documents and deterministic chunking."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+_SENTENCE_SPLIT = re.compile(r"(?<=[.!?])\s+")
+
+
+@dataclass
+class Document:
+    """A source document before chunking."""
+
+    id: str
+    text: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def title(self) -> str:
+        return str(self.metadata.get("title", ""))
+
+
+@dataclass
+class Chunk:
+    """A retrievable unit produced by chunking a Document."""
+
+    id: str  # f"{doc_id}#c{n}"
+    doc_id: str
+    text: str
+    position: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+def chunk_text(
+    text: str,
+    doc_id: str,
+    size: int = 500,
+    overlap: int = 50,
+    metadata: dict | None = None,
+) -> list[Chunk]:
+    """Split text into sentence-aware sliding-window chunks.
+
+    Deterministic: same input always yields the same chunks.
+    """
+    text = text.strip()
+    if not text:
+        return []
+    sentences = [s.strip() for s in _SENTENCE_SPLIT.split(text) if s.strip()]
+    if not sentences:
+        sentences = [text]
+
+    chunks: list[Chunk] = []
+    buf: list[str] = []
+    buf_len = 0
+    n = 0
+
+    def flush() -> None:
+        nonlocal n, buf, buf_len
+        if buf:
+            chunks.append(
+                Chunk(
+                    id=f"{doc_id}#c{n}",
+                    doc_id=doc_id,
+                    text=" ".join(buf),
+                    position=n,
+                    metadata=dict(metadata or {}),
+                )
+            )
+            n += 1
+
+    for sent in sentences:
+        # A single over-long sentence becomes its own chunk (hard-split).
+        while len(sent) > size:
+            flush()
+            chunks.append(
+                Chunk(
+                    id=f"{doc_id}#c{n}",
+                    doc_id=doc_id,
+                    text=sent[:size],
+                    position=n,
+                    metadata=dict(metadata or {}),
+                )
+            )
+            n += 1
+            sent = sent[size - overlap :]
+        if buf_len + len(sent) + 1 > size and buf:
+            flush()
+            # overlap: carry the tail of the previous chunk forward
+            tail = " ".join(buf)
+            buf = [tail[-overlap:]] if len(tail) > overlap else []
+            buf_len = sum(len(b) + 1 for b in buf)
+        buf.append(sent)
+        buf_len += len(sent) + 1
+    flush()
+    return chunks

rerankgate/embeddings.py

@@ -0,0 +1,76 @@
+"""Pluggable embedding backends.
+
+The default path is local-first and LLM-free. Heavy model dependencies
+are lazy: `SentenceTransformerBackend` only imports sentence-transformers
+when instantiated, and the pipeline degrades to BM25-only if embeddings
+are unavailable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol
+
+
+class EmbeddingBackend(Protocol):
+    """Contract every embedding backend implements."""
+
+    @property
+    def dim(self) -> int: ...
+    @property
+    def name(self) -> str: ...
+    def embed(self, texts: Sequence[str]) -> list[list[float]]: ...
+
+
+class SentenceTransformerBackend:
+    """Local dense embeddings via sentence-transformers (lazy import)."""
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None:
+        try:
+            from sentence_transformers import SentenceTransformer  # type: ignore
+        except ImportError as exc:
+            raise RuntimeError(
+                "sentence-transformers is not installed. "
+                "Install it with `pip install rerankgate[local]`, or run "
+                "the pipeline with use_vectors=False for BM25-only mode."
+            ) from exc
+        self._model = SentenceTransformer(model_name)
+        self._model_name = model_name
+
+    @property
+    def name(self) -> str:
+        return f"sentence-transformers:{self._model_name}"
+
+    @property
+    def dim(self) -> int:
+        return int(self._model.get_sentence_embedding_dimension())
+
+    def embed(self, texts: Sequence[str]) -> list[list[float]]:
+        vecs = self._model.encode(list(texts), normalize_embeddings=True)
+        return [list(map(float, v)) for v in vecs]
+
+
+class FakeEmbeddings:
+    """Deterministic toy backend for tests — NOT for production use."""
+
+    def __init__(self, dim: int = 16) -> None:
+        self._dim = dim
+
+    @property
+    def name(self) -> str:
+        return "fake-test-embeddings"
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    def embed(self, texts: Sequence[str]) -> list[list[float]]:
+        import hashlib
+
+        out: list[list[float]] = []
+        for t in texts:
+            h = hashlib.sha256(t.encode()).digest()
+            vec = [((h[i % len(h)] / 255.0) - 0.5) * 2 for i in range(self._dim)]
+            norm = sum(x * x for x in vec) ** 0.5 or 1.0
+            out.append([x / norm for x in vec])
+        return out

rerankgate/gate.py

@@ -0,0 +1,171 @@
+"""Heuristic sufficiency gate: "is this retrieval enough to answer?".
+
+Works WITHOUT any LLM call. Every decision carries machine-readable reason
+codes plus a human explanation, so a rejection can be audited and the
+threshold tuned with `rerankgate eval`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from time import perf_counter
+
+from rerankgate.config import GateConfig
+from rerankgate.retrieval import ScoredChunk
+from rerankgate.tokenize import content_tokens, tokenize
+
+# Reason codes (stable strings; safe to match on in calling code).
+PASS = "PASS"
+LOW_TOP_SCORE = "LOW_TOP_SCORE"
+LOW_COVERAGE = "LOW_COVERAGE"
+NO_MARGIN = "NO_MARGIN"
+NO_CONSENSUS = "NO_CONSENSUS"
+NO_EVIDENCE = "NO_EVIDENCE"
+
+_EXPLANATIONS = {
+    PASS: "Retrieved evidence is sufficient to answer.",
+    LOW_TOP_SCORE: "Best chunk score is below the minimum — nothing looks strongly relevant.",
+    LOW_COVERAGE: "Too few query terms appear in the evidence — likely off-topic retrieval.",
+    NO_MARGIN: "Top chunks score almost identically — no clear winner, ambiguous retrieval.",
+    NO_CONSENSUS: "Top chunks don't corroborate each other — evidence is fragmented.",
+    NO_EVIDENCE: "Retrieval returned no candidates at all.",
+}
+
+
+@dataclass
+class GateFeatures:
+    top_score: float = 0.0  # normalized top-1 rerank score [0,1]
+    coverage: float = 0.0  # fraction of query content-terms in evidence
+    margin: float = 0.0  # (s0 - s1) / s0 — decisiveness of the ranking
+    support: float = 0.0  # mean pairwise token overlap of top-3 chunks
+    depth: int = 0  # chunks above the score floor
+
+    def to_dict(self) -> dict:
+        return {
+            "top_score": round(self.top_score, 4),
+            "coverage": round(self.coverage, 4),
+            "margin": round(self.margin, 4),
+            "support": round(self.support, 4),
+            "depth": self.depth,
+        }
+
+
+@dataclass
+class GateDecision:
+    passed: bool
+    score: float
+    features: GateFeatures
+    reasons: list[str] = field(default_factory=list)
+    explanation: str = ""
+    threshold: float = 0.0
+    ms: float = 0.0
+
+    def to_dict(self) -> dict:
+        return {
+            "passed": self.passed,
+            "score": round(self.score, 4),
+            "threshold": self.threshold,
+            "reasons": self.reasons,
+            "explanation": self.explanation,
+            "features": self.features.to_dict(),
+            "ms": self.ms,
+        }
+
+
+class SufficiencyGate:
+    """Decide whether reranked evidence is sufficient to answer the query."""
+
+    def __init__(self, config: GateConfig | None = None) -> None:
+        self.config = config or GateConfig()
+
+    def evaluate(self, query: str, evidence: list[ScoredChunk]) -> GateDecision:
+        t0 = perf_counter()
+        cfg = self.config
+        features = self._features(query, evidence)
+        reasons: list[str] = []
+
+        if not evidence:
+            return GateDecision(
+                passed=False,
+                score=0.0,
+                features=features,
+                reasons=[NO_EVIDENCE],
+                explanation=_EXPLANATIONS[NO_EVIDENCE],
+                threshold=cfg.threshold,
+                ms=_ms(t0),
+            )
+
+        # Hard floors first: these fail regardless of the weighted score.
+        if features.top_score < cfg.min_top_score:
+            reasons.append(LOW_TOP_SCORE)
+        if features.coverage < cfg.min_coverage:
+            reasons.append(LOW_COVERAGE)
+
+        score = (
+            cfg.w_top_score * features.top_score
+            + cfg.w_coverage * features.coverage
+            + cfg.w_margin * features.margin
+            + cfg.w_support * features.support
+        )
+
+        passed = score >= cfg.threshold and not reasons
+        if passed:
+            reasons = [PASS]
+        else:
+            # Soft diagnostics: explain *why* a near-miss failed.
+            if features.margin < 0.05:
+                reasons.append(NO_MARGIN)
+            if features.support < 0.05 and features.depth >= 2:
+                reasons.append(NO_CONSENSUS)
+            reasons = reasons or [f"SCORE_BELOW_THRESHOLD ({score:.2f} < {cfg.threshold:.2f})"]
+
+        explanation = " ".join(_EXPLANATIONS.get(r, r) for r in reasons)
+        return GateDecision(
+            passed=passed,
+            score=round(score, 4),
+            features=features,
+            reasons=reasons,
+            explanation=explanation,
+            threshold=cfg.threshold,
+            ms=_ms(t0),
+        )
+
+    def _features(self, query: str, evidence: list[ScoredChunk]) -> GateFeatures:
+        q_terms = set(content_tokens(query))
+        scores = [c.score for c in evidence]
+        top = scores[0] if scores else 0.0
+        # Normalize top score to [0,1] against the observed max for stability.
+        norm_top = min(1.0, top)
+
+        ev_toks: set[str] = set()
+        for c in evidence:
+            ev_toks |= set(tokenize(c.chunk.text))
+        coverage = len(q_terms & ev_toks) / len(q_terms) if q_terms else 0.0
+
+        margin = 0.0
+        if len(scores) >= 2 and scores[0] > 0:
+            margin = max(0.0, (scores[0] - scores[1]) / scores[0])
+
+        support = 0.0
+        top3 = evidence[:3]
+        if len(top3) >= 2:
+            pairs = 0
+            acc = 0.0
+            sets = [set(content_tokens(c.chunk.text)) for c in top3]
+            for i in range(len(sets)):
+                for j in range(i + 1, len(sets)):
+                    a, b = sets[i], sets[j]
+                    union = a | b
+                    if union:
+                        acc += len(a & b) / len(union)
+                    pairs += 1
+            support = acc / pairs if pairs else 0.0
+
+        depth = sum(1 for s in scores if s >= 0.2)
+        return GateFeatures(
+            top_score=norm_top, coverage=coverage, margin=margin, support=support, depth=depth
+        )
+
+
+def _ms(t0: float) -> float:
+    return round((perf_counter() - t0) * 1000, 2)