@ai-sdk/mistral/src/mistral-provider.ts

import {
  NoSuchModelError,
  type EmbeddingModelV4,
  type LanguageModelV4,
  type ProviderV4,
  type SpeechModelV4,
  type TranscriptionModelV4,
} from '@ai-sdk/provider';
import {
  resolveWith,
  type Resolvable,
  loadApiKey,
  withoutTrailingSlash,
  withUserAgentSuffix,
  type FetchFunction,
  type WebSocketConstructor,
} from '@ai-sdk/provider-utils';
import { MistralChatLanguageModel } from './mistral-chat-language-model';
import { MistralConversationLanguageModel } from './mistral-conversation-language-model';
import type { MistralChatModelId } from './mistral-chat-language-model-options';
import { MistralEmbeddingModel } from './mistral-embedding-model';
import type { MistralEmbeddingModelId } from './mistral-embedding-model-options';
import { MistralSpeechModel } from './mistral-speech-model';
import type { MistralSpeechModelId } from './mistral-speech-model-options';
import { MistralTranscriptionModel } from './mistral-transcription-model';
import type { MistralTranscriptionModelId } from './mistral-transcription-model-options';
import { VERSION } from './version';
import { mistralTools } from './tool';
import {
  createMistralRealtimeFactory,
  type MistralRealtimeFactory,
} from './realtime/mistral-realtime-factory';

export interface MistralProvider extends ProviderV4 {
  tools: typeof mistralTools;
  /** Creates managed voice agents and their server-side sessions. */
  experimental_realtime: MistralRealtimeFactory;
  (modelId: MistralChatModelId): LanguageModelV4;

  /**
   * Creates a model for text generation.
   */
  languageModel(modelId: MistralChatModelId): LanguageModelV4;

  /**
   * Creates a model for text generation.
   */
  chat(modelId: MistralChatModelId): LanguageModelV4;

  /**
   * Creates a model using the Conversations API, with support for web search.
   */
  conversation(modelId: MistralChatModelId): LanguageModelV4;

  /**
   * Creates a model for text embeddings.
   */
  embedding(modelId: MistralEmbeddingModelId): EmbeddingModelV4;

  /**
   * Creates a model for text embeddings.
   */
  embeddingModel: (modelId: MistralEmbeddingModelId) => EmbeddingModelV4;

  /**
   * Creates a model for speech generation (text-to-speech).
   */
  speech(modelId: MistralSpeechModelId): SpeechModelV4;

  /**
   * Creates a model for speech generation (text-to-speech).
   */
  speechModel(modelId: MistralSpeechModelId): SpeechModelV4;

  /**
   * Creates a model for audio transcription.
   */
  transcription(modelId: MistralTranscriptionModelId): TranscriptionModelV4;

  /**
   * Creates a model for audio transcription.
   */
  transcriptionModel(
    modelId: MistralTranscriptionModelId,
  ): TranscriptionModelV4;

  /**
   * @deprecated Use `embedding` instead.
   */
  textEmbedding(modelId: MistralEmbeddingModelId): EmbeddingModelV4;

  /**
   * @deprecated Use `embeddingModel` instead.
   */
  textEmbeddingModel(modelId: MistralEmbeddingModelId): EmbeddingModelV4;
}

export interface MistralProviderSettings {
  /**
   * Use a different URL prefix for API calls, e.g. to use proxy servers.
   * The default prefix is `https://api.mistral.ai/v1`.
   */
  baseURL?: string;

  /**
   * API key that is being send using the `Authorization` header.
   * It defaults to the `MISTRAL_API_KEY` environment variable.
   */
  apiKey?: Resolvable<string>;

  /**
   * Custom headers to include in the requests.
   */
  headers?: Record<string, string>;

  /**
   * Custom fetch implementation. You can use it as a middleware to intercept requests,
   * or to provide a custom fetch implementation for e.g. testing.
   */
  fetch?: FetchFunction;

  /**
   * WebSocket implementation for realtime transcription. Use a header-capable
   * implementation such as `ws` to send Mistral's Authorization header.
   */
  webSocket?: WebSocketConstructor;

  generateId?: () => string;
}

/**
 * Create a Mistral AI provider instance.
 */
export function createMistral(
  options: MistralProviderSettings = {},
): MistralProvider {
  const baseURL =
    withoutTrailingSlash(options.baseURL) ?? 'https://api.mistral.ai/v1';

  const getHeaders = () =>
    resolveWith(options.apiKey, apiKey =>
      withUserAgentSuffix(
        {
          Authorization: `Bearer ${loadApiKey({
            apiKey,
            environmentVariableName: 'MISTRAL_API_KEY',
            description: 'Mistral',
          })}`,
          ...options.headers,
        },
        `ai-sdk-mistral/${VERSION}`,
      ),
    );

  const createChatModel = (modelId: MistralChatModelId) =>
    new MistralChatLanguageModel(modelId, {
      provider: 'mistral.chat',
      baseURL,
      headers: getHeaders,
      fetch: options.fetch,
      generateId: options.generateId,
    });

  const createEmbeddingModel = (modelId: MistralEmbeddingModelId) =>
    new MistralEmbeddingModel(modelId, {
      provider: 'mistral.embedding',
      baseURL,
      headers: getHeaders,
      fetch: options.fetch,
    });

  const createConversationModel = (modelId: MistralChatModelId) =>
    new MistralConversationLanguageModel(modelId, {
      provider: 'mistral.conversation',
      baseURL,
      headers: getHeaders,
      fetch: options.fetch,
      generateId: options.generateId,
    });

  const createSpeechModel = (modelId: MistralSpeechModelId) =>
    new MistralSpeechModel(modelId, {
      provider: 'mistral.speech',
      baseURL,
      headers: getHeaders,
      fetch: options.fetch,
    });

  const createTranscriptionModel = (modelId: MistralTranscriptionModelId) =>
    new MistralTranscriptionModel(modelId, {
      provider: 'mistral.transcription',
      baseURL,
      headers: getHeaders,
      fetch: options.fetch,
      webSocket: options.webSocket,
    });

  const provider = function (modelId: MistralChatModelId) {
    if (new.target) {
      throw new Error(
        'The Mistral model function cannot be called with the new keyword.',
      );
    }

    return createChatModel(modelId);
  };

  provider.specificationVersion = 'v4' as const;
  provider.languageModel = createChatModel;
  provider.chat = createChatModel;
  provider.conversation = createConversationModel;
  provider.tools = mistralTools;
  provider.experimental_realtime = createMistralRealtimeFactory({
    baseURL,
    headers: getHeaders,
    fetch: options.fetch,
  });
  provider.embedding = createEmbeddingModel;
  provider.embeddingModel = createEmbeddingModel;
  provider.textEmbedding = createEmbeddingModel;
  provider.textEmbeddingModel = createEmbeddingModel;
  provider.speech = createSpeechModel;
  provider.speechModel = createSpeechModel;
  provider.transcription = createTranscriptionModel;
  provider.transcriptionModel = createTranscriptionModel;

  provider.imageModel = (modelId: string) => {
    throw new NoSuchModelError({ modelId, modelType: 'imageModel' });
  };

  return provider;
}

/**
 * Default Mistral provider instance.
 */
export const mistral = createMistral();