RouterBench

https://github.com/withmartian/routerbench
Admin
RouterBench is a benchmark framework for evaluating and comparing multi-LLM routing systems,...
Tokens:9,500
Snippets:68
Trust Score:9.5
License:MIT
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install Project Packages

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Install the project's Python packages by running this command in the root directory.

```bash
pip install -e .
```

--------------------------------

### Router Evaluation Configuration Example

Source: https://context7.com/withmartian/routerbench/llms.txt

An example YAML configuration file for `evaluate_routers.py`, demonstrating how to set parameters for data loading, model selection, and router-specific settings. This provides a structured way to manage complex evaluation setups.

```yaml
data_path: routerbench_5shot.pkl
train_fraction: 0.7
local: true
local_cache: true
data_name: routerbench
embedding_models: ["all-MiniLM-L12-v2"]
wanted_eval_name: ["grade-school-math", "hellaswag", "mmlu", "arc-challenge"]
knn:
  neighbors: [5, 10, 40]
  metrics: ["cosine"]
mlp:
  hidden_layer_sizes: [100, [100, 100]]
  activation: ["relu"]
  learning_rates: [0.001]
cascading_router:
  max_cost_per_response_list: [0.001, 0.01, 0.1, 1.0, 10.0]
  error_rates: [0, 0.05, 0.1, 0.2]
```

--------------------------------

### Install RouterBench and Configure Environment

Source: https://context7.com/withmartian/routerbench/llms.txt

Install the package in editable mode and create a .env file for MongoDB connection string if needed.

```bash
pip install -e .

# Create a .env file with your MongoDB connection string (optional – omit for local-only mode)
echo "CONNECTION_STRING='mongodb+srv://user:pass@cluster.mongodb.net/'" > .env
```

--------------------------------

### Environment Variable Setup

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Create a .env file in the root directory to configure your MongoDB connection string for embedding caching.

```bash
CONNECTION_STRING='your mongodb connection string'
```

--------------------------------

### CLI entry point for router evaluation

Source: https://context7.com/withmartian/routerbench/llms.txt

Provides a command-line interface for evaluating routers, typically using a YAML configuration file for setup. Executes the `evaluate_routers.py` script.

```bash
# Local run using YAML config
python evaluate_routers.py --config=configs/evaluate_routers.yaml
```

--------------------------------

### Code Formatting and Linting

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Run these commands to ensure code quality by applying flake8, black, and isort to Python files tracked by Git. Install them if necessary.

```bash
flake8 $(git ls-files '*.py')
black $(git ls-files '*.py')
isort $(git ls-files '*.py')
```

--------------------------------

### Deploy Modal App

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Deploy the updated modal application using this command.

```bash
modal deploy modal_router.py
```

--------------------------------

### Visualize Results CLI Entry Point

Source: https://context7.com/withmartian/routerbench/llms.txt

Command-line interface for visualizing evaluation results using `visualize_results.py`. This script can be run with a configuration file or inline arguments to specify input files and desired evaluations.

```bash
python visualize_results.py --config=configs/visualize.yaml

# Inline
python visualize_results.py \
  --input-file data/routerbench/eval_results/eval_results__01-15-14__routerbench.pkl \
  --data-path data/routerbench/input_wide.pkl \
  --data-name routerbench \
  --wanted-eval-names '["grade-school-math", "mmlu"]'
```

--------------------------------

### Deploy Modal App for Serverless Evaluation

Source: https://context7.com/withmartian/routerbench/llms.txt

Deploys the RouterBench Modal app for large-scale parallel experiments. After deployment, `evaluate_routers.py` automatically uses Modal unless `--local` is specified.

```bash
# Deploy the Modal app
modal deploy modal_router.py

# After deployment, evaluate_routers.py automatically uses Modal when --local is NOT set
python evaluate_routers.py \
  --data-path data/routerbench/input_wide.pkl \
  --data-name routerbench \
  --gcp-token gcs_credentials.json \
  --gcp-bucket my-gcs-bucket \
  --embedding-models '["all-MiniLM-L12-v2"]' \
  --knn.neighbors '[5, 10]' \
  --knn.metrics '["cosine"]'
```

--------------------------------

### Evaluate Routers with Command-Line Arguments

Source: https://context7.com/withmartian/routerbench/llms.txt

Execute router evaluations using the `evaluate_routers.py` script with various command-line arguments to specify data, models, and evaluation parameters. This is useful for running batch evaluations with custom configurations.

