rag-integration

RAG integration plugin for mloda. Composes modular FeatureGroups into text and image processing pipelines with PII redaction, chunking, deduplication, embedding, FAISS vector search, retrieval, evaluation, and LLM response generation.

See the demo notebook for an interactive walkthrough or the CLI for command-line usage.

Project Structure

rag_integration/
  feature_groups/
    rag_pipeline/       # Text: source, PII, chunk, dedup, embed, index, retrieve, LLM
    image_pipeline/     # Image: source, PII, preprocess, dedup, embed
    datasets/           # BEIR (SciFact) and image (Flickr30k) loaders
    evaluation/         # Retrieval metrics (precision, recall, NDCG, MAP)
cli/                    # Command-line demo tools (see cli/README.md)
tests/                  # Unit and integration tests
docs/                   # Detailed guides

Quick Start

1. Define your documents

Documents are loaded through a DataCreator-based FeatureGroup. The simplest option is DictDocumentSource, which loads from a Python list:

documents = [
    {"doc_id": "1", "text": "Contact support@example.com for help."},
    {"doc_id": "2", "text": "Our office is at 123 Main St."},
]

2. Build the pipeline

Each pipeline stage is expressed as a feature name. Stages chain with __:

docs                                    # raw documents
docs__pii_redacted                      # PII removed
docs__pii_redacted__chunked             # text split into chunks
docs__pii_redacted__chunked__deduped    # duplicates removed
docs__pii_redacted__chunked__deduped__embedded  # vector embeddings

3. Run with mlodaAPI

from mloda.user import mlodaAPI, PluginCollector, Feature, Options
from mloda_plugins.compute_framework.base_implementations.python_dict.python_dict_framework import (
    PythonDictFramework,
)

from rag_integration.feature_groups.rag_pipeline import (
    DictDocumentSource,
    RegexPIIRedactor,
    FixedSizeChunker,
    ExactHashDeduplicator,
    MockEmbedder,
)

providers = {
    DictDocumentSource,
    RegexPIIRedactor,
    FixedSizeChunker,
    ExactHashDeduplicator,
    MockEmbedder,
}

results = mlodaAPI.run_all(
    features=["docs__pii_redacted__chunked__deduped__embedded"],
    compute_frameworks={PythonDictFramework},
    plugin_collector=PluginCollector.enabled_feature_groups(providers),
)

4. Configure pipeline stages

Use Options to select specific implementations and tune parameters:

feature = Feature(
    "docs__pii_redacted__chunked__deduped__embedded",
    options=Options(context={
        "redaction_method": "regex",        # or "simple", "pattern", "presidio"
        "chunking_method": "sentence",      # or "fixed_size", "paragraph", "semantic"
        "deduplication_method": "exact_hash",  # or "normalized", "ngram"
        "embedding_method": "sentence_transformer",  # or "hash", "tfidf", "mock"
        "chunk_size": 512,
        "chunk_overlap": 128,
    }),
)

Available Components

Text Pipeline

Image Pipeline

Connector families

Alongside the build-your-own stage pipeline, the connectors/ package wraps whole external open-source RAG tools under one mloda surface, organized into six families by query-contract shape (retrieve, rerank, generate, graph_rag, structured, orchestrator). You swap backends by changing options, not by rewriting a pipeline.

The two layers share one seam: the FAISS retrieval stage is the native dense path of the retrieve family (retrieve_backend="faiss"), and a stage and its connector counterpart emit the same passage / answer row shape under the same canonical feature name, so migrating between them is an option swap. See "Relationship to the stage pipeline" in docs/rag-connector-base-classes.md.

See feature_groups/connectors/README.md for the family map (per-family contract, backends, no-Docker concrete, and pedigree), runnable examples, and links to the contract suites. The design rationale is in docs/rag-connector-base-classes.md.

from mloda.user import mlodaAPI, Feature, Options, PluginCollector
from mloda_plugins.compute_framework.base_implementations.python_dict.python_dict_framework import (
    PythonDictFramework,
)
from rag_integration.feature_groups.connectors.retrieve import Bm25sRetriever

feature = Feature(
    "retrieved_passages",
    options=Options(context={
        "retrieve_backend": "bm25s",
        "query_text": "cat pet",
        "corpus": [
            {"doc_id": "d1", "text": "A cat is an independent and curious pet."},
            {"doc_id": "d2", "text": "Cars need regular engine oil and maintenance."},
        ],
        "top_k": 3,
    }),
)
results = mlodaAPI.run_all(
    [feature],
    compute_frameworks={PythonDictFramework},
    plugin_collector=PluginCollector.enabled_feature_groups({Bm25sRetriever}),
)

Install a family's backend with uv sync --extra connectors (or rerank / graph / structured / orchestrator).

Swapping one backend for another is an option change, not a pipeline rewrite. python -m cli.swap_demo runs that swap within a family (retrieve_backend="bm25s" -> "tfidf") and across families (retrieve vs orchestrator over the same inputs); the contract is written up under "Swapping backends" in docs/rag-connector-base-classes.md.

Installation

Clone the repository and install with uv:

git clone https://github.com/mloda-ai/rag_integration.git
cd rag_integration
uv venv
source .venv/bin/activate
uv sync --all-extras

To install only specific extras, use uv sync --extra <name>:

CLI

A command-line interface is available for running pipelines interactively. See cli/README.md for full usage.

python3 -m cli.rag_demo run --input cli/docs/ --pii regex --chunking sentence --embedding tfidf -v

Development Setup

uv venv
source .venv/bin/activate
uv sync --all-extras

Run all checks (pytest, ruff, mypy, bandit):

tox

Run individual checks

pytest
ruff format --check --line-length 120 .
ruff check .
mypy --strict --ignore-missing-imports .
bandit -c pyproject.toml -r -q .

Related

• Getting Started Guide for a detailed walkthrough
• GitHub Workflows for CI/CD setup and required secrets
• mloda core library
• mloda-registry for plugin guides and community plugins