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).

Important

ai4rag is designed to be provider-agnostic: user may provide his own implementation for foundation model, embedding model or vector store and use them for the experiment. Out of the box ai4rag is designed to work with OGX. To use the full capabilities of ai4rag, you'll need access to an OGX server configured with at least one foundation model, one embedding model, and a vector database.

OGX

ai4RAG can run experiments using an OGX server for embeddings, vector storage, and text generation. Use the official client and API docs to connect and extend:

• Client: ogx-client >= 0.7.0 (Python package used by ai4RAG; installs with this project).
• Server: OGX >= 0.7.0.
• API reference: OGX API docs — HTTP API used by the client.

Features used by ai4rag

When using the OGX backend, ai4rag relies on:

• Embeddings — Text embeddings via the client (e.g. for indexing and query encoding). See Embeddings API in the docs.
• Vector stores — Create, retrieve, and delete vector store instances (e.g. Milvus) with a chosen embedding model and dimension. See Vector stores in the API docs.
• Vector IO — Insert document chunks (with embeddings) into a store and run similarity search (query) for retrieval. See Vector IO and insert/query endpoints.
• Chat / responses — Foundation model integration for answer generation (e.g. chat completions or responses API) when evaluating RAG patterns.

Quick start

1. Provide an instance of ogx-client to integrate with OGX.
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 ogx-client

To enable full integration with OGX, instantiate an OgxClient. This allows ai4rag to use the models and vector stores available on your OGX server.

Tip

Store your credentials securely in a .env file.

import os
from dotenv import load_dotenv, find_dotenv
from ogx_client import OgxClient

client = OgxClient(base_url=os.getenv("BASE_URL"), api_key=os.getenv("API_KEY"))

Prepare knowledge base documents

Prepare a set of documents to serve as the knowledge base for retrieval. These documents will be used to ground the LLM's responses and should be stored in a local directory.

Note

If you are using the project locally, you can load documents using the FileStore class from the dev_utils module. Supported document formats can be found in the FileStore implementation.

from pathlib import Path
from dev_utils.file_store import FileStore

documents_path = Path("<path to the documents folder>")
documents = FileStore(documents_path).load_as_documents()

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_ids": ["<list of documents ids 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_ids": ["<list of documents ids based on which correct answers were generated>"]
	}
]

All benchmark questions and answers must be derived from your knowledge base documents.

from dev_utils.utils import read_benchmark_from_json

benchmark_data_path = Path("<path to benchmark_data.json>")
benchmark_data = read_benchmark_from_json(benchmark_data_path)

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 ai4rag.rag.foundation_models.ogx import OGXFoundationModel
from ai4rag.rag.embedding.ogx import OGXEmbeddingModel


search_space = AI4RAGSearchSpace(
    params=[
        Parameter(
            name="foundation_model",
            param_type="C",
            values=[OGXFoundationModel(model_id="ollama/llama3.2:3b", client=client)],
        ),
        Parameter(
            name="embedding_model",
            param_type="C",
            values=[
                OGXEmbeddingModel(
                    model_id="ollama/nomic-embed-text:latest",
                    client=client,
                    params={"embedding_dimension": 768, "context_length": 8192},
                )
            ]
        )
    ]
)

Tip

To run automatic models discovery with OGX you may use prepare_search_space_with_ogx() from ai4rag.search_space.prepare_search_space.

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.

Note

For OGX vector stores, use vector_store_type="ogx" and specify the provider with ogx_vector_io_provider_id (e.g., ogx_vector_io_provider_id="milvus", ogx_vector_io_provider_id="qdrant"). To use ChromaDB in-memory, specify vector_store_type="chroma".

from ai4rag.core.experiment.experiment import AI4RAGExperiment
from ai4rag.utils.event_handler import LocalEventHandler

experiment = AI4RAGExperiment(
    client=client,
    documents=documents,
    benchmark_data=benchmark_data,
    search_space=search_space,
    vector_store_type="ogx",
    ogx_vector_io_provider_id="milvus",
    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(best_eval.rag_pattern.generate("What ai4rag can be used for?"))

Tip

For production use, implement your own custom EventHandler to handle status changes and artifacts produced during the experiment. See the BaseEventHandler implementation for reference.

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.