```bash
python evaluate_routers.py \
  --data-path data/routerbench/input_wide.pkl \
  --data-name routerbench \
  --embedding-models '["all-MiniLM-L12-v2"]' \
  --train-fraction 0.7 \
  --local \
  --local-cache \
  --wanted-eval-name '["grade-school-math", "mmlu"]' \
  --knn.neighbors '[5, 10, 40]' \
  --knn.metrics '["cosine"]' \
  --mlp.hidden_layer_sizes '[100, [100,100]]' \
  --mlp.activation '["relu"]' \
  --mlp.learning_rates '[0.001]' \
  --cascading_router.max_cost_per_response_list '[0.001, 0.01, 0.1, 1.0]' \
  --cascading_router.error_rates '[0, 0.05, 0.2]'
```

--------------------------------

### Convert Data Script

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Run this script to process various data formats into a common format for evaluation. Ensure MongoDB is running or set local_cache to true if not.

```bash
convert_data.py --config=configs/convert_data.yaml
```

--------------------------------

### CLI for data conversion with convert_data.py

Source: https://context7.com/withmartian/routerbench/llms.txt

Convert raw data logs using the convert_data.py script. Supports configuration via YAML files or inline arguments for input format, data path, embedding models, and cache settings.

```bash
# Using a YAML config file (recommended)
python convert_data.py --config=configs/convert_data.yaml

# Equivalent inline arguments
python convert_data.py \
  --data-path routerbench_raw.pkl \
  --input-format openai \
  --data-name routerbench \
  --embedding-models '["all-MiniLM-L12-v2"]' \
  --local-cache
```

```yaml
data_path: routerbench_raw.pkl
input_format: openai          # or "martian-evals"
data_name: routerbench
embedding_models: ["all-MiniLM-L12-v2"]
local_cache: true
```

--------------------------------

### Implement a KNN Router

Source: https://context7.com/withmartian/routerbench/llms.txt

A K-Nearest Neighbors router that fits a NearestNeighbors model on prompt embeddings. It retrieves similar training samples to select the best model, with an option to balance performance and cost using willingness_to_pay.

```python
import pandas as pd
from embedding.cache import EmbeddingCache
from routers.knn_router import KNNRouter

cache = EmbeddingCache(local_mode=True)
train_df = pd.read_pickle("data/routerbench/input_wide.pkl")
models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview", "meta/llama-2-70b-chat"]

router = KNNRouter(
    train_file=train_df,
    embedding_model="all-MiniLM-L6-v2",
    cache=cache,
    n_neighbors=10,
    distance_metric="cosine",
    models_to_route=models,
)

prompts = ["Solve: 3x + 5 = 20", "Summarize the French Revolution in 3 sentences."]

# Route without cost penalty (pure performance)
best_models = router.batch_route_prompts(prompts, willingness_to_pay=0)
print(best_models)  # e.g. ['gpt-4-1106-preview' 'gpt-4-1106-preview']

# Route with cost-awareness (lower willingness_to_pay → cheaper models preferred)
cost_aware = router.batch_route_prompts(prompts, willingness_to_pay=0.001)
print(cost_aware)  # e.g. ['gpt-3.5-turbo-1106' 'meta/llama-2-70b-chat']
```

--------------------------------

### Load, Inspect, and Export Evaluation Results

Source: https://context7.com/withmartian/routerbench/llms.txt

Demonstrates loading a saved `EvaluationResultCollection`, inspecting available evaluations and router types, and exporting all results to a CSV file. This is essential for post-evaluation analysis and data management.

```python
from evaluation.eval import DefaultEvaluationResult, EvaluationResultCollection
import pandas as pd

# Load a previously saved collection
collection = EvaluationResultCollection.load(
    "data/routerbench/eval_results/eval_results__01-15-14__routerbench.pkl"
)

# Inspect available eval names and router types
print(collection.get_eval_names())
# ['grade-school-math', 'mmlu', 'hellaswag', ...]
print(collection.get_router_names())
# {'knn': [...], 'mlp': [...], 'cascading router': [...]} 

# Export all results to CSV
collection.to_csv("data/routerbench/eval_results/combined_results.csv")

# Save/load a single result
result: DefaultEvaluationResult = collection.evaluation_results[0]
result.save("data/routerbench/eval_results/knn_result.pkl")
loaded = DefaultEvaluationResult.load("data/routerbench/eval_results/knn_result.pkl")
```

--------------------------------

### Estimate token cost for prompt-response pairs

Source: https://context7.com/withmartian/routerbench/llms.txt

