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.

RAG templates

Each candidate pattern is evaluated through a rag_template passed to AI4RAGExperiment — a BaseRAGTemplate subclass. The default, AgenticRAG, is a LangChain agent that can rewrite its search query and retrieve again before answering; SimpleRAG (a single fixed retrieve-then-generate pass) is also available via rag_template=SimpleRAG. See the RAG Components architecture doc.

Published patterns persist the selected template as template_id, which generated inference notebooks use to recreate it. Pattern publication currently supports the built-in AgenticRAG and SimpleRAG classes; custom subclasses need registry support. LocalEventHandler writes conversations.json alongside each agentic pattern's pattern.json and evaluation_results.json, recording retrieved context and tool activity.

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

1. On an internet-connected machine, download the artifacts:

   pip install 'ai4rag[text-extraction]'
   
   # Trigger Docling model download
   python -c "from docling.document_converter import DocumentConverter; \
             converter = DocumentConverter(); \
             converter.convert_document_string('/tmp/test.txt')"
   
   # Pre-download HuggingFace models
   export HF_HOME=/path/to/hf_cache
   python -c "from transformers import AutoTokenizer; \
             AutoTokenizer.from_pretrained('BAAI/bge-m3')"
   

2. Transfer the artifacts to your disconnected cluster:

   rsync -av ~/.cache/docling/ cluster:/offline/docling/
   rsync -av /path/to/hf_cache/ cluster:/offline/hf_cache/
   

3. On the disconnected cluster, set environment variables:

   export DOCLING_ARTIFACTS_PATH=/offline/docling
   export HF_MODEL_DIR=/offline/models/whisper-tiny
   export HF_HOME=/offline/hf_cache
   export HF_HUB_OFFLINE=1  # Enforce offline mode
   

   HF_MODEL_DIR must contain a complete local Transformers-format Whisper model. AI4RAG does not download audio-transcription weights at runtime. You can override the environment setting with DoclingExtractionConfig(asr_model_path="/models/whisper-tiny").

Indexing Notebooks

The provided maas_indexing_template.ipynb includes a "Prerequisites for Disconnected Clusters" section with:

• Validation cells to verify artifact availability
• Configuration examples for custom model paths
• Step-by-step download instructions in the appendix

The MaaS indexing notebook selects model requirements after document discovery: PDF/image inputs require DOCLING_ARTIFACTS_PATH, audio inputs require HF_MODEL_DIR, and text-only corpora require neither. The knowledge-graph indexing notebook also includes S3 discovery and extraction, and retains its Docling artifact prerequisite.

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.