ragfoundry 0.1.0__py3-none-any.whl

ragfoundry/__init__.py

@@ -0,0 +1,36 @@
+"""ragfoundry: configurable production RAG on LangChain and LangGraph."""
+
+from __future__ import annotations
+
+from .chunking import chunk_documents, register_chunker
+from .config import RAGConfig
+from .errors import (
+    ConfigurationError,
+    IngestionError,
+    MissingDependencyError,
+    RAGFoundryError,
+)
+from .evaluation import EvalReport, EvalSample, load_dataset
+from .graphs import IngestReport
+from .loaders import load_sources, register_loader
+from .pipeline import RAGPipeline, RAGResult
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ConfigurationError",
+    "EvalReport",
+    "EvalSample",
+    "IngestReport",
+    "IngestionError",
+    "MissingDependencyError",
+    "RAGConfig",
+    "RAGFoundryError",
+    "RAGPipeline",
+    "RAGResult",
+    "chunk_documents",
+    "load_dataset",
+    "load_sources",
+    "register_chunker",
+    "register_loader",
+]

ragfoundry/_utils.py

@@ -0,0 +1,88 @@
+"""Small helpers shared across the package."""
+
+from __future__ import annotations
+
+import importlib
+import json
+import re
+from typing import Any
+
+from langchain_core.messages import BaseMessage
+
+from .errors import MissingDependencyError
+
+_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
+_METADATA_KEY = re.compile(r"[^0-9a-zA-Z_]")
+
+
+def require(module: str, extra: str) -> Any:
+    """Import an optional module, or explain which extra provides it."""
+    try:
+        return importlib.import_module(module)
+    except ImportError as exc:
+        raise MissingDependencyError(
+            f"'{module}' is required for this feature. "
+            f"Install it with: pip install 'ragfoundry[{extra}]'"
+        ) from exc
+
+
+def message_text(message: BaseMessage | str) -> str:
+    """Return the plain text of a model reply, without reasoning blocks."""
+    content: Any = message if isinstance(message, str) else message.content
+    if isinstance(content, list):
+        parts = []
+        for block in content:
+            if isinstance(block, str):
+                parts.append(block)
+            elif isinstance(block, dict) and block.get("type") == "text":
+                parts.append(str(block.get("text", "")))
+        content = "".join(parts)
+    return _THINK_BLOCK.sub("", str(content)).strip()
+
+
+def parse_json(text: str) -> Any | None:
+    """Extract the first JSON array or object from model output, or ``None``."""
+    decoder = json.JSONDecoder()
+    for match in re.finditer(r"[\[{]", text):
+        try:
+            value, _ = decoder.raw_decode(text[match.start() :])
+        except json.JSONDecodeError:
+            continue
+        return value
+    return None
+
+
+def string_items(value: Any) -> list[str] | None:
+    """Read a model's JSON list of strings, tolerating ``[{"claim": "..."}]`` style objects.
+
+    Returns ``None`` when the value is not a list at all.
+    """
+    if not isinstance(value, list):
+        return None
+    items = []
+    for item in value:
+        if isinstance(item, dict):
+            item = next((v for v in item.values() if isinstance(v, str)), None)
+        if isinstance(item, str) and item.strip():
+            items.append(item.strip())
+    return items
+
+
+def sanitize_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
+    """Make metadata safe for every vector store.
+
+    Stores differ in what they accept, so keys are reduced to identifier characters, ``None``
+    values are dropped and non-scalar values are JSON encoded.
+    """
+    clean: dict[str, Any] = {}
+    for key, value in metadata.items():
+        if value is None:
+            continue
+        name = _METADATA_KEY.sub("_", str(key))
+        if not name or name[0].isdigit():
+            name = f"m_{name}"
+        if isinstance(value, (str, int, float, bool)):
+            clean[name] = value
+        else:
+            clean[name] = json.dumps(value, default=str, ensure_ascii=False)
+    return clean

ragfoundry/chunk_store.py