Calculate the USD cost of a prompt-response pair using per-token pricing. Supports both single-turn and multi-turn conversations.

```python
from routers.abstract_router import calculate_cost_for_prompt_and_response

cost = calculate_cost_for_prompt_and_response(
    prompt="Explain the difference between supervised and unsupervised learning.",
    response="Supervised learning uses labeled data...",
    model_name="gpt-4-1106-preview",
)
print(f"Estimated cost: ${cost:.8f}")
# e.g. Estimated cost: $0.00004200

# Multi-turn (list of strings)
multi_turn_cost = calculate_cost_for_prompt_and_response(
    prompt=["Turn 1: Hello", "Turn 2: What can you do?"],
    response=["Hi there!", "I can help with many tasks."],
    model_name="claude-v2",
)
print(f"Multi-turn cost: ${multi_turn_cost:.8f}")
```

--------------------------------

### calculate_cost_for_prompt_and_response

Source: https://context7.com/withmartian/routerbench/llms.txt

Computes the USD cost of a single prompt-response pair for a given model using per-token pricing.

```APIDOC
## calculate_cost_for_prompt_and_response

### Description
Computes the USD cost of a single prompt–response pair for a given model using per-token pricing from `routers/common.py`.

### Usage
```python
from routers.abstract_router import calculate_cost_for_prompt_and_response

cost = calculate_cost_for_prompt_and_response(
    prompt="Explain the difference between supervised and unsupervised learning.",
    response="Supervised learning uses labeled data...",
    model_name="gpt-4-1106-preview",
)
print(f"Estimated cost: ${cost:.8f}")
# e.g. Estimated cost: $0.00004200

# Multi-turn (list of strings)
multi_turn_cost = calculate_cost_for_prompt_and_response(
    prompt=["Turn 1: Hello", "Turn 2: What can you do?"],
    response=["Hi there!", "I can help with many tasks."],
    model_name="claude-v2",
)
print(f"Multi-turn cost: ${multi_turn_cost:.8f}")
```
```

--------------------------------

### Implement a Random Router

Source: https://context7.com/withmartian/routerbench/llms.txt

A basic router implementation that randomly selects a model from a provided list for each prompt. Inherits from AbstractRouter and requires a list of models to route.

```python
import numpy as np
import pandas as pd
from numpy.typing import NDArray
from routers.abstract_router import AbstractRouter

class RandomRouter(AbstractRouter):
    def __init__(self, models_to_route: list[str], **kwargs) -> None:
        self.models_to_route = models_to_route

    def batch_route_prompts(self, prompts: list[str], **kwargs) -> NDArray[str]:
        return np.random.choice(self.models_to_route, size=len(prompts))

models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview", "meta/llama-2-70b-chat"]
router = RandomRouter(models_to_route=models)
decisions = router.batch_route_prompts(["What is 2+2?", "Write a sorting algorithm."])
print(decisions)  # e.g. ['gpt-3.5-turbo-1106' 'gpt-4-1106-preview']
```

--------------------------------

### Custom Data Conversion with AbstractConvertor

Source: https://context7.com/withmartian/routerbench/llms.txt

Implement a custom data convertor by inheriting from AbstractConvertor and implementing the convert method. The convert_and_embed method handles conversion and embedding.

```python
import pandas as pd
from embedding.cache import EmbeddingCache
from convertors.abstract_convertor import AbstractConvertor

class MyCustomConvertor(AbstractConvertor):
    def convert(self, input_data: pd.DataFrame) -> pd.DataFrame:
        # Normalize your format into the standard schema:
        # Required columns: sample_id, prompt, eval_name, model_name, performance, model_response
        input_data = input_data.rename(columns={"question": "prompt", "score": "performance"})
        input_data["eval_name"] = "my_eval"
        return input_data

cache = EmbeddingCache(local_mode=True)
convertor = MyCustomConvertor()

raw_df = pd.DataFrame({
    "sample_id": ["q1", "q2"],
    "question": ["What is 2+2?", "Name a planet."],
    "model_name": ["gpt-3.5-turbo-1106", "gpt-3.5-turbo-1106"],
    "score": [1.0, 1.0],
    "model_response": ["4", "Mars"],
})

converted_df = convertor.convert_and_embed(
    raw_df, cache, embedding_models=["all-MiniLM-L6-v2"]
)
print(converted_df.columns.tolist())
```

--------------------------------

### Visualize Results Script

Source: https://github.com/withmartian/routerbench/blob/main/README.md

