Using a custom model

Training custom routers is not limited to using models that we support. You can also include custom models that you've fine-tuned, or even an agent. Any arbitrary inference endpoint may be specified. Just include its evaluation results in the training data and you'll be able to use it in your routing decisions.

To define a custom model in your training, include it in your llm_providers list with "is_custom": true. You must define:

context_length
input_price
output_price
latency

so that we can reason about quality, cost, and latency when routing.

🚧
The custom provider/model string needs to be unique
When defining a custom model, make sure the provider/model pair is unique:

It cannot match any supported model.

It cannot match any other custom model.

It cannot have more than one forward slash /

Download the evaluation data

curl -L "https://drive.google.com/uc?export=download&id=16gZsbU71HQue59-A4hFK_XwPIebR-CnF" -o evaluation_data.csv

Training with custom models

When training a custom router, include your custom model configuration in the llm_providers array alongside standard models.

import os
import json
from notdiamond import Notdiamond

# Initialize ND client
client = Notdiamond(api_key=os.environ.get("NOTDIAMOND_API_KEY"))

# Define your custom model
custom_model = {
    "provider": "custom",
    "model": "my-finetuned-model",
    "is_custom": True,
    "context_length": 200000,
    "input_price": 0.1,   # USD per million input tokens
    "output_price": 0.2,  # USD per million output tokens
    "latency": 0.01       # time to first token (seconds)
}

# LLMs to train router on (mix of standard and custom models)
llm_providers = [
    {"provider": "openai", "model": "gpt-4o"},
    {"provider": "openai", "model": "gpt-4o-mini"},
    {"provider": "google", "model": "gemini-2.5-pro"},
    {"provider": "anthropic", "model": "claude-opus-4-20250514"},
    custom_model,  # Add your custom model to the list
]

# Train the custom router (assuming the evaluation_data.csv exists)
with open("evaluation_data.csv", "rb") as f:
    response = client.custom_router.train_custom_router(
        dataset_file=f,
        language="english",
        llm_providers=json.dumps(llm_providers),
        maximize=True,
        prompt_column="prompt",
    )

preference_id = response.preference_id
print("Custom router preference ID:", preference_id)

import NotDiamond from 'notdiamond';
import * as fs from 'fs';

// Initialize ND client
const client = new NotDiamond({api_key: process.env.NOTDIAMOND_API_KEY});

// Define your custom model
const customModel = {
  provider: 'custom',
  model: 'my-finetuned-model',
  is_custom: true,
  context_length: 200000,
  input_price: 0.1,   // USD per million input tokens
  output_price: 0.2,  // USD per million output tokens
  latency: 0.01       // time to first token (seconds)
};

// LLMs to train router on (mix of standard and custom models)
const llmProviders = [
  { provider: 'openai', model: 'gpt-4o' },
  { provider: 'openai', model: 'gpt-4o-mini' },
  { provider: 'google', model: 'gemini-2.5-pro' },
  { provider: 'anthropic', model: 'claude-opus-4-20250514' },
  customModel,  // Add your custom model to the list
];

// Train the custom router (assuming the evaluation_data.csv exists)
const datasetFile = fs.createReadStream('evaluation_data.csv');

const response = await client.customRouter.trainCustomRouter({
  dataset_file: datasetFile,
  language: 'english',
  llm_providers: JSON.stringify(llmProviders),
  maximize: true,
  prompt_column: 'prompt',
});

const preferenceId = response.preference_id;
console.log('Custom router preference ID:', preferenceId);

Routing with custom models

Once your custom router is trained, you can use it for routing by including your custom model configuration in the llm_providers list when calling model_router.select_model().

messages = [
    {"role": "user", "content": "Write merge sort in 3 lines."}
]

result = client.model_router.select_model(
    messages=messages,
    llm_providers=llm_providers,
    preference_id=preference_id,  # Preference ID from training
)

print("LLM selected:", result)

const messages = [
  { role: 'user', content: 'Write merge sort in 3 lines.' }
];

const result = await client.modelRouter.selectModel({
  messages: messages,
  llmProviders: llmProviders,
  preferenceId: preferenceId,  // Preference ID from training
});

console.log('LLM selected:', result);