rag-blocks 0.6.0__py3-none-any.whl

rag_blocks/__init__.py

@@ -0,0 +1,190 @@
+"""rag-blocks: composable building blocks for production RAG pipelines.
+
+Design in one paragraph: stages are swappable Strategies registered under a
+(kind, name) key; they communicate only through typed data contracts
+(Source → Page → Document → Chunk → …); ingestion is streaming-first so
+memory never scales with document size; every component carries a config
+fingerprint that later powers the auto-tuning suite's cross-pipeline caching.
+
+Quick start:
+
+    import rag_blocks as rk
+
+    doc = rk.ingest("report.pdf")                       # sane defaults
+    doc = rk.ingest("scan.pdf", ocr_engine="mistral",   # cloud OCR
+                    ocr_policy=rk.OcrPolicy.FORCE)
+
+    parser = rk.AutoParser()                            # streaming access
+    for page in parser.iter_pages(rk.Source.from_path("huge.pdf")):
+        ...
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from .core import (
+    Answer,
+    Chunk,
+    Citation,
+    Component,
+    Document,
+    EmbeddingError,
+    GenerationError,
+    Page,
+    PageSpan,
+    Query,
+    RagToolkitError,
+    ScoredChunk,
+    Source,
+    SourceFormat,
+    SparseVector,
+    StorageError,
+    VectorSpec,
+    VectorValue,
+    registry,
+)
+from .embedding import (
+    CachingEmbedder,
+    Embedder,
+    HashingEmbedder,
+    SentenceTransformerEmbedder,
+    SparseEncoder,
+)
+from .indexing import ChunkIndex, ChunkSink, DocumentCatalog, DocumentRef
+from .enrichment import (
+    ContextualEnricher,
+    Enricher,
+    HeadingEnricher,
+)
+from .generation import AnthropicGenerator, ExtractiveGenerator, Generator
+from .retrieval import (
+    FusionRetriever,
+    HybridRetriever,
+    HydeRetriever,
+    IndexRetriever,
+    MultiQueryRetriever,
+    Retriever,
+)
+from .ingestion import (
+    AutoParser,
+    DoclingParser,
+    OcrEngine,
+    OcrPolicy,
+    OcrResult,
+    PageImage,
+    Parser,
+    PlainTextParser,
+    detect_format,
+)
+from .chunking import Chunker, FixedChunker, MarkdownChunker
+from .refinement import (
+    CrossEncoderReranker,
+    KeywordRefiner,
+    NeighborExpander,
+    Refiner,
+    ScoreThreshold,
+)
+from .pipeline import IndexingPipeline, QueryPipeline, RagPipeline, TraceEvent
+from .storage import (
+    BM25Index,
+    BlobStore,
+    LexicalIndex,
+    LocalBlobStore,
+    MemoryVectorStore,
+    MinioBlobStore,
+    QdrantVectorStore,
+    VectorStore,
+)
+
+__version__ = "0.6.0"
+
+__all__ = [
+    "ingest",
+    "registry",
+    "Component",
+    "Source",
+    "SourceFormat",
+    "Page",
+    "PageSpan",
+    "Document",
+    "Chunk",
+    "Query",
+    "ScoredChunk",
+    "Citation",
+    "Answer",
+    "Parser",
+    "AutoParser",
+    "DoclingParser",
+    "PlainTextParser",
+    "OcrEngine",
+    "OcrPolicy",
+    "OcrResult",
+    "PageImage",
+    "detect_format",
+    "Chunker",
+    "FixedChunker",
+    "MarkdownChunker",
+    "Embedder",
+    "SparseEncoder",
+    "CachingEmbedder",
+    "HashingEmbedder",
+    "SentenceTransformerEmbedder",
+    "SparseVector",
+    "VectorSpec",
+    "VectorValue",
+    "ChunkIndex",
+    "ChunkSink",
+    "DocumentCatalog",
+    "DocumentRef",
+    "Enricher",
+    "HeadingEnricher",
+    "ContextualEnricher",
+    "BlobStore",
+    "LocalBlobStore",
+    "MinioBlobStore",
+    "VectorStore",
+    "MemoryVectorStore",
+    "QdrantVectorStore",
+    "LexicalIndex",
+    "BM25Index",
+    "Retriever",
+    "IndexRetriever",
+    "FusionRetriever",
+    "HybridRetriever",
+    "MultiQueryRetriever",
+    "HydeRetriever",
+    "Refiner",
+    "CrossEncoderReranker",
+    "KeywordRefiner",
+    "NeighborExpander",
+    "ScoreThreshold",
+    "Generator",
+    "ExtractiveGenerator",
+    "AnthropicGenerator",
+    "IndexingPipeline",
+    "QueryPipeline",
+    "RagPipeline",
+    "TraceEvent",
+    "RagToolkitError",
+    "StorageError",
+    "EmbeddingError",
+    "GenerationError",
+    "EnrichmentError",
+]
+
+
+def ingest(path: str | Path, **docling_overrides) -> Document:
+    """One-call Facade: file path → markdown Document with provenance.
+
+    Keyword arguments are forwarded to the PDF/office parser (DoclingParser),
+    since that is where 95% of tuning happens:
+
+        ingest("scan.pdf", ocr_engine="mistral", page_batch_size=4)
+
+    For full control (custom routes, streaming, in-memory sources), drop one
+    level down to AutoParser / DoclingParser directly.
+    """
+    parser_configs = {"docling": docling_overrides} if docling_overrides else {}
+    parser = AutoParser(parser_configs=parser_configs)
+    return parser.parse(Source.from_path(path))

rag_blocks/chunking/__init__.py

@@ -0,0 +1,17 @@
+"""Chunking subsystem: Document → stream of retrieval Chunks.
+
+Importing this package registers the built-in chunkers (module import is the
+registration side effect the registry relies on). Strategies decide only WHERE
+to cut (char-offset spans); the Template Method in `base.Chunker.chunk` owns all
+bookkeeping and provenance.
+"""
+
+from .base import Chunker
+from .fixed import FixedChunker
+from .markdown import MarkdownChunker
+
+__all__ = [
+    "Chunker",
+    "FixedChunker",
+    "MarkdownChunker",
+]

rag_blocks/chunking/base.py

@@ -0,0 +1,77 @@
+"""Chunker: the retrieval-unit Strategy interface.
+
+The single design decision mirrored from the Parser: the *primitive* a strategy
+implements is `iter_spans(document) -> Iterator[(start, end)]` — WHERE to cut,
+as half-open character offsets into `document.markdown`. `chunk()` is a Template
+Method layered on top that does every piece of bookkeeping once, correctly, for
+every strategy.
+
+Why spans (coordinates), not strings (copies)?
+    Return strings and you throw away everything that makes RAG citations work:
+    provenance (which chars → which pages), overlap (two chunks sharing a
+    region), and neighbor merging (fetch index±1 at query time). All three fall
+    out naturally when a strategy only decides *coordinates* and the base class
+    owns the slicing. This is the whole reason the interface emits offsets.
+
+What the Template Method guarantees (so strategies never re-implement it):
+    - `text = document.markdown[start:end]` — the base slices, not the strategy.
+    - whitespace-only slices are skipped WITHOUT advancing the index, so
+      `index` stays contiguous 0-based with NO holes (a manual counter, never
+      `enumerate` over raw spans). Query-time neighbor expansion fetches
+      `index ± 1`, so a hole would silently drop a real neighbor.
+    - `id = f"{document.id}:{index}"` — deterministic ⇒ re-indexing overwrites
+      instead of duplicating.
+    - `char_start`/`char_end` are the primary provenance; `page_start`/
+      `page_end` are derived from them via `Document.pages_for_span` and are
+      ALWAYS filled for a doc-derived chunk.
+
+Contract for strategies: yield spans in reading order of `start`. Overlapping
+spans are LEGAL (that is how overlap strategies express themselves — in
+coordinates). Strategies are configured by their `Config`, never by subclassing.
+"""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Iterator
+
+from ..core.component import Component
+from ..core.contracts import Chunk, Document
+
+__all__ = ["Chunker"]
+
+
+class Chunker(Component):
+    """Turns a Document into a stream of retrieval Chunks."""
+
+    kind = "chunker"
+
+    @abstractmethod
+    def iter_spans(self, document: Document) -> Iterator[tuple[int, int]]:
+        """Yield half-open `(start, end)` char offsets into
+        `document.markdown`, in reading order of `start`. Overlaps allowed."""
+
+    def chunk(self, document: Document) -> Iterator[Chunk]:
+        """Template Method: turn spans into Chunks, owning all bookkeeping.
+
+        Strategies implement `iter_spans`; this method — written once — is the
+        only place chunk ids, index contiguity, and provenance are decided.
+        """
+        index = 0
+        for start, end in self.iter_spans(document):
+            text = document.markdown[start:end]
+            if not text.strip():
+                # Skip empties WITHOUT advancing index (contiguity, no holes).
+                continue
+            pages = document.pages_for_span(start, end)
+            yield Chunk(
+                id=f"{document.id}:{index}",
+                doc_id=document.id,
+                text=text,
+                index=index,
+                char_start=start,
+                char_end=end,
+                page_start=pages[0] if pages else None,
+                page_end=pages[-1] if pages else None,
+            )
+            index += 1

rag_blocks/chunking/fixed.py

@@ -0,0 +1,69 @@
+"""FixedChunker: fixed-size character windows with overlap.
+
+The workhorse baseline: walk the markdown in windows of ~`chunk_chars`, with
+`overlap_chars` of trailing context repeated into the next window so a fact
+split across a boundary still lands whole in at least one chunk.
+
+The one bit of finesse — cut on a boundary, not mid-sentence — reuses the exact
+logic proven in `PlainTextParser._cut_point`: prefer to end a window at a
+paragraph break (`\n\n`), else a line break (`\n`), but REFUSE a soft cut that
+would leave a chunk shorter than half the target (otherwise a document full of
+short lines would produce a swarm of tiny chunks). No boundary in range ⇒ a
+hard cut at `chunk_chars`.
+
+Overlap is expressed purely in coordinates (the next window starts
+`overlap_chars` before the previous end), which is exactly why the Chunker
+interface emits spans rather than strings — overlap and provenance survive.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterator
+
+from ..core.contracts import Document
+from ..core.registry import registry
+from .base import Chunker
+
+__all__ = ["FixedChunker"]
+
+
+@registry.register
+class FixedChunker(Chunker):
+    name = "fixed"
+    version = "0.1.0"
+
+    @dataclass
+    class Config:
+        chunk_chars: int = 1600
+        overlap_chars: int = 200
+
+    def iter_spans(self, document: Document) -> Iterator[tuple[int, int]]:
+        text = document.markdown
+        n = len(text)
+        start = 0
+        while start < n:
+            end = self._cut_end(text, start, n)
+            yield start, end
+            if end >= n:
+                break
+            # Step back by the overlap for the next window; guarantee forward
+            # progress even under a pathological overlap >= window size.
+            next_start = end - self.config.overlap_chars
+            start = next_start if next_start > start else end
+
+    def _cut_end(self, text: str, start: int, n: int) -> int:
+        """Where to end the window that begins at `start` (mirror of
+        PlainTextParser._cut_point, adapted to absolute offsets)."""
+        hard_end = start + self.config.chunk_chars
+        if hard_end >= n:
+            return n
+        # Only accept a soft cut in the back half of the window.
+        floor = start + self.config.chunk_chars // 2
+        para = text.rfind("\n\n", floor, hard_end)
+        if para != -1:
+            return para + 2
+        line = text.rfind("\n", floor, hard_end)
+        if line != -1:
+            return line + 1
+        return hard_end