This script visualizes evaluation results using the EvaluationCollection, creating a performance-vs-cost plot.

```bash
visualize_results.py --config=configs/visualize.yaml
```

--------------------------------

### SVMRouter for prompt routing

Source: https://context7.com/withmartian/routerbench/llms.txt

Trains a LinearSVR for each model to regress performance from embeddings. Selects the highest-scoring model per prompt, with optional cost adjustment based on willingness to pay. Requires `pandas`, `embedding.cache.EmbeddingCache`, and `routers.svm_router.SVMRouter`.

```python
import pandas as pd
from embedding.cache import EmbeddingCache
from routers.svm_router import SVMRouter

cache = EmbeddingCache(local_mode=True)
train_df = pd.read_pickle("data/routerbench/input_wide.pkl")
models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview", "mistralai/mistral-7b-chat"]

router = SVMRouter(
    train_file=train_df,
    embedding_model="all-MiniLM-L6-v2",
    cache=cache,
    models_to_route=models,
)

decisions = router.batch_route_prompts(
    ["Define recursion.", "Write a haiku about autumn."],
    willingness_to_pay=1.0,
)
print(decisions)
```

--------------------------------

### Evaluate Routers Script

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Use this script with processed data to evaluate different routers. It generates a CSV results file and an EvaluationCollection.

```bash
evaluate_routers.py --config=configs/evaluate_routers.yaml
```

--------------------------------

### Run cascading router for evaluation

Source: https://context7.com/withmartian/routerbench/llms.txt

Simulates a cascading strategy where models are tried in ascending cost order. Responses are accepted if they meet a performance threshold (estimated by an evaluator with a configurable error rate) or rejected, and the next model is tried until the budget is exhausted. Returns aggregate performance and total cost statistics. Requires `pandas` and `routers.run_cascading_router.run_cascading_router_for_eval`.

```python
import pandas as pd
from routers.run_cascading_router import run_cascading_router_for_eval

eval_df = pd.read_pickle("data/routerbench/input_wide.pkl")
eval_df = eval_df[eval_df["eval_name"] == "grade-school-math"]
models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview"]

result = run_cascading_router_for_eval(
    inputs_df=eval_df,
    performance_threshold_quantile=0.95,   # accept if in top 5% of performances
    max_cost_per_response=0.01,            # budget cap per query in USD
    evaluator_error_rate=0.05,             # 5% chance evaluator is wrong
    eval_name="grade-school-math",
    models_to_route=models,
)
print(result)
# {
#   'model_name': 'cascading router',
#   'eval_name': 'grade-school-math',
#   'performance': 0.812,
#   'total_cost': 0.427,
#   'max_cost_per_response': 0.01,
#   'evaluator_error_rate': 0.05,
#   'performance_threshold': 0.91
# }
```

--------------------------------

### Parse Router Name to Parameters

Source: https://context7.com/withmartian/routerbench/llms.txt

Decodes a pipe-separated router name string into a structured dictionary of hyperparameters. Ensure the `utils` module is available.

```python
from utils import parse_model_parameters

model_string = (
    "knn|neighbors:500.0|metric:cosine|fraction:1.0"
    "|embedding:all-MiniLM-L6-v2|willingness_to_pay:2.0|total_cost"
)
params = parse_model_parameters(model_string)
print(params)
```

--------------------------------

### RandomRouter

Source: https://context7.com/withmartian/routerbench/llms.txt

A simple routing algorithm that randomly selects a model from a predefined list for each prompt.

```APIDOC
## RandomRouter

### Description
Base class for routing algorithms. All routers must implement `__init__` (fit/load weights) and `batch_route_prompts(prompts, **kwargs) -> NDArray[str]`. The return value is an array of model name strings (one per prompt) drawn from the pool of routable models.

### Usage
```python
import numpy as np
import pandas as pd
from numpy.typing import NDArray
from routers.abstract_router import AbstractRouter

class RandomRouter(AbstractRouter):
    def __init__(self, models_to_route: list[str], **kwargs) -> None:
        self.models_to_route = models_to_route

    def batch_route_prompts(self, prompts: list[str], **kwargs) -> NDArray[str]:
        return np.random.choice(self.models_to_route, size=len(prompts))

models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview", "meta/llama-2-70b-chat"]
router = RandomRouter(models_to_route=models)
decisions = router.batch_route_prompts(["What is 2+2?", "Write a sorting algorithm."])
print(decisions)  # e.g. ['gpt-3.5-turbo-1106' 'gpt-4-1106-preview']
```
```

