Rabbit SDK Logo
Rabbit SDK
GitHub
Documentation

Advanced

Custom Providers

Build a custom provider for Rabbit SDK

You can easily integrate any LLM API by extending the abstract Provider class.

Why build a custom provider?

While Rabbit SDK ships with OpenAI, Groq, Gemini, and Anthropic providers out of the box, you might want to use:

  • A custom internal LLM deployment
  • An open-source model hosted on vLLM or Ollama
  • A new or specialized AI API like Mistral or Cohere

Extending the Provider Class

Extend the abstract Provider class and implement the generate method:

typescript
import { Agent, Provider, ProviderRequest, ProviderResponse } from "@rabbit-agent-sdk/rabbit-agent-sdk";
 
class MistralProvider extends Provider {
  public name = "mistral";
  private apiKey: string;
 
  constructor(config: { apiKey: string; model?: string }) {
    super(config.model || "mistral-large-latest");
    this.apiKey = config.apiKey;
  }
 
  public async generate(
    request: ProviderRequest
  ): Promise<ProviderResponse | AsyncIterable<ProviderResponse>> {
 
    // 1. Map the standardized messages to Mistral's format
    const messages = request.messages.map((m) => ({
      role: m.role,
      content: m.content,
    }));
 
    if (request.systemPrompt) {
      messages.unshift({ role: "system", content: request.systemPrompt });
    }
 
    // 2. Call the Mistral API
    const response = await fetch("https://api.mistral.ai/v1/chat/completions", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": `Bearer ${this.apiKey}`,
      },
      body: JSON.stringify({
        model: request.model,
        messages,
        max_tokens: request.maxTokens ?? 1024,
      }),
    });
 
    const data = await response.json();
 
    // 3. Return the standardized ProviderResponse
    return {
      message: {
        role: "assistant",
        content: data.choices[0]?.message?.content || "",
      },
      usage: {
        promptTokens: data.usage?.prompt_tokens || 0,
        completionTokens: data.usage?.completion_tokens || 0,
        totalTokens: data.usage?.total_tokens || 0,
      },
    };
  }
}
 
// Use it just like any other provider
const agent = new Agent({
  provider: new MistralProvider({
    apiKey: process.env.MISTRAL_API_KEY!,
    model: "mistral-large-latest",
  }),
});