@@ -0,0 +1,109 @@
+"""SQLite record of indexed chunks.
+
+The record makes re-ingestion incremental, supplies the corpus for BM25, and detects an embedding
+model change that would otherwise mix incompatible vectors in one collection.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import sqlite3
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+from langchain_core.documents import Document
+
+from .errors import ConfigurationError
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS chunks (
+    id TEXT PRIMARY KEY,
+    source TEXT NOT NULL,
+    content TEXT NOT NULL,
+    metadata TEXT NOT NULL
+);
+CREATE INDEX IF NOT EXISTS idx_chunks_source ON chunks(source);
+CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
+"""
+
+
+class ChunkStore:
+    """Durable map of chunk ids to their content. Safe to share between threads."""
+
+    def __init__(self, persist_dir: str | Path, collection: str) -> None:
+        directory = Path(persist_dir)
+        directory.mkdir(parents=True, exist_ok=True)
+        self.path = directory / f"{re.sub(r'[^0-9A-Za-z_-]', '_', collection)}.chunks.sqlite"
+        with self._connect() as conn:
+            conn.executescript(_SCHEMA)
+
+    @contextmanager
+    def _connect(self) -> Iterator[sqlite3.Connection]:
+        conn = sqlite3.connect(self.path, timeout=30)
+        try:
+            with conn:
+                yield conn
+        finally:
+            conn.close()
+
+    def ensure_fingerprint(self, fingerprint: str) -> None:
+        """Bind the collection to one embedding model, or fail if it already has another."""
+        with self._connect() as conn:
+            row = conn.execute("SELECT value FROM meta WHERE key='fingerprint'").fetchone()
+            if row is None:
+                conn.execute("INSERT INTO meta VALUES ('fingerprint', ?)", (fingerprint,))
+            elif row[0] != fingerprint:
+                raise ConfigurationError(
+                    f"This collection was indexed with embeddings '{row[0]}' but the pipeline "
+                    f"is configured with '{fingerprint}'. Use a new collection_name or re-index."
+                )
+
+    def revision(self) -> int:
+        with self._connect() as conn:
+            row = conn.execute("SELECT value FROM meta WHERE key='revision'").fetchone()
+        return int(row[0]) if row else 0
+
+    def ids_for_source(self, source: str) -> set[str]:
+        with self._connect() as conn:
+            rows = conn.execute("SELECT id FROM chunks WHERE source=?", (source,)).fetchall()
+        return {row[0] for row in rows}
+
+    def sources(self) -> dict[str, int]:
+        with self._connect() as conn:
+            rows = conn.execute(
+                "SELECT source, COUNT(*) FROM chunks GROUP BY source ORDER BY source"
+            ).fetchall()
+        return {row[0]: row[1] for row in rows}
+
+    def documents(self) -> list[Document]:
+        with self._connect() as conn:
+            rows = conn.execute("SELECT content, metadata FROM chunks ORDER BY id").fetchall()
+        return [Document(page_content=row[0], metadata=json.loads(row[1])) for row in rows]
+
+    def upsert(self, chunks: Iterable[Document]) -> None:
+        rows = [
+            (
+                chunk.metadata["chunk_id"],
+                str(chunk.metadata.get("source", "")),
+                chunk.page_content,
+                json.dumps(chunk.metadata, default=str, ensure_ascii=False),
+            )
+            for chunk in chunks
+        ]
+        with self._connect() as conn:
+            conn.executemany("INSERT OR REPLACE INTO chunks VALUES (?, ?, ?, ?)", rows)
+            self._bump(conn)
+
+    def delete(self, ids: Iterable[str]) -> None:
+        with self._connect() as conn:
+            conn.executemany("DELETE FROM chunks WHERE id=?", [(item,) for item in ids])
+            self._bump(conn)
+
+    @staticmethod
+    def _bump(conn: sqlite3.Connection) -> None:
+        conn.execute(
+            "INSERT INTO meta VALUES ('revision', '1') "
+            "ON CONFLICT(key) DO UPDATE SET value = CAST(value AS INTEGER) + 1"
+        )

ragfoundry/chunking/__init__.py

@@ -0,0 +1,23 @@
+"""Chunking strategies: fixed, recursive, semantic and LLM-guided."""
+
+from __future__ import annotations
+
+from .base import (
+    ChunkerResources,
+    assign_chunk_ids,
+    build_splitter,
+    chunk_documents,
+    register_chunker,
+)
+from .llm import LLMChunker
+from .semantic import SemanticChunker
+
+__all__ = [
+    "ChunkerResources",
+    "LLMChunker",
+    "SemanticChunker",
+    "assign_chunk_ids",
+    "build_splitter",
+    "chunk_documents",
+    "register_chunker",
+]