--------------------------------

### Build Reproducible Train/Eval Dataset

Source: https://context7.com/withmartian/routerbench/llms.txt

Splits a dataset by `sample_id` or `source` using a fixed random seed, ensuring no overlap between train and validation sets. Supports out-of-distribution evaluation.

```python
import pandas as pd
from utils import build_train_eval_dataset

dataset_df = pd.read_pickle("data/routerbench/input_wide.pkl")

train_df, eval_df = build_train_eval_dataset(
    wanted_eval_name="grade-school-math",
    other_eval_names=dataset_df["eval_name"].unique().tolist(),
    dataset_df=dataset_df,
    fraction=0.7,
    out_of_distribution=False,
)

print(f"Train samples: {len(train_df)}")   # ~70% of grade-school-math samples
print(f"Eval samples:  {len(eval_df)}")    # ~30% of grade-school-math samples
# Verify no overlap
assert len(set(train_df.sample_id) & set(eval_df.sample_id)) == 0
```

--------------------------------

### EmbeddingCache for Prompt Embedding with Caching

Source: https://context7.com/withmartian/routerbench/llms.txt

Computes and caches sentence-transformer embeddings for prompts. Supports local-only mode or MongoDB caching for large-scale runs.

```python
from embedding.cache import EmbeddingCache

# Local-only mode: no MongoDB required
cache = EmbeddingCache(local_mode=True)

prompts = (
    "What is the capital of France?",
    "Explain quantum entanglement in simple terms.",
    "Write a Python function to reverse a linked list.",
)

# Returns a numpy array of shape (len(prompts), embedding_dim)
embeddings = cache.batch_get_embedding(prompts, embedding_model="all-MiniLM-L6-v2")
print(embeddings.shape)  # e.g. (3, 384)

# With MongoDB caching (for large-scale runs)
cache_mongo = EmbeddingCache(
    connection_string="mongodb+srv://user:pass@cluster.mongodb.net/",
    local_mode=False,
)
embeddings_cached = cache_mongo.batch_get_embedding(prompts, embedding_model="all-MiniLM-L12-v2")
```

--------------------------------

### Run batch cascading router sweep

Source: https://context7.com/withmartian/routerbench/llms.txt

Sweeps over combinations of `max_cost_per_response_list` and `error_rates`, runs `run_cascading_router_for_eval` for each, and saves results to a CSV file. Returns the path to the saved file. Requires `pandas` and `routers.run_cascading_router.run_cascading_router_for_eval_dataframe`.

```python
import pandas as pd
from routers.run_cascading_router import run_cascading_router_for_eval_dataframe

eval_df = pd.read_pickle("data/routerbench/input_wide.pkl")
models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview"]

csv_path = run_cascading_router_for_eval_dataframe(
    eval_df=eval_df,
    eval_name="grade-school-math",
    max_cost_per_response_list=[0.001, 0.01, 0.1, 1.0],
    error_rates=[0.0, 0.05, 0.2],
    models_to_route=models,
)
print(f"Results saved to: {csv_path}")
results = pd.read_csv(csv_path)
print(results.head())
```

--------------------------------

### MartianEvalConvertor

Source: https://context7.com/withmartian/routerbench/llms.txt

Converts raw Martian Evals JSONL logs into a wide-format DataFrame. It parses events per sample, normalizes prompts, and embeds data.

```APIDOC
## MartianEvalConvertor

### Description
Walks a directory tree of `.jsonl` log files produced by `martian-evals`, parses sampling/match/metrics events per sample, normalizes multi-turn prompts, and outputs the same wide-format DataFrame as `OpenAIEvalsConvertor`.

### Usage
```python
from convertors.martian_evals_convertor import MartianEvalConvertor
from embedding.cache import EmbeddingCache

cache = EmbeddingCache(local_mode=True)
convertor = MartianEvalConvertor()

# data_path is a directory containing .jsonl log files
wide_df = convertor.convert_and_embed(
    "path/to/martian_evals_output/",
    cache,
    embedding_models=["all-MiniLM-L6-v2"],
)
print(wide_df.head())
```
```

--------------------------------

### Convert Martian Evals JSONL logs to wide DataFrame

Source: https://context7.com/withmartian/routerbench/llms.txt

Use MartianEvalConvertor to parse and normalize multi-turn prompts from Martian Evals JSONL logs. Requires an EmbeddingCache for embedding prompts.

