# RAG Optimization Plugin

<br>

## Overview

The **RAG Optimization Plugin** is a Dataiku plugin that automatically optimizes RAG system configurations using Bayesian optimization techniques by leveraging the [Optuna](https://optuna.org/) framework.
The plugin automatically searches for optimal RAG configurations by testing different parameter combinations in an intelligent way, using Optuna library to maximize answer quality while minimizing execution costs.

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

<div align="center">
  <img src="source/img/rag_llm_score_plot.png"/>
</div>


<br>

## Installation

**Install the plugin**: From the Plugins section in Dataiku, import or add the RAG Optimization plugin.

Once installed, the plugin components ("RAG Optimization Recipe") become available for use.



<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

#### What is [Optuna](https://optuna.org/)?

**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. Optimization parameters configuration
   ↓
2. Embedding Optimization
   ├── Create embedding recipes with different configs
   ├── Evaluate each configuration
   └── Select best configuration
   ↓
3. Build optimized Knowledge Bank
   ↓
4. RA-LLM Optimization
   ├── Create RA-LLMs with different configs
   ├── Evaluate each configuration
   └── Select best configuration
   ↓
5. Return results and metrics
```


#### 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]

#### 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

#### 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

#### 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>

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

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

<div  style="display: flex; justify-content: center; gap: 10px; flex-wrap: wrap;">
  <img src="source/img/plugin_recipe_flow.png"  height="400"/>
  <img src="source/img/input_folder.png" />
  <img src="source/img/eval_data.png"  />
</div>


<br>

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)

<div  style="display: flex; justify-content: center; gap: 10px; flex-wrap: wrap;">
  <img src="source/img/recipe_common_params.png" />
  <img src="source/img/embedding_opti_params.png" />
  <img src="source/img/rag_llm_opti_params.png"  />
</div>

<br>

3. **Run the recipe** to start the optimization process

During the optimization run, multiple objects (recipes, Knowledge Banks, RA-LLMs) may appear in the flow. This is normal and is due to the different simulations being launched for each trial. These temporary objects are automatically cleaned up at the end of the optimization process.

<div align="center">
  <img src="source/img/embedding_opti_run.png" height="600" />
</div>

4. **Review the outputs** to analyze the results and identify optimal configurations. At the end of the optimization process, five datasets are generated:
   
   **Optimization Results Datasets:**
   - **Embedding Optimization – Detailed Results**: Per-trial results with parameters, retrieved content, and scores
   - **RA-LLM Optimization – Detailed Results**: Per-trial results with retrieval/generation parameters and performance
   - **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.

<div align="center">
  <img src="source/img/optimized_embed_and_rag_llm_flow.png" />
</div>


<br>

### 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


<br>

## Best Practices

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

2. **Optimization Speed**: Use a subset/sample to speed up optimization, then apply the optimal config to full data

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

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

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

6. **Default Parameters**: Avoid modifying default parameters as there may be dependencies between parameters that could break the optimization process

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

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