ragfoundry/chunking/_text.py

@@ -0,0 +1,20 @@
+"""Sentence segmentation that keeps offsets into the original text."""
+
+from __future__ import annotations
+
+import re
+
+_BOUNDARY = re.compile(r"(?<=[.!?。！？])\s+|\n\s*\n")
+
+
+def sentence_spans(text: str) -> list[tuple[int, int]]:
+    """Return ``(start, end)`` offsets of non-blank sentences."""
+    spans = []
+    start = 0
+    for match in _BOUNDARY.finditer(text):
+        if text[start : match.start()].strip():
+            spans.append((start, match.start()))
+        start = match.end()
+    if text[start:].strip():
+        spans.append((start, len(text)))
+    return spans

ragfoundry/chunking/base.py

@@ -0,0 +1,171 @@
+"""Chunker registry and stable chunk identity."""
+
+from __future__ import annotations
+
+import hashlib
+import uuid
+from collections import Counter
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models import BaseChatModel
+from langchain_text_splitters import (
+    RecursiveCharacterTextSplitter,
+    TextSplitter,
+    TokenTextSplitter,
+)
+
+from .._utils import require
+from ..config import ChunkingConfig
+from ..errors import ConfigurationError
+from .llm import LLMChunker
+from .semantic import SemanticChunker
+
+_NAMESPACE = uuid.UUID("5f0c1e0a-7c55-4d0b-9c55-6f1f3f0d9a11")
+
+
+@dataclass
+class ChunkerResources:
+    """Models a strategy may need; resolved only by the strategies that use them."""
+
+    embeddings: Callable[[], Embeddings] | None = None
+    llm: Callable[[], BaseChatModel] | None = None
+
+    def get_embeddings(self) -> Embeddings:
+        if self.embeddings is None:
+            raise ConfigurationError("The 'semantic' strategy requires embeddings")
+        return self.embeddings()
+
+    def get_llm(self) -> BaseChatModel:
+        if self.llm is None:
+            raise ConfigurationError("The 'llm' strategy requires an LLM")
+        return self.llm()
+
+
+class FixedWindowSplitter(TextSplitter):
+    """Exact sliding windows of ``chunk_size`` characters overlapping by ``chunk_overlap``."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        kwargs["strip_whitespace"] = False
+        super().__init__(**kwargs)
+
+    def split_text(self, text: str) -> list[str]:
+        step = self._chunk_size - self._chunk_overlap
+        windows = (text[i : i + self._chunk_size] for i in range(0, len(text), step))
+        chunks = [window for window in windows if window.strip()]
+        # Drop a final window that lies entirely inside the previous one.
+        if len(chunks) > 1 and len(chunks[-1]) <= self._chunk_overlap:
+            chunks.pop()
+        return chunks
+
+
+ChunkerFactory = Callable[[ChunkingConfig, ChunkerResources], TextSplitter]
+_CHUNKERS: dict[str, ChunkerFactory] = {}
+
+
+def register_chunker(name: str) -> Callable[[ChunkerFactory], ChunkerFactory]:
+    """Register a strategy selectable through ``chunking.strategy``."""
+
+    def decorator(factory: ChunkerFactory) -> ChunkerFactory:
+        _CHUNKERS[name] = factory
+        return factory
+
+    return decorator
+
+
+def _length_function(config: ChunkingConfig) -> Callable[[str], int]:
+    if config.length_unit == "characters":
+        return len
+    tiktoken = require("tiktoken", "tokens")
+    encoding = tiktoken.get_encoding(config.encoding_name)
+    return lambda text: len(encoding.encode(text, disallowed_special=()))
+
+
+@register_chunker("fixed")
+def _fixed(config: ChunkingConfig, resources: ChunkerResources) -> TextSplitter:
+    if config.length_unit == "tokens":
+        require("tiktoken", "tokens")
+        return TokenTextSplitter(
+            encoding_name=config.encoding_name,
+            chunk_size=config.chunk_size,
+            chunk_overlap=config.chunk_overlap,
+            add_start_index=True,
+        )
+    return FixedWindowSplitter(
+        chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap, add_start_index=True
+    )
+
+
+@register_chunker("recursive")
+def _recursive(config: ChunkingConfig, resources: ChunkerResources) -> TextSplitter:
+    return RecursiveCharacterTextSplitter(
+        chunk_size=config.chunk_size,
+        chunk_overlap=config.chunk_overlap,
+        length_function=_length_function(config),
+        separators=config.separators,
+        add_start_index=True,
+    )
+
+
+@register_chunker("semantic")
+def _semantic(config: ChunkingConfig, resources: ChunkerResources) -> TextSplitter:
+    return SemanticChunker(
+        resources.get_embeddings(),
+        breakpoint_percentile=config.semantic.breakpoint_percentile,
+        buffer_size=config.semantic.buffer_size,
+        chunk_size=config.chunk_size,
+        chunk_overlap=config.chunk_overlap,
+        length_function=_length_function(config),
+    )
+
+
+@register_chunker("llm")
+def _llm(config: ChunkingConfig, resources: ChunkerResources) -> TextSplitter:
+    return LLMChunker(
+        resources.get_llm(),
+        window_chars=config.llm.window_chars,
+        chunk_size=config.chunk_size,
+        chunk_overlap=config.chunk_overlap,
+        length_function=_length_function(config),
+    )
+
+
+def build_splitter(
+    config: ChunkingConfig, resources: ChunkerResources | None = None
+) -> TextSplitter:
+    factory = _CHUNKERS.get(config.strategy)
+    if factory is None:
+        raise ConfigurationError(
+            f"Unknown chunking strategy '{config.strategy}'. "
+            f"Available: {', '.join(sorted(_CHUNKERS))}"
+        )
+    return factory(config, resources or ChunkerResources())
+
+
+def assign_chunk_ids(chunks: Sequence[Document]) -> list[Document]:
+    """Give each chunk an id derived from its source and text.
+
+    An unchanged chunk keeps its id across runs even if earlier text in the file changes, so
+    re-ingestion only embeds what is new. The occurrence count separates identical chunks.
+    """
+    occurrences: Counter[tuple[str, str]] = Counter()
+    for chunk in chunks:
+        source = str(chunk.metadata.get("source", ""))
+        digest = hashlib.sha256(chunk.page_content.encode("utf-8")).hexdigest()
+        occurrences[(source, digest)] += 1
+        key = f"{source}\0{digest}\0{occurrences[(source, digest)]}"
+        chunk.metadata["chunk_id"] = str(uuid.uuid5(_NAMESPACE, key))
+    return list(chunks)
+
+
+def chunk_documents(
+    documents: Sequence[Document],
+    config: ChunkingConfig | None = None,
+    resources: ChunkerResources | None = None,
+) -> list[Document]:
+    """Split documents with the configured strategy and assign stable ids."""
+    splitter = build_splitter(config or ChunkingConfig(), resources)
+    return assign_chunk_ids(splitter.split_documents(list(documents)))

ragfoundry/chunking/llm.py

@@ -0,0 +1,88 @@
+"""Ask a chat model where the topic boundaries in a text are."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from langchain_core.language_models import BaseChatModel
+from langchain_text_splitters import RecursiveCharacterTextSplitter, TextSplitter
+
+from .._utils import message_text, parse_json
+from ..prompts import LLM_CHUNK_PROMPT
+from ._text import sentence_spans
+
+logger = logging.getLogger(__name__)
+
+
+class LLMChunker(TextSplitter):
+    """Group sentences into topical chunks chosen by an LLM.
+
+    Long texts are processed in windows. The final group of each window is carried into the next,
+    so a topic is not cut merely because a window ended. If a reply cannot be parsed, that window
+    falls back to recursive splitting and a warning is logged.
+    """
+
+    def __init__(self, llm: BaseChatModel, window_chars: int = 8000, **kwargs: Any) -> None:
+        kwargs.setdefault("add_start_index", True)
+        super().__init__(**kwargs)
+        self._llm = llm
+        self._window_chars = window_chars
+        self._fallback = RecursiveCharacterTextSplitter(
+            chunk_size=self._chunk_size,
+            chunk_overlap=self._chunk_overlap,
+            length_function=self._length_function,
+        )
+
+    def split_text(self, text: str) -> list[str]:
+        spans = sentence_spans(text)
+        if len(spans) < 2:
+            return self._bounded(text)
+
+        chunks: list[str] = []
+        start = 0
+        while start < len(spans):
+            end = self._window_end(spans, start)
+            boundaries = self._ask(text, spans, start, end)
+            if boundaries is None:
+                logger.warning("Unusable LLM chunking reply; using recursive split for a window")
+                chunks.extend(self._fallback.split_text(text[spans[start][0] : spans[end - 1][1]]))
+                start = end
+                continue
+            finished = end == len(spans)
+            stop = end if finished or not boundaries else boundaries[-1]
+            edges = [start] + [b for b in boundaries if b < stop] + [stop]
+            for first, last in zip(edges, edges[1:], strict=False):
+                chunks.extend(self._bounded(text[spans[first][0] : spans[last - 1][1]]))
+            start = stop
+        return chunks
+
+    def _window_end(self, spans: list[tuple[int, int]], start: int) -> int:
+        end = start + 1
+        while end < len(spans) and spans[end][1] - spans[start][0] <= self._window_chars:
+            end += 1
+        return end
+
+    def _ask(
+        self, text: str, spans: list[tuple[int, int]], start: int, end: int
+    ) -> list[int] | None:
+        """Return sorted sentence indices inside ``(start, end)`` that open a chunk."""
+        if end - start < 2:
+            return []
+        numbered = "\n".join(
+            f"[{index}] {text[spans[index][0] : spans[index][1]].strip()}"
+            for index in range(start, end)
+        )
+        reply = self._llm.invoke(LLM_CHUNK_PROMPT.format(sentences=numbered))
+        parsed = parse_json(message_text(reply))
+        if not isinstance(parsed, list) or not all(isinstance(item, int) for item in parsed):
+            return None
+        return sorted({item for item in parsed if start < item < end})
+
+    def _bounded(self, text: str) -> list[str]:
+        text = text.strip()
+        if not text:
+            return []
+        if self._length_function(text) <= self._chunk_size:
+            return [text]
+        return self._fallback.split_text(text)