```python
from convertors.martian_evals_convertor import MartianEvalConvertor
from embedding.cache import EmbeddingCache

cache = EmbeddingCache(local_mode=True)
convertor = MartianEvalConvertor()

# data_path is a directory containing .jsonl log files
wide_df = convertor.convert_and_embed(
    "path/to/martian_evals_output/",
    cache,
    embedding_models=["all-MiniLM-L6-v2"],
)
print(wide_df.head())
```

--------------------------------

### Convert OpenAI Evals Data with OpenAIEvalsConvertor

Source: https://context7.com/withmartian/routerbench/llms.txt

Converts long-format DataFrame from OpenAI Evals to RouterBench wide format, adds costs, and computes oracle routing labels.

```python
import pandas as pd
from convertors.openai_evals_convertor import OpenAIEvalsConvertor
from embedding.cache import EmbeddingCache

# Input: long-format DataFrame produced by OpenAI Evals
raw_df = pd.read_pickle("routerbench_raw.pkl")

cache = EmbeddingCache(local_mode=True)
convertor = OpenAIEvalsConvertor()

wide_df = convertor.convert_and_embed(
    raw_df, cache, embedding_models=["all-MiniLM-L12-v2"]
)
# wide_df now has columns like:
#   sample_id, prompt, eval_name,

```

--------------------------------

### BibTeX Citation

Source: https://github.com/withmartian/routerbench/blob/main/README.md

Use this BibTeX entry for citing the RouterBench paper in academic work.

```bibtex
@article{hu2024routerbench,
  title   = {ROUTERBENCH: A Benchmark for Multi-LLM Routing System},
  author  = {Qitian Jason Hu and Jacob Bieker and Xiuyu Li and Nan Jiang and Benjamin Keigwin and Gaurav Ranganath and Kurt Keutzer and Shriyash Kaustubh Upadhyay},
  year    = {2024},
  journal = {arXiv preprint arXiv: 2403.12031}
}
```

--------------------------------

### Visualize Performance vs. Cost Curves

Source: https://context7.com/withmartian/routerbench/llms.txt

Generates performance-cost Pareto curves for routers using `EvaluationResultCollection.plot_performance_vs_cost`. This function helps visualize the trade-offs between performance and cost, with options to show the Non-Decreasing Convex Hull (NDCH).

```python
from evaluation.eval import EvaluationResultCollection
import pandas as pd

collection = EvaluationResultCollection.load(
    "data/routerbench/eval_results/eval_results__01-15-14__routerbench.pkl"
)
dataset_df = pd.read_pickle("data/routerbench/input_wide.pkl")
models_to_route = dataset_df.model_name.unique().tolist()

# Plot for a single eval, save PNG, show NDCH
collection.plot_performance_vs_cost(
    eval_name="grade-school-math",
    save_file_path="data/routerbench/analysis_results/",
    show_plot=False,
    plot_ndch=True,
    models_to_route=models_to_route,
)

# Plot aggregate across all evals (eval_name=None)
collection.plot_performance_vs_cost(
    eval_name=None,
    show_plot=True,
    plot_ndch=False,
    models_to_route=models_to_route,
)
```

--------------------------------

### KNNRouter

Source: https://context7.com/withmartian/routerbench/llms.txt

A K-Nearest Neighbors router that uses prompt embeddings to select the best model based on historical performance and cost.

```APIDOC
## KNNRouter

### Description
K-Nearest Neighbors router. Fits a `sklearn.neighbors.NearestNeighbors` model on training prompt embeddings. At inference time, it retrieves the k most similar training samples and averages their per-model performance scores to select the best model. Supports a `willingness_to_pay` parameter that trades off performance against cost.

### Usage
```python
import pandas as pd
from embedding.cache import EmbeddingCache
from routers.knn_router import KNNRouter

cache = EmbeddingCache(local_mode=True)
train_df = pd.read_pickle("data/routerbench/input_wide.pkl")
models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview", "meta/llama-2-70b-chat"]

router = KNNRouter(
    train_file=train_df,
    embedding_model="all-MiniLM-L6-v2",
    cache=cache,
    n_neighbors=10,
    distance_metric="cosine",
    models_to_route=models,
)

prompts = ["Solve: 3x + 5 = 20", "Summarize the French Revolution in 3 sentences."]

# Route without cost penalty (pure performance)
best_models = router.batch_route_prompts(prompts, willingness_to_pay=0)
print(best_models)  # e.g. ['gpt-4-1106-preview' 'gpt-4-1106-preview']

