ai4rag

RAG Templates Optimization Engine

Initializes RAG Templates with optimal parameters

Getting Started • User Guide • API Reference • Development

🎯 What is ai4RAG?

ai4RAG is an optimization engine for RAG Templates that is LLM and vector database provider-agnostic. It accepts a variety of RAG Templates and a search space definition, then returns an initialized RAG Template with optimal parameter values (called a RAG Pattern).

Model providers

ai4rag reaches foundation and embedding models through the stock openai SDK, so it works with any OpenAI-compatible endpoint — a hosted API, a self-managed server (vLLM, TGI, Ollama, …), or an OpenShift AI Models-as-a-Service (MaaS) deployment. Prefer something else entirely? Implement BaseFoundationModel / BaseEmbeddingModel and pass your own models straight into an experiment.

MaaS is the integration ai4rag ships helpers for, so the walkthrough below uses it:

• SDK: openai >= 2, < 3 (Python package used by ai4RAG; installs with this project).
• Deployment: an OpenShift AI MaaS instance exposing at least one foundation model and one embedding model.
• Endpoints: MaaS serves everything from a single OpenAI-compatible endpoint — MAAS_BASE_URL (host-only or /v1-suffixed; normalized automatically). One client lists the available models (models.list()) and serves chat/completions and embeddings for all of them. Model ids are used verbatim, exactly as models.list() reports them.

Features used by ai4rag

When using the MaaS backend, ai4rag relies on:

• Embeddings — Text embeddings via the embeddings endpoint (e.g. for indexing and query encoding). Because models.list() carries no metadata, embedding dimension and context length are auto-detected at construction (or supplied via params).
• Chat / completions — Foundation model integration for answer generation when evaluating RAG patterns.

Vector storage is independent of MaaS: ai4rag connects directly to remote Milvus, embedded Milvus Lite, or PostgreSQL/pgvector via the config classes in ai4rag.rag.vector_store (see Vector stores below).

Vector stores

ai4RAG talks to the vector store directly through provider-specific clients — no MaaS deployment is required for this part. Pick a provider and pass its config to AI4RAGExperiment as vector_store_config:

• MilvusConfig — remote Milvus server or Zilliz Cloud only. uri must be a http(s):// URL (TLS and self-signed CAs via ca_cert); anything else (a bare host, a file path, an empty string) raises ValueError. This is a deliberate safety check: a mistyped or unreachable MILVUS_URI now fails loudly instead of silently falling back to a throwaway local database. Supports hybrid search (dense + BM25).
• MilvusLiteConfig — the embedded, zero-server Milvus Lite engine, backed by a local db_path file (default "./ai4rag_milvus_lite.db") — no setup required, ideal for local development and small-scale workloads. Also supports hybrid search (dense + BM25); rejects http(s):// values (use MilvusConfig for those).
• PGVectorConfig — PostgreSQL with the pgvector extension. Hybrid search (dense + tsvector full-text).
• Neo4jConfig — Neo4j (with APOC), for graph-based RAG. For a server certificate issued by a publicly trusted CA, use a neo4j+s:// or bolt+s:// URI. For a private or self-signed CA, supply ca_cert; ai4rag automatically configures encrypted, CA-verified transport. add_documents() extracts entities and relationships into a per-collection knowledge graph; search_mode="graph" is the only supported mode (no vector/hybrid search). See the Neo4j Indexing and Graph Search architecture doc.

Each config is a frozen dataclass with a .from_env() constructor and an env_vars attribute listing the environment variables it reads (e.g. MILVUS_URI for MilvusConfig, MILVUS_LITE_DB_PATH for MilvusLiteConfig, PGVECTOR_HOST for PGVectorConfig).

Document processing

ai4RAG uses docling-core for document representation and chunking. Documents are represented as DoclingDocument instances, and the DoclingChunker leverages docling's HybridChunker for structure-aware, token-aware chunking. docling-core, openai, langchain/langchain-openai, and the vector store clients (pymilvus with Milvus Lite, pgvector, asyncpg) are all installed automatically with ai4rag.

For OCR and audio ingestion, install ai4rag[text-extraction]. This extra includes rapidocr~=3.9.2, onnxruntime~=1.25.0, and huggingface-hub>=0.30,<2 along with Docling's conversion and audio dependencies. Both the core package and this extra require docling-slim~=2.121.0 for ai4rag and Kubeflow Pipelines compatibility. Generated indexing notebooks install OCR packages through this extra. Audio extraction also requires the system ffmpeg executable on PATH with support for the input codec.

extract_text() and the repository's development-only FileStore helper parse PDFs with PDFium to recover text that the native parser can emit as NUL characters. FileStore defaults PDF model inference to CPU; pass device="auto" or a specific accelerator to enable GPU inference. It automatically rebuilds legacy PDF caches using the .pdfium.json suffix.

RAG templates

Each candidate selects its template through the categorical template_id search parameter. Vector backends (Milvus, Milvus Lite, pgvector) accept agentic_rag (default) or sequential_rag; Neo4j requires agentic_graph_rag and graph search. Both agentic IDs use AgenticRAG; sequential uses SequentialRAG. The selected ID is retained in results, streamed patterns, inference notebooks, and starter kits.