ragfoundry/chunking/semantic.py

@@ -0,0 +1,67 @@
+"""Split text where the embedding similarity of neighbouring sentences drops."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from langchain_core.embeddings import Embeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter, TextSplitter
+
+from ._text import sentence_spans
+
+
+class SemanticChunker(TextSplitter):
+    """Break at the largest semantic shifts; ``chunk_size`` is enforced as an upper bound."""
+
+    def __init__(
+        self,
+        embeddings: Embeddings,
+        breakpoint_percentile: float = 95,
+        buffer_size: int = 1,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.setdefault("add_start_index", True)
+        super().__init__(**kwargs)
+        self._embeddings = embeddings
+        self._percentile = breakpoint_percentile
+        self._buffer = buffer_size
+        self._oversize = RecursiveCharacterTextSplitter(
+            chunk_size=self._chunk_size,
+            chunk_overlap=self._chunk_overlap,
+            length_function=self._length_function,
+        )
+
+    def split_text(self, text: str) -> list[str]:
+        spans = sentence_spans(text)
+        if len(spans) < 3:
+            return self._bounded(text)
+
+        windows = [
+            text[
+                spans[max(0, i - self._buffer)][0] : spans[min(len(spans) - 1, i + self._buffer)][1]
+            ]
+            for i in range(len(spans))
+        ]
+        vectors = np.asarray(self._embeddings.embed_documents(windows), dtype=float)
+        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+        vectors = vectors / np.where(norms == 0, 1, norms)
+        distances = 1 - np.sum(vectors[:-1] * vectors[1:], axis=1)
+        threshold = np.percentile(distances, self._percentile)
+
+        chunks: list[str] = []
+        first = 0
+        for index, distance in enumerate(distances):
+            if distance >= threshold and distance > 0:
+                chunks.extend(self._bounded(text[spans[first][0] : spans[index][1]]))
+                first = index + 1
+        chunks.extend(self._bounded(text[spans[first][0] : spans[-1][1]]))
+        return chunks
+
+    def _bounded(self, text: str) -> list[str]:
+        text = text.strip()
+        if not text:
+            return []
+        if self._length_function(text) <= self._chunk_size:
+            return [text]
+        return self._oversize.split_text(text)