# Route with cost-awareness (lower willingness_to_pay → cheaper models preferred)
cost_aware = router.batch_route_prompts(prompts, willingness_to_pay=0.001)
print(cost_aware)  # e.g. ['gpt-3.5-turbo-1106' 'meta/llama-2-70b-chat']
```
```

--------------------------------

### AbstractConvertor

Source: https://context7.com/withmartian/routerbench/llms.txt

Base class for data format converters. Users can subclass this to create custom converters.

```APIDOC
## AbstractConvertor

### Description
All convertors inherit from `AbstractConvertor` and must implement `convert(input_data) -> pd.DataFrame`. The convenience method `convert_and_embed` runs conversion and then triggers embedding computation in one call.

### Method
`AbstractConvertor()`

### Method
`convert(input_data: pd.DataFrame) -> pd.DataFrame`

### Parameters
#### Method Parameters
- **input_data** (pd.DataFrame) - Required - The raw input data to convert.

### Response
#### Success Response
- **pd.DataFrame** - A DataFrame in the standard RouterBench format with required columns: `sample_id`, `prompt`, `eval_name`, `model_name`, `performance`, `model_response`.

### Method
`convert_and_embed(raw_df: pd.DataFrame, cache: EmbeddingCache, embedding_models: list) -> pd.DataFrame`

### Parameters
#### Method Parameters
- **raw_df** (pd.DataFrame) - Required - The raw input data to convert.
- **cache** (EmbeddingCache) - Required - An instance of `EmbeddingCache` for prompt embedding.
- **embedding_models** (list) - Required - A list of embedding model names to use.

### Response
#### Success Response
- **pd.DataFrame** - The converted DataFrame with embeddings computed.

### Request Example
```python
import pandas as pd
from embedding.cache import EmbeddingCache
from convertors.abstract_convertor import AbstractConvertor

class MyCustomConvertor(AbstractConvertor):
    def convert(self, input_data: pd.DataFrame) -> pd.DataFrame:
        # Normalize your format into the standard schema:
        # Required columns: sample_id, prompt, eval_name, model_name, performance, model_response
        input_data = input_data.rename(columns={"question": "prompt", "score": "performance"})
        input_data["eval_name"] = "my_eval"
        return input_data

cache = EmbeddingCache(local_mode=True)
convertor = MyCustomConvertor()

raw_df = pd.DataFrame({
    "sample_id": ["q1", "q2"],
    "question": ["What is 2+2?", "Name a planet."],
    "model_name": ["gpt-3.5-turbo-1106", "gpt-3.5-turbo-1106"],
    "score": [1.0, 1.0],
    "model_response": ["4", "Mars"],
})

converted_df = convertor.convert_and_embed(
    raw_df, cache, embedding_models=["all-MiniLM-L6-v2"]
)
print(converted_df.columns.tolist())
```
```

--------------------------------

### MLPRouter for prompt routing

Source: https://context7.com/withmartian/routerbench/llms.txt

Trains an MLP regressor for each model to predict performance scores from prompt embeddings. Selects the model with the highest adjusted score, optionally penalized by cost. Requires `pandas`, `embedding.cache.EmbeddingCache`, and `routers.mlp_router.MLPRouter`.

```python
import pandas as pd
from embedding.cache import EmbeddingCache
from routers.mlp_router import MLPRouter

cache = EmbeddingCache(local_mode=True)
train_df = pd.read_pickle("data/routerbench/input_wide.pkl")
models = ["gpt-3.5-turbo-1106", "gpt-4-1106-preview", "mistralai/mixtral-8x7b-chat"]

router = MLPRouter(
    train_file=train_df,
    embedding_model="all-MiniLM-L12-v2",
    cache=cache,
    hidden_layer_sizes=(100, 100),
    activation_function="relu",
    learning_rate=0.001,
    models_to_route=models,
)

prompts = ["What causes thunder?", "Implement a binary search tree in Python."]
decisions = router.batch_route_prompts(prompts, willingness_to_pay=0.5)
print(decisions)
```

--------------------------------

### OpenAIEvalsConvertor

Source: https://context7.com/withmartian/routerbench/llms.txt

Converts data formatted by OpenAI Evals into the wide RouterBench format and computes embeddings.