AgenticRAG can retry a recognized context-window overflow once with query-relevant excerpts from retrieved documents. This fallback is enabled by default; see context compression for its limits and configuration.

LocalEventHandler writes conversations.json alongside each agentic pattern's pattern.json and evaluation_results.json, recording retrieved context and tool activity. Conversations are also available as EvaluationResult.conversations and included by to_dict(); the field is None when no conversations are emitted. Per-question evaluation records include an optional conversation_id linking to these records.

generate_starter_kit() packages the selected pattern into a microservice with OpenShell deployment for Milvus, PostgreSQL/pgvector, or Neo4j. Deployment forwards database CA certificates and retrieves from an already indexed collection. See the Starter Kit guide and template migration guidance.

Running on Disconnected Clusters

If you are running ai4rag on a disconnected cluster (no internet access), you must pre-download the required ML models before execution.

What to Pre-Download

ai4rag depends on models from two sources:

Setup Instructions

For the MaaS and knowledge-graph indexing notebooks, use the approved Docling artifact bundle from the documents-indexing pipeline image. A generic Docling download alone does not provide the pinned PP-OCRv4 files required for OCR parity. The selected OCR bundle must include detection, classification, recognition, and matching dictionary files under RapidOcr/. Only stage bundles needed by your corpus: PDF/image inputs need Docling/RapidOCR, audio needs Whisper, and text-only inputs need neither.

1. On an internet-connected machine, stage the approved Docling/RapidOCR bundle under ./offline/docling/ for PDF/image inputs. For audio, download the pinned Transformers-format Whisper Tiny snapshot:

   pip install 'ai4rag[text-extraction]'
   

   from huggingface_hub import snapshot_download
   
   snapshot_download(
       repo_id="openai/whisper-tiny",
       revision="169d4a4341b33bc18d8881c4b69c2e104e1cc0af",
       allow_patterns=["*.json", "*.txt", "*.safetensors"],
       local_dir="./offline/models/whisper-tiny",
   )
   

2. Transfer the required artifacts to your disconnected cluster:

   # Copy only the bundles required by your corpus.
   rsync -av ./offline/docling/ cluster:/offline/docling/
   rsync -av ./offline/models/whisper-tiny/ cluster:/offline/models/whisper-tiny/
   

3. On the disconnected cluster, set environment variables:

   export DOCLING_ARTIFACTS_PATH=/offline/docling
   export HF_MODEL_DIR=/offline/models/whisper-tiny
   export HF_HUB_OFFLINE=1  # Enforce offline mode
   export TRANSFORMERS_OFFLINE=1
   export RHOAI_NOTEBOOK_CLUSTER_MODE=disconnected
   

   HF_MODEL_DIR must contain a complete local Transformers-format Whisper model on disconnected clusters. Connected indexing notebooks can download the pinned Whisper model when no local model is configured. You can override the environment setting with DoclingExtractionConfig(asr_model_path="/models/whisper-tiny"). Provide ffmpeg in the workbench image for audio, and mirror the current ai4rag[text-extraction] dependencies in the disconnected environment.

Indexing Notebooks

Both indexing notebooks default to connected mode. Their early Disconnected cluster setup cell lists everything to prepare before running offline. Set RHOAI_NOTEBOOK_CLUSTER_MODE=disconnected in the workbench, or pass NOTEBOOK_CLUSTER_MODE=disconnected to an AutoX Job. There are no environment tags or cells to skip; Jupyter Run All uses the same mode setting.

After document discovery, prepare_extraction() from ai4rag.utils.data.text_extraction checks the selected file types and returns a PreparedExtraction with config, artifacts_path, needs_ocr, and needs_asr. PDFs and images need the pipeline-compatible PP-OCRv4 bundle; audio needs a Whisper model. Text-only corpora need neither. Connected runs download pinned OCR or Whisper files only when no local model path is configured. Disconnected runs must mount the approved OCR bundle under DOCLING_ARTIFACTS_PATH and, for audio, a complete Whisper model under HF_MODEL_DIR. The pattern's selected language chooses English or Chinese OCR automatically. Each notebook uses one extraction worker so a small workbench can process scanned PDFs and audio without exhausting memory. For direct Python use, pass disconnected=True to prepare_extraction() and forward its config and artifacts_path to extract_text(); notebook mode variables apply only to the generated notebooks. extract_text() itself does not download Whisper weights.

Quick start

1. Prepare a MaaS client to integrate with your models.
2. Prepare your knowledge base documents for the experiment.
3. Prepare benchmark_data.json with evaluation questions and answers.
4. Define and constrain your search space.
5. Configure the optimizer.
6. Create and run the experiment.

Prepare the MaaS client

To enable full integration with MaaS, build a single client that lists the available models and serves them all — ai4rag reuses it for every foundation and embedding model wrapper. The dev_utils helper create_dev_maas_client() reads MAAS_BASE_URL / MAAS_API_KEY and builds that client for you.

from dotenv import load_dotenv, find_dotenv
from dev_utils.utils import create_dev_maas_client

