Skip to main content

Core utilities for document processing, RAG configuration, querying, and evaluation.

Project description

ragbandit-core

Test Coverage

Core utilities for:

  • Document ingestion & processing (OCR, chunking, embedding)
  • Building and running Retrieval-Augmented Generation (RAG) pipelines
  • Evaluating answers with automated metrics

Test Coverage

The codebase maintains 87% test coverage with comprehensive integration tests covering all major components. See tests/README.md for details on running tests and coverage reports.

Quick start

pip install ragbandit-core
from ragbandit.documents import (
    DocumentPipeline,
    ReferencesRefiner,
    FootnoteRefiner,
    MistralOCR,
    MistralEmbedder,
    SemanticChunker
)
import os
import logging
from dotenv import load_dotenv
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)

MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")

file_path = "./data/raw/[document_name].pdf"

doc_pipeline = DocumentPipeline(
    chunker=SemanticChunker(api_key=MISTRAL_API_KEY, min_chunk_size=500),
    embedder=MistralEmbedder(api_key=MISTRAL_API_KEY, model="mistral-embed"),
    ocr_processor=MistralOCR(api_key=MISTRAL_API_KEY),
    refiners=[
        ReferencesRefiner(api_key=MISTRAL_API_KEY),
        FootnoteRefiner(api_key=MISTRAL_API_KEY),
    ],
)

extended_response = doc_pipeline.process(file_path)

Using Alternative OCR and Embedding Providers

The package supports multiple OCR and embedding providers:

from ragbandit.documents import (
    DocumentPipeline,
    DatalabOCR,
    OpenAIEmbedder,
    FixedSizeChunker
)
import os