```APIDOC
## OpenAIEvalsConvertor

### Description
Converts a long-format DataFrame (one row per model per sample) with standard OpenAI Evals columns into the wide RouterBench format, adds per-model costs, and computes oracle routing labels.

### Method
`OpenAIEvalsConvertor()`

### Method
`convert_and_embed(raw_df: pd.DataFrame, cache: EmbeddingCache, embedding_models: list) -> pd.DataFrame`

### Parameters
#### Method Parameters
- **raw_df** (pd.DataFrame) - Required - Input: long-format DataFrame produced by OpenAI Evals.
- **cache** (EmbeddingCache) - Required - An instance of `EmbeddingCache` for prompt embedding.
- **embedding_models** (list) - Required - A list of embedding model names to use.

### Response
#### Success Response
- **pd.DataFrame** - A DataFrame in the wide RouterBench format with columns like: `sample_id`, `prompt`, `eval_name`, `model_name`, `performance`.

### Request Example
```python
import pandas as pd
from convertors.openai_evals_convertor import OpenAIEvalsConvertor
from embedding.cache import EmbeddingCache

# Input: long-format DataFrame produced by OpenAI Evals
raw_df = pd.read_pickle("routerbench_raw.pkl")

cache = EmbeddingCache(local_mode=True)
convertor = OpenAIEvalsConvertor()

wide_df = convertor.convert_and_embed(
    raw_df, cache, embedding_models=["all-MiniLM-L12-v2"]
)
# wide_df now has columns like:
#   sample_id, prompt, eval_name,
```
```

--------------------------------

### EmbeddingCache

Source: https://context7.com/withmartian/routerbench/llms.txt

Computes and caches sentence-transformer embeddings for prompts. Supports in-memory, local file, and MongoDB caching.

```APIDOC
## EmbeddingCache

### Description
`EmbeddingCache` computes and caches sentence-transformer embeddings for batches of prompts. It checks an in-memory LRU cache, then a local pickle file, then MongoDB (if not in local mode) before calling the embedding model. The cache is always written back to `data/embedding_cache_<model>.pkl` for future runs.

### Method
`EmbeddingCache(local_mode: bool = True, connection_string: str = None)`

### Parameters
#### Constructor Parameters
- **local_mode** (bool) - Optional - If True, disables MongoDB caching.
- **connection_string** (str) - Optional - MongoDB connection string. Required if `local_mode` is False.

### Method
`batch_get_embedding(prompts: tuple, embedding_model: str) -> numpy.ndarray`

### Parameters
#### Method Parameters
- **prompts** (tuple) - Required - A tuple of strings, where each string is a prompt.
- **embedding_model** (str) - Required - The name of the sentence-transformer model to use for embedding.

### Response
#### Success Response
- **numpy.ndarray** - A numpy array of shape (len(prompts), embedding_dim) containing the embeddings.

### Request Example
```python
from embedding.cache import EmbeddingCache

# Local-only mode
cache = EmbeddingCache(local_mode=True)
prompts = (
    "What is the capital of France?",
    "Explain quantum entanglement in simple terms.",
    "Write a Python function to reverse a linked list.",
)
embeddings = cache.batch_get_embedding(prompts, embedding_model="all-MiniLM-L6-v2")
print(embeddings.shape)

# With MongoDB caching
cache_mongo = EmbeddingCache(
    connection_string="mongodb+srv://user:pass@cluster.mongodb.net/",
    local_mode=False,
)
embeddings_cached = cache_mongo.batch_get_embedding(prompts, embedding_model="all-MiniLM-L12-v2")
```
```

--------------------------------

### Calculate Area-under-NDCH Quality (AIQ) Score

Source: https://context7.com/withmartian/routerbench/llms.txt

Computes the AIQ score for router results using the `calc_AIQ` function. This metric quantifies the overall performance-cost trade-off for a given evaluation and router, useful for comparing different configurations.

```python
import pandas as pd
import numpy as np
from evaluation.AIQ import calc_AIQ

results = pd.read_csv("data/routerbench/eval_results/combined_results.csv")
eval_name = "mmlu"
router_name = "knn"

eval_results = results[results["eval_name"] == eval_name]
router_results = eval_results[eval_results["model_name"] == router_name]

# Split by embedding model for comparison
router_dfs, router_names = [], []
for embedding in router_results["embedding"].unique():
    sub_df = router_results[router_results["embedding"] == embedding]
    if len(sub_df[["total_cost", "performance"]].drop_duplicates()) > 2:
        router_dfs.append(sub_df)
        router_names.append(f"embedding:{embedding}")

aiq_scores = calc_AIQ(
    router_dfs,
    plot_graph=True,
    graph_title_eval_name=eval_name,
    router_results_names=router_names,
)
for name, score in zip(router_names, aiq_scores):
    print(f"{name}: AIQ={score:.4f}")
```
Tokens:
Raw