load_dotenv(find_dotenv())

client = create_dev_maas_client()  # reads MAAS_BASE_URL / MAAS_API_KEY

Prepare knowledge base documents

Prepare a set of documents to serve as the knowledge base for retrieval. Documents are represented as DoclingDocument instances (from the docling-core library). A local folder is all you need — ai4rag does not require object storage.

Convert the folder with Docling, naming each document by its path relative to that folder:

from pathlib import Path
from docling.document_converter import DocumentConverter

documents_root = Path("<path to the documents folder>")
converter = DocumentConverter()

documents = []
for file_path in sorted(p for p in documents_root.rglob("*") if p.is_file()):
    document = converter.convert(file_path).document
    # The document's key: what benchmark data references.
    document.name = str(file_path.relative_to(documents_root))
    documents.append(document)

Already keeping your corpus in a bucket? discover_documents() and extract_text() name each document by its full S3 object key, so that is what correct_answer_document_keys must reference — prefix included. Discovery can span several locations in one bucket (prefixes=["manuals/", "reports/"]) and fails with a BenchmarkKeyError if a benchmark key matches none of the documents it found.

Prepare benchmark_data.json

Create a benchmark_data.json file following this schema:

[
	{
		"question": "<question_1>",
		"correct_answers": [
			"<answer 1 for question 1>",
			"<answer 2 for question 1>"
		],
		"correct_answer_document_keys": ["<list of document keys based on which correct answers were generated>"]
	},
	{
		"question": "<question_2>",
		"correct_answers": [
			"<answer 1 for question 2>",
			"<answer 2 for question 2>"
		],
		"correct_answer_document_keys": ["<list of document keys based on which correct answers were generated>"]
	}
]

All benchmark questions and answers must be derived from your knowledge base documents, and every correct_answer_document_keys entry must equal the name of one of the documents you loaded above.

import pandas as pd

benchmark_data = pd.read_json("<path to benchmark_data.json>")

Define and constrain search space

The search space defines all possible parameter combinations, where each combination creates a unique RAG Pattern. During the experiment, the engine will optimize the RAG Pattern for the selected metric over the given search space, using an objective function to evaluate each configuration.

from ai4rag.search_space.src.parameter import Parameter
from ai4rag.search_space.src.search_space import AI4RAGSearchSpace
from dev_utils.utils import build_maas_model


search_space = AI4RAGSearchSpace(
    params=[
        Parameter(
            name="foundation_model",
            param_type="C",
            values=[build_maas_model(client, model_id="qwen3-8b-fp8-dynamic", model_type="llm")],
        ),
        Parameter(
            name="embedding_model",
            param_type="C",
            values=[
                build_maas_model(
                    client,
                    model_id="bge-m3",
                    model_type="embedding",
                    embedding_params={"embedding_dimension": 1024, "context_length": 8192},
                )
            ],
        ),
        Parameter(
            name="chunking_method",
            param_type="C",
            values=["recursive", "hybrid"],
        ),
        Parameter(
            name="chunk_size",
            param_type="C",
            values=[512, 1024, 2048],
        ),
        Parameter(
            name="chunk_overlap",
            param_type="C",
            values=[0, 128, 256],
        ),
    ]
)

Configure optimizer

You have full control over the optimization algorithm. Configure the GAMOptimizer by adjusting GAMOptSettings.

from ai4rag.core.hpo.gam_opt import GAMOptSettings

optimizer_settings = GAMOptSettings(
    max_evals=10, n_random_nodes=4
)

Run the experiment

Using the information from the previous steps, create an experiment and run the ai4rag optimization engine.

from ai4rag.core.experiment.experiment import AI4RAGExperiment
from ai4rag.rag.vector_store import MilvusConfig
from ai4rag.utils.event_handler import LocalEventHandler

experiment = AI4RAGExperiment(
    documents=documents,
    benchmark_data=benchmark_data,
    search_space=search_space,
    vector_store_config=MilvusConfig.from_env(),
    optimizer_settings=optimizer_settings,
    event_handler=LocalEventHandler(output_path="<local-path-to-store-your-output-files>"),
)

experiment.search()
best_eval = experiment.results.get_best_evaluations(k=1)[0]
print(best_eval)

print(f"Best pattern: {best_eval.pattern_name} (score: {best_eval.final_score})")

Contribution

Pull requests are very welcome! Make sure your patches are well tested. Ideally create a topic branch for every separate change you make.

Development setup

This project uses uv for dependency management.

# Clone the repository
git clone https://github.com/IBM/ai4rag.git
cd ai4rag

# Install all development dependencies
uv sync --extra dev

# Run tests
uv run pytest tests/unit/

# Check code style
uv run black --check ai4rag/
uv run pylint ai4rag/

# Build and serve documentation locally
uv run mkdocs serve

Pull request workflow

1. Fork the repo
2. Create your feature branch (git checkout -b my-new-feature)
3. Commit your changes (git commit -s -am 'Added some feature')
4. Push to the branch (git push origin my-new-feature)
5. Create new Pull Request

See more details in contributing section.