rag_blocks/chunking/markdown.py

@@ -0,0 +1,61 @@
+"""MarkdownChunker: cut on heading boundaries.
+
+This is the payoff of normalizing every input to markdown during ingestion:
+document *structure* survives all the way to the cutting decision. Instead of
+slicing blindly every N characters, we cut at ATX headings (`#`..`######`), so
+each chunk is a coherent section that starts with its own heading — far better
+retrieval units than arbitrary windows.
+
+Coordinates only (like every Chunker): we locate heading line offsets and emit
+the spans *between* them. Content before the first heading is its own span; a
+document with no headings yields a single span covering the whole thing. The
+base Template Method then skips any whitespace-only section without disturbing
+index contiguity.
+
+Deliberately simple for v0.2: no secondary size cap on a very long section yet
+(a `max_chars` split is an easy, non-breaking add later). Structure first.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Iterator
+
+from ..core.contracts import Document
+from ..core.registry import registry
+from .base import Chunker
+
+__all__ = ["MarkdownChunker"]
+
+#: An ATX heading: 1–6 '#', then whitespace, then the title. Up to three
+#: leading spaces are allowed by CommonMark; `# no-space` is NOT a heading.
+_HEADING = re.compile(r"^ {0,3}#{1,6}\s")
+
+
+def _is_heading(line: str) -> bool:
+    return _HEADING.match(line) is not None
+
+
+@registry.register
+class MarkdownChunker(Chunker):
+    name = "markdown-aware"
+    version = "0.1.0"
+
+    def iter_spans(self, document: Document) -> Iterator[tuple[int, int]]:
+        text = document.markdown
+        n = len(text)
+        if n == 0:
+            return
+
+        # Collect cut points: the document start, plus every heading's offset.
+        cut_points = [0]
+        offset = 0
+        for line in text.splitlines(keepends=True):
+            if offset != 0 and _is_heading(line):
+                cut_points.append(offset)
+            offset += len(line)
+        cut_points.append(n)
+
+        for start, end in zip(cut_points, cut_points[1:]):
+            if end > start:  # consecutive headings can coincide; skip empties
+                yield start, end

rag_blocks/core/__init__.py

@@ -0,0 +1,66 @@
+"""Core layer: contracts, component model, registry, errors.
+
+Nothing in here knows about RAG specifics — it is the framework the stages
+are built on. Zero third-party dependencies by design.
+"""
+
+from .component import Component
+from .contracts import (
+    Answer,
+    Chunk,
+    Citation,
+    Document,
+    Page,
+    PageSpan,
+    Query,
+    ScoredChunk,
+    Source,
+    SourceFormat,
+    SparseVector,
+    VectorSpec,
+    VectorValue,
+)
+from .errors import (
+    ComponentNotFoundError,
+    ConfigError,
+    DuplicateComponentError,
+    EmbeddingError,
+    EnrichmentError,
+    GenerationError,
+    OcrError,
+    ParseError,
+    RagToolkitError,
+    StorageError,
+    UnsupportedFormatError,
+)
+from .registry import Registry, registry
+
+__all__ = [
+    "Component",
+    "Registry",
+    "registry",
+    "Source",
+    "SourceFormat",
+    "Page",
+    "PageSpan",
+    "Document",
+    "Chunk",
+    "Query",
+    "ScoredChunk",
+    "Citation",
+    "Answer",
+    "SparseVector",
+    "VectorSpec",
+    "VectorValue",
+    "RagToolkitError",
+    "ComponentNotFoundError",
+    "DuplicateComponentError",
+    "ConfigError",
+    "UnsupportedFormatError",
+    "ParseError",
+    "OcrError",
+    "StorageError",
+    "EmbeddingError",
+    "GenerationError",
+    "EnrichmentError",
+]

rag_blocks/core/component.py

@@ -0,0 +1,139 @@
+"""Component: the common base of every swappable building block.
+
+Why one base class?
+-------------------
+Every stage implementation (a parser, an OCR engine, a chunker, a reranker…)
+needs the same plumbing: a (kind, name) identity for the registry, a typed
+config, and a stable *fingerprint*. Centralizing that here keeps each concrete
+component focused purely on its algorithm (Single Responsibility), and it
+gives the evaluation suite a uniform way to describe/cache/compare components
+without knowing anything about them (Dependency Inversion: the tuner depends
+on `Component`, never on `DoclingParser`).
+
+The fingerprint is the quiet superpower: sha256(kind | name | version |
+canonical-config). Two pipeline variants that share the same parser config
+share the same fingerprint ⇒ the tuner can reuse the cached parse output
+instead of re-parsing 5 GB of PDFs per combination. Bump `version` whenever
+you change a component's behavior so stale caches invalidate themselves.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import json
+from abc import ABC
+from enum import Enum
+from typing import Any, ClassVar, Optional, Type
+
+from .errors import ConfigError
+
+
+def _plain(value: Any) -> Any:
+    """Normalize config values to log/JSON-friendly primitives.
+
+    Enums become their `.value` (so `OcrPolicy.AUTO` logs as "auto"), and
+    containers are walked recursively. Keeps describe() output — and thus
+    trial logs — clean and diffable.
+    """
+    if isinstance(value, Enum):
+        return value.value
+    if isinstance(value, dict):
+        return {k: _plain(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_plain(v) for v in value]
+    return value
+
+__all__ = ["Component"]
+
+# Config field names containing these substrings are redacted from
+# describe()/fingerprint(): secrets must never leak into logs or cache keys,
+# and rotating an API key must not invalidate caches.
+_SECRET_MARKERS = ("key", "token", "secret", "password", "credential")
+
+
+class Component(ABC):
+    """Base class for every pluggable building block.
+
+    Subclass contract:
+        kind:    the stage slot this fills ("parser", "ocr", "chunker", ...)
+        name:    unique name within the kind ("docling", "mistral", ...)
+        version: bump on behavior changes (cache invalidation)
+        Config:  optional dataclass type describing the configuration
+    """
+
+    kind: ClassVar[str]
+    name: ClassVar[str]
+    version: ClassVar[str] = "0.1.0"
+    Config: ClassVar[Optional[Type[Any]]] = None
+
+    #: The resolved config instance (or None when Config is undeclared).
+    #: Declared as Any so subclasses read `self.config.<field>` without mypy
+    #: narrowing it to `Any | None` at every access — the config's real shape
+    #: is the nested `Config` dataclass, checked where it is constructed.
+    config: Any
+
+    def __init__(self, config: Any = None, **overrides: Any) -> None:
+        if self.Config is None:
+            if config is not None or overrides:
+                raise ConfigError(
+                    f"{type(self).__name__} takes no configuration"
+                )
+            self.config = None
+            return
+
+        if config is not None and not isinstance(config, self.Config):
+            raise ConfigError(
+                f"{type(self).__name__} expected config of type "
+                f"{self.Config.__name__}, got {type(config).__name__}"
+            )
+
+        # Ergonomics: accept a ready Config, keyword overrides, or both.
+        # `MyParser(page_batch_size=4)` beats building a config object by hand.
+        base = config if config is not None else self._default_config()
+        try:
+            self.config = (
+                dataclasses.replace(base, **overrides) if overrides else base
+            )
+        except TypeError as exc:  # unknown field name → clear error, fail fast
+            raise ConfigError(f"{type(self).__name__}: {exc}") from exc
+
+    def _default_config(self) -> Any:
+        try:
+            return self.Config()  # type: ignore[misc]
+        except TypeError as exc:
+            raise ConfigError(
+                f"{type(self).__name__}.Config has required fields; "
+                f"pass them explicitly: {exc}"
+            ) from exc
+
+    # -- identity ------------------------------------------------------------
+
+    def describe(self) -> dict:
+        """Loggable, secret-free description of this exact component setup.
+
+        This dict is what lands in every evaluation Trial record, so a result
+        is always reproducible from its log line alone.
+        """
+        cfg: dict[str, Any] = {}
+        if self.config is not None:
+            for k, v in dataclasses.asdict(self.config).items():
+                lowered = k.lower()
+                if any(m in lowered for m in _SECRET_MARKERS):
+                    cfg[k] = "<redacted>"
+                else:
+                    cfg[k] = _plain(v)
+        return {
+            "kind": self.kind,
+            "name": self.name,
+            "version": self.version,
+            "config": cfg,
+        }
+
+    def fingerprint(self) -> str:
+        """Stable hash of (kind, name, version, config) — the cache key."""
+        canonical = json.dumps(self.describe(), sort_keys=True, default=str)
+        return hashlib.sha256(canonical.encode()).hexdigest()[:16]
+
+    def __repr__(self) -> str:  # helpful in logs and notebooks
+        return f"<{type(self).__name__} {self.kind}:{self.name} v{self.version}>"