pluto

RAG hyperparameter optimisation with Multi-Armed Bandits — a Python library that uses UCB bandits and agglomerative question clustering to find the best Retrieval-Augmented Generation configuration for each question cluster, without expensive grid search.

Table of Contents

• Overview
• Key Concepts
• Repository Structure
• Installation
• Quick Start
  • 1 — RAG pipeline
  • 2 — MAB configuration selector
  • 3 — Cluster-aware selection pipeline
• Module Reference
• Scripts
• Notebooks
• Development
• Version History
• Contributing

Overview

pluto solves a practical problem: which RAG configuration (chunk size, LLM temperature, …) gives the best answers for a given set of questions?

Instead of running an expensive brute-force grid search, it uses:

1. Multi-Armed Bandits (UCB) to efficiently explore the configuration space and converge on the best (chunk_size, llm_temperature) pair.
2. Agglomerative question clustering to group semantically similar questions, so each cluster gets its own optimal configuration rather than a single global one.
3. A KNN predictor that maps new, unseen questions to the nearest cluster — and therefore to the right configuration — at inference time with a single call.

Key Concepts

Repository Structure

pluto/
├── src/pluto/               # Installable Python package
│   ├── rag/                 # RAG pipeline
│   │   ├── config.py        #   RAGConfig, LLMProvider
│   │   ├── rag.py           #   ChromaRAG end-to-end pipeline
│   │   ├── rag_interface.py #   ChromaRAGInterface (low-level)
│   │   ├── llm_loader.py    #   LLM factory (Ollama / Groq)
│   │   └── utils.py         #   Prompt templates, helpers
│   ├── rl/                  # Action space and reward scorers
│   │   └── rl.py            #   RagAction, DatasetRewardScorer, HumanRewardScorer
│   ├── clustering/          # Question clustering
│   │   ├── clusters.py      #   AgglomerativeQuestionClusterer
│   │   ├── embeders.py      #   QuestionEmbedder
│   │   ├── cluster_predictor.py  # WeightedKNNClusterPredictor
│   │   └── viz.py           #   Cluster visualisation helpers
│   ├── eval/                # Evaluation
│   │   ├── evaluators.py    #   RAGBruteForceEvaluator, RAGMABEvaluator
│   │   ├── experiment.py    #   ExperimentRunner
│   │   ├── datasets.py      #   SquadV2Loader
│   │   └── viz.py           #   RewardHeatMap, plot_time_comparison
│   ├── selection/           # Configuration selection
│   │   ├── config_selector.py    # BaseSelector, MABSelector, SelectionResult
│   │   ├── config_predictor.py   # ClusterConfigPredictor, RoutedSelection
│   │   └── selection_pipeline.py # ClusterSelectionPipeline, PipelineResult
│   └── pipelines/           # End-to-end entry points
│       ├── training_pipeline.py  # pluto-train (offline MAB training)
│       └── online_pipeline.py    # pluto-ask  (inference from artifact)
├── scripts/                 # Runnable experiment scripts
├── notebooks/               # Jupyter notebooks
├── tests/                   # pytest test suite
├── artifacts/               # Saved pipeline artifacts (.pkl)
└── docs/                    # Documentation

Installation

Requirements: Python 3.10+, a running Ollama instance (or a Groq API key).

Editable install (recommended for development)

pip install -e ".[dev]"

Production install

pip install .

All runtime and development dependencies are declared in pyproject.toml. There is no requirements.txt.

Quick Start

1 — RAG pipeline

from pluto.rag.config import RAGConfig
from pluto.rag.rag import ChromaRAG
from pluto.rag.utils import SHORT_ANSWER_PROMPT

config = RAGConfig(
    texts=["Your long document goes here..."],
    llm_model_name="llama3.1:8b",
    llm_provider="ollama",
    embedding_model_name="sentence-transformers/all-mpnet-base-v2",
    custom_prompt_template=SHORT_ANSWER_PROMPT,
    chunk_size=350,
    llm_temperature=0.2,
)