DATALAB_API_KEY = os.getenv("DATALAB_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

file_path = "./data/raw/[document_name].pdf"

# Using Datalab OCR and OpenAI embeddings
doc_pipeline = DocumentPipeline(
    ocr_processor=DatalabOCR(
        api_key=DATALAB_API_KEY,
        model="marker",
        mode="balanced"  # Options: fast, balanced, accurate
    ),
    chunker=FixedSizeChunker(chunk_size=500, overlap=100),
    embedder=OpenAIEmbedder(
        api_key=OPENAI_API_KEY,
        model="text-embedding-3-small"  # or text-embedding-3-large
    ),
)

result = doc_pipeline.process(file_path)

Running Steps Manually

For more control, you can run each pipeline step independently:

from ragbandit.documents import (
    DocumentPipeline,
    ReferencesRefiner,
    MistralOCR,
    MistralEmbedder,
    SemanticChunker
)
import os
from dotenv import load_dotenv
load_dotenv()

MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
file_path = "./data/raw/[document_name].pdf"

# Create pipeline with only the components you need
pipeline = DocumentPipeline(
    ocr_processor=MistralOCR(api_key=MISTRAL_API_KEY),
    refiners=[ReferencesRefiner(api_key=MISTRAL_API_KEY)],
    chunker=SemanticChunker(api_key=MISTRAL_API_KEY, min_chunk_size=500),
    embedder=MistralEmbedder(api_key=MISTRAL_API_KEY, model="mistral-embed"),
)

# Step 1: Run OCR
ocr_result = pipeline.run_ocr(file_path)

# Step 2: Run refiners (optional)
refining_results = pipeline.run_refiners(ocr_result)
final_doc = refining_results[-1]  # Get the last refiner's output

# Step 3: Chunk the document
chunk_result = pipeline.run_chunker(final_doc)

# Step 4: Embed chunks
embedding_result = pipeline.run_embedder(chunk_result)

You can also use components independently without a pipeline:

# Run OCR directly - Mistral
ocr = MistralOCR(api_key=MISTRAL_API_KEY)
ocr_result = ocr.process(file_path)

# Or use Datalab OCR
from ragbandit.documents import DatalabOCR
datalab_ocr = DatalabOCR(
    api_key=DATALAB_API_KEY,
    mode="accurate",
    max_pages=10  # Optional: limit pages processed
)
ocr_result = datalab_ocr.process(file_path)

# Run refiners directly
refiner = FootnoteRefiner(api_key=MISTRAL_API_KEY)
refined_result = refiner.process(ocr_result)

# Run chunker directly
chunker = SemanticChunker(api_key=MISTRAL_API_KEY, min_chunk_size=500)
chunk_result = chunker.chunk(refined_result)

# Run embedder directly - Mistral
embedder = MistralEmbedder(api_key=MISTRAL_API_KEY)
embedding_result = embedder.embed_chunks(chunk_result)

# Or use OpenAI embeddings
from ragbandit.documents import OpenAIEmbedder
openai_embedder = OpenAIEmbedder(
    api_key=OPENAI_API_KEY,
    model="text-embedding-3-large"  # Higher quality, larger dimensions
)
embedding_result = openai_embedder.embed_chunks(chunk_result)

Available Components

OCR

Class Provider Models Key params
MistralOCR Mistral mistral-ocr-2512 (default) api_key, model
DatalabOCR Datalab marker api_key, mode (fast / balanced / accurate), max_pages, page_range

Refiners

Class What it does
ReferencesRefiner Detects and extracts the references/bibliography section. Stores in extracted_data["references_markdown"].
FootnoteRefiner Detects footnotes, inlines explanations, and collects citations.
TableOfContentsRefiner Detects and removes the table of contents. Stores in extracted_data["toc_markdown"].

Chunkers

Class Params (defaults) When to use
FixedSizeChunker chunk_size=1000, overlap=200 Fast, deterministic splitting by character count
SentenceChunker sentences_per_chunk=5, sentence_overlap=1, min_chunk_size=100 Sentence-aware sliding window, no external deps
RecursiveMarkdownChunker chunk_size=1000, overlap=100 Heading-aware hierarchical splitting (H1→H2→H3→H4→paragraph→sentence)
SemanticChunker api_key, min_chunk_size=500 LLM-based semantic boundary detection (uses Mistral)

Embedders

Class Provider Models Cost / 1M tokens
MistralEmbedder Mistral mistral-embed $0.10
OpenAIEmbedder OpenAI text-embedding-3-small, text-embedding-3-large $0.02 / $0.13
VoyageAIEmbedder Voyage AI voyage-3.5 (default), voyage-3.5-lite, voyage-3-large, voyage-3, voyage-3-lite $0.06 / $0.02 / $0.18 / $0.06 / $0.02
CohereEmbedder Cohere embed-v4.0 $0.12

Examples & Notebooks

Example scripts (examples/)

File What it shows
01_basic_pipeline.py End-to-end DocumentPipeline.process() with MistralOCR + FixedSizeChunker + MistralEmbedder
02_choosing_components.py Same doc with two combos (Mistral-only vs mixed providers) — compares chunks, dims, cost
03_step_by_step.py Manual run_ocr()run_refiners()run_chunker()run_embedder() with intermediate inspection
04_cost_tracking.py TokenUsageTracker standalone — per-model breakdown and total cost
.venv/bin/python examples/01_basic_pipeline.py

Notebooks (notebooks/)

File What it shows
getting_started.ipynb Full pipeline walkthrough — one cell per stage
component_comparison.ipynb Compares FixedSize vs Sentence vs RecursiveMarkdown chunking strategies
component_explorer.ipynb Exercises every component with all valid configurations

Each notebook has a setup cell where you set PDF_PATH and ENV_PATH to point to your own document and API keys.

Package layout

ragbandit-core/
├── examples/          # Runnable example scripts
├── notebooks/         # Jupyter notebooks
├── src/ragbandit/
│   ├── config/        # Pricing and model configuration
│   ├── documents/     # Document ingestion, OCR, chunking, embedding
│   │   ├── chunkers/
│   │   ├── embedders/
│   │   ├── ocr/
│   │   └── refiners/
│   ├── prompt_tools/  # LLM-based tools
│   └── utils/         # Token tracking, logging, client managers
└── tests/

License

MIT

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

ragbandit_core-0.4.0.tar.gz (49.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

ragbandit_core-0.4.0-py3-none-any.whl (69.7 kB view details)

Uploaded Python 3

File details

Details for the file ragbandit_core-0.4.0.tar.gz.

File metadata

  • Download URL: ragbandit_core-0.4.0.tar.gz
  • Upload date:
  • Size: 49.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for ragbandit_core-0.4.0.tar.gz
Algorithm Hash digest
SHA256 e4e35dd6584263ceabe427465d5e240e9faf0f5b89756997e2743e3b60bfe94f
MD5 0474c30741fc9f00bb45419cfc00414c
BLAKE2b-256 d47184a435791f9e724b8b663036a2bf05be4d9a77908eb8f7dcb0b11272bf0a

See more details on using hashes here.

File details

Details for the file ragbandit_core-0.4.0-py3-none-any.whl.

File metadata

  • Download URL: ragbandit_core-0.4.0-py3-none-any.whl
  • Upload date:
  • Size: 69.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for ragbandit_core-0.4.0-py3-none-any.whl
Algorithm Hash digest
SHA256 8b1daeb130f292d7a47968a1d549c5b1a90066252619dc163a1d8ffd08fa6e9c
MD5 d0ff9aeb4e719fd64e6bf09b9c75f442
BLAKE2b-256 28af3069ad290ef632c317b6caec3a4e61154f5a9ee3a63a2fcb1c137c4714d6

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page