<span id="version" style="color: grey; float: right">Version 0.1.0</span>

# RAG Optimization Plugin - Documentation

<br>

## Table of Contents

1. [Overview](#overview)
2. [How It Works](#how-it-works)
3. [Usage and Configuration](#usage-and-configuration)
4. [Main Features](#main-features)
5. [Best Practices](#best-practices)
6. [Limitations](#limitations)
7. [Technical Requirements](#technical-requirements)
8. [License](#license)



<br>

## Overview
---

The **RAG Optimization Plugin** is a Dataiku plugin that automatically optimizes RAG system configurations using Bayesian optimization techniques (Optuna). The plugin optimizes two main components:

1. **Embedding Optimization**: Optimizes document embedding recipe parameters
2. **RA-LLM Optimization**: Optimizes RA-LLM parameters for search and generation

### Objectives

- **Automate** the search for optimal RAG configurations
- **Maximize** answer quality (correctness)
- **Minimize** execution costs
- **Systematically test** different parameter combinations

<br>


## How It Works
---

### Key Concepts

- **Trial**: A single test run with one specific parameter combination
- **Study**: The complete optimization process, containing all trials
- **Objective Function**: The function that evaluates a configuration and returns a score (higher is better)
- **Score**: A composite metric combining correctness (answer quality) and cost efficiency
- **Combin**: Short for "combination" - a specific set of parameter values being tested



<br>

### What is Optuna?

**Optuna** is a hyperparameter optimization framework that uses Bayesian optimization to efficiently search for the best parameter combinations. Instead of testing every possible combination (which would be computationally expensive), Optuna:

1. **Learns from previous trials**: After each test, it analyzes which parameters led to better results
2. **Suggests promising configurations**: Uses machine learning to predict which parameter combinations are likely to perform well
3. **Focuses on promising regions**: Gradually narrows down the search space to areas with high potential

**Why Optuna?**
- **Efficient**: Finds good solutions with fewer trials compared to random or grid search
- **Adaptive**: Learns and adapts as it explores the parameter space
- **Smart sampling**: Uses Tree-structured Parzen Estimator (TPE) to intelligently suggest next parameters


<br>

### Plugin Workflow Overview

The plugin automates the optimization process through the following high-level steps:

#### 1. **Parameter Space Definition**
You define which parameters to optimize and their possible values. For example:
- Embedding chunk size: [2000, 3000, 4000] characters
- RAG search type: [SIMILARITY, MMR, SIMILARITY_THRESHOLD]

<br>

#### 2. **Trial Execution**
For each trial (test run):
- **Optuna suggests** a parameter combination based on previous results
- **Plugin creates** a temporary embedding recipe or RA-LLM with those parameters
- **All temporary objects** (recipes, Knowledge Banks, RA-LLMs) are stored in a dedicated **work flow zone** to keep the main flow clean
- **System evaluates** the configuration by:
  - Building a Knowledge Bank (for embeddings)
  - Running queries from your evaluation dataset
  - Getting answers from the RAG system
  - Comparing answers to expected outputs (correctness)
  - Calculating estimated costs
- **Plugin scores** the configuration: `score = f(correctness, cost)`
- **Temporary objects are cleaned up** at the end of the optimization process

<br>

#### 3. **Learning and Iteration**
- Optuna records the score and parameter combination
- It learns which parameters correlate with better scores
- It suggests more promising combinations for the next trial
- This repeats for `n_trials` iterations

<br>

#### 4. **Best Configuration Selection**
After all trials:
- Optuna identifies the configuration with the highest score
- The plugin creates or updates the final optimized recipe/LLM with these best parameters
- Results and metrics are returned for analysis

<br>



### Why This Approach Works

1. **Automated**: No manual parameter tuning needed
2. **Systematic**: Tests configurations in an intelligent order
3. **Balanced**: Considers both quality (correctness) and cost
4. **Adaptive**: Learns from each trial to improve suggestions
5. **Practical**: Returns ready-to-use optimized configurations

<br>

### Execution Flow

```
1. Optimization parameters configuration
   ↓
2. Embedding Optimization (Optuna)
   ├── Create embedding recipes with different configs
   ├── Evaluate each configuration
   └── Select best configuration
   ↓
3. Build optimized Knowledge Bank
   ↓
4. RA-LLM Optimization (Optuna)
   ├── Create RA-LLMs with different configs
   ├── Evaluate each configuration
   └── Select best configuration
   ↓
5. Return results and metrics
```

<br>

### Metrics and Scoring

#### Composite Score

The final score combines **correctness** and **cost efficiency**:

```python
if correctness < hard_min_correctness:
    score = -1.0  # Rejected
else:
    cost_score = 1.0 / (1.0 + np.log1p(cost * 1e6))
    score = 2 * correctness * cost_score / (correctness + cost_score + 1e-9)
```

**Interpretation**:
- High score = High correctness + Low cost
- Score = -1.0 = Insufficient correctness (rejected)

#### Returned Metrics

Each trial returns:
- `answerCorrectness`: Average correctness score (0-1)
- `estimatedCost`: Total estimated cost
- `duration_ms`: Execution duration in milliseconds
- `n_eval_rows`: Number of evaluated rows


<br>

### Optimization Parameters

#### Embedding Parameters

| Parameter | Type | Description | Possible Values |
|-----------|------|-------------|-----------------|
| `actionToPerform` | categorical | Extraction type | `['VLM', 'STRUCTURED']` |
| `chunkSizeCharacters` | categorical | Chunk size | `[2000, 3000, 4000]` |
| `chunkOverlapCharacters` | categorical | Overlap | `[0, 120, 300]` |
| `vlm.splitUnit` | categorical | VLM split unit | `['PAGE', 'CUSTOM']` |
| `vlm.useCustomPrompt` | categorical | Use custom prompt | `[False, True]` |
| `vlm.customNbPages` | categorical | Custom number of pages | `[1, 2, 3]` |
| `vlm.customPagesOverlap` | categorical | Custom pages overlap | `[0, 1]` |
| `structured.splitUnit` | categorical | STRUCTURED split unit | `['DOCUMENT', 'SECTION']` |
| `structured.imageHandlingMode` | categorical | Image handling | `['IGNORE', 'OCR', 'VLM_ANNOTATE']` |
| `structured.maxSectionDepth` | categorical | Max section depth | `[2, 6, 10]` |

<br>

#### RA-LLM Parameters

| Parameter | Type | Description | Possible Values |
|-----------|------|-------------|-----------------|
| `searchType` | categorical | Search type | `['SIMILARITY', 'SIMILARITY_THRESHOLD', 'MMR']` |
| `retrievalSource` | categorical | Retrieval source | `['MULTIMODAL', 'EMBEDDING']` |
| `queryStrategy` | categorical | Query strategy | `['RAW_QUERY', 'REWRITE_QUERY']` |
| `maxDocuments` | categorical | Max number of documents | `[2, 5, 10, 15, 20]` |
| `similarityThreshold` | categorical | Similarity threshold | `[0.2, 0.4, 0.6, 0.8]` |
| `mmrK` | categorical | K for MMR | `[5, 10, 15, 20, 25]` |
| `mmrDiversity` | categorical | MMR diversity | `[0.0, 0.2, 0.4, 0.6, 0.8, 1.0]` |

<br>

**Note**: `similarityThreshold` is only used if `searchType == 'SIMILARITY_THRESHOLD'`.  
`mmrK` and `mmrDiversity` are only used if `searchType == 'MMR'`.


<br>


## Usage and Configuration
---

### Using the Plugin in a Flow

The RAG Optimization plugin is used as a custom code recipe in your Dataiku flow. To use it correctly:

1. **Add the recipe** to your flow and connect the required inputs (Input Folder and Evaluation Dataset)
2. **Configure the parameters** in the recipe settings:
   - Select the evaluation dataset columns (User Query and Ground Truth)
   - Choose your embedding and completion models
   - Review or customize the optimization parameters (defaults are recommended)
   - Set execution parameters (number of trials, parallelism)
3. **Run the recipe** to start the optimization process
4. **Review the outputs** to analyze the results and identify optimal configurations

**Requirements for proper functioning**:
- The input folder must contain documents in supported formats (PDF, DOCX, PPTX, images, etc.)
- The evaluation dataset must have at least the User Query and Ground Truth columns
- Valid LLM connections must be configured for both embedding and completion models
- The work zone name must be specified (or use the default) to isolate temporary objects during optimization

See [Optimization Parameters](#optimization-parameters) section for parameter details.

<div class="alert">
Use a subset/sample to speed up optimization, then apply the optimal config to full data 
</div>

<br>

### Outputs
At the end of the optimization process, five datasets are generated:

**Optimization Results Datasets:**
1. **Embedding Optimization – Detailed Results**: Per-trial results with parameters, retrieved content, and scores
2. **RA-LLM Optimization – Detailed Results**: Per-trial results with retrieval/generation parameters and performance

**Evaluation Dataset:**
3. **Evaluation Detailed Results**: All question-answer pairs from all trials (embedding + RA-LLM) stacked together. Contains:
   - All original columns from the evaluation dataset
   - `answer_correctness`: Correctness score (0-1) for each question-answer pair
   - `estimated_cost`: Estimated cost for each question-answer pair
   - `trial_number`: Trial number identifier
   - `optimization_type`: Type of optimization ("embedding" or "rag_llm")

The plugin also automatically creates an optimized embedding recipe, its associated Knowledge Bank, and a RA-LLM with the best configurations.




<br>

## Main Features
---

### 1. Embedding Optimization

Optimizes document embedding recipe parameters, including:

- **Action to perform**: VLM or STRUCTURED
- **Chunk size**: Number of characters per chunk
- **Chunk overlap**: Overlap between chunks
- **VLM parameters**: Split unit, custom pages, prompts
- **STRUCTURED parameters**: Split unit, image handling, section depth

### 2. RA-LLM Optimization

Optimizes RA-LLM parameters for search and generation:

- **Search type**: SIMILARITY, SIMILARITY_THRESHOLD, MMR
- **Retrieval source**: MULTIMODAL, EMBEDDING
- **Query strategy**: RAW_QUERY, REWRITE_QUERY
- **LLM parameters**: LLM ID, temperature
- **Search parameters**: Number of documents, thresholds, MMR diversity

### 3. Automatic Evaluation

For each tested configuration:

1. Create recipe/LLM with the configuration
2. Build Knowledge Bank (for embeddings)
3. Execute queries on evaluation dataset
4. Calculate **correctness** via LLM judge
5. Calculate **estimated cost**
6. Calculate **composite score** (correctness × cost efficiency)
7. Automatic cleanup of temporary objects


<br>


## Best Practices
---

1. **Evaluation Dataset**: Use a representative but not too large dataset (10-50 examples often sufficient)

2. **Number of Trials**: Start with `n_trials=5-10` for testing, then increase gradually

3. **Parallelism**: Use `n_jobs=1` by default to avoid resource conflicts

4. **Parameters to Optimize**: Limit the number of parameters to reduce search space

5. **Discrete Values**: Use lists of discrete values rather than continuous ranges

6. **Cleanup**: The plugin automatically cleans up temporary objects, but verify in case of interruption


<br>


## Limitations
---

- Optimizations are sequential (embedding then RA-LLM)
- Automatic cleanup may fail on abrupt interruption
- Correctness is evaluated via LLM judge (may vary)


<br>

## Technical Requirements
---

Before using this plugin, make sure the following requirements are met.


### Dataiku Instance

This plugin has been tested and is compatible with **Dataiku 14.3+**.


### LLM Connections & Vector Storage

* **Completion LLM**: Used to generate answers from RAG queries and evaluate answer correctness via LLM judge.
  *Examples of tested technologies in this plugin:* **GPT-4.1**, **GPT-4.1 mini**, **OpenAI GPT models**.

* **Embedding LLM**: Used to generate embeddings for building the Knowledge Bases (KBs).
  *Examples of tested technologies in this plugin:* **Vertex AI text-embedding-004**, **OpenAI embedding models**.

* **Vector Storage**: Stores and manages the vectorized representations of documents for efficient retrieval.
  *Example of tested technology in this plugin:* **Dataiku Knowledge Banks** (supports various vector storage backends).



### Code Environment

The code environment should use **Python 3.11+** and include all the required libraries for the plugin to function correctly.

Required packages for this plugin are:

>scikit-learn
optuna
langchain
pydantic


<br>

## License
---

This plugin is licensed under the **Apache Software License** and **MIT License**.