rag = ChromaRAG(config)
answer = rag.query("What is the capital of France?")
print(answer)

2 — MAB configuration selector

Find the best RAG configuration for a set of questions using the UCB bandit:

from pluto.rag.config import RAGConfig
from pluto.rl.rl import DatasetRewardScorer, RagAction
from pluto.selection.config_selector import MABSelector

action_space = [
    RagAction(chunk_size=k, llm_temperature=t)
    for k in [100, 300, 500, 700, 1000]
    for t in [0.0, 0.3, 0.7, 1.0]
]

selector = MABSelector(
    config_template=config,          # RAGConfig from step 1
    reward_scorer=DatasetRewardScorer(similarity_threshold=0.5),
    action_space=action_space,
    alpha=1.0,
    seed=42,
)

result = selector.select(questions=items, trials=100)
print("Best config:", result.best_action)
print("Best score: ", result.best_score)

3 — Cluster-aware selection pipeline

Group questions by topic and find the optimal configuration per cluster:

from pluto.clustering.embeders import QuestionEmbedder
from pluto.clustering.clusters import AgglomerativeQuestionClusterer
from pluto.clustering.cluster_predictor import WeightedKNNClusterPredictor
from pluto.selection.selection_pipeline import ClusterSelectionPipeline

embedder  = QuestionEmbedder("sentence-transformers/all-mpnet-base-v2")
clusterer = AgglomerativeQuestionClusterer(embedder=embedder)

pipeline = ClusterSelectionPipeline(clusterer=clusterer, selector=selector)
pipeline_result = pipeline.run(
    items=items,            # list of {"question": ..., "answers": ...} dicts
    min_clusters=2,
    max_clusters=5,
    trials=100,
)

# Predict optimal config for a new question
predictor = WeightedKNNClusterPredictor(
    embedder=embedder,
    train_embeddings=pipeline_result.cluster_result.embeddings,
    train_labels=pipeline_result.cluster_result.labels,
    k=3,
)

cluster_id = predictor.predict_cluster("What year was Python created?")
optimal = pipeline_result.cluster_selections[cluster_id].selection
print("Optimal action:", optimal.best_action)

Module Reference

Scripts

Pre-built runnable scripts are in scripts/. Run them from the repository root after installing the package.

# Example
python scripts/single_question_analysis.py

Notebooks

Interactive tutorials are in notebooks/:

Development

# Install all dev dependencies
pip install -e ".[dev]"

# Lint & format
ruff check .
ruff format .

# Run tests
pytest tests/

Version History

[0.1.0] — 2026-06-14 · Initial Release

Core infrastructure

• RAGConfig dataclass — unified control surface for all RAG hyperparameters (LLM provider, model, chunk size, temperature, embedding model, prompt template)
• ChromaRAG end-to-end pipeline and ChromaRAGInterface low-level retrieval layer backed by ChromaDB
• LLM factory supporting Ollama and Groq providers

Reward scoring

• RagAction action space (chunk size × LLM temperature)
• UCB Multi-Armed Bandit selector (MABSelector) with configurable exploration coefficient
• DatasetRewardScorer (cosine similarity) and HumanRewardScorer reward functions

Question clustering

• QuestionEmbedder wrapping sentence-transformers
• AgglomerativeQuestionClusterer with silhouette-based automatic cluster count selection
• WeightedKNNClusterPredictor for routing unseen questions to the nearest cluster

Selection pipeline

• ClusterSelectionPipeline — clusters questions and runs a per-cluster MAB selector
• ClusterConfigPredictor — single-call inference: question in, optimal RagAction out

Evaluation & tooling

• Brute-force and MAB evaluators with reward heatmaps and step-wise snapshots
• ExperimentRunner for multi-question, multi-method comparisons
• SquadV2Loader dataset adapter
• Runnable scripts and Jupyter notebooks
• Full pytest suite and GitHub Actions CI

Contributing

Want to contribute to pluto?
Open an issue or a pull request — all contributions are welcome.

Licensed under Apache 2.0.