Writing a Custom Provider

import {
  generateId,
  loadApiKey,
  withoutTrailingSlash,
} from '@ai-sdk/provider-utils';
import { ProviderV2 } from '@ai-sdk/provider';
import { CustomChatLanguageModel } from './custom-chat-language-model';

// Define your provider interface extending ProviderV2
interface CustomProvider extends ProviderV2 {
  (modelId: string, settings?: CustomChatSettings): CustomChatLanguageModel;

  // Add specific methods for different model types
  languageModel(
    modelId: string,
    settings?: CustomChatSettings,
  ): CustomChatLanguageModel;
}

// Provider settings
interface CustomProviderSettings {
  /**
   * Base URL for API calls
   */
  baseURL?: string;

  /**
   * API key for authentication
   */
  apiKey?: string;

  /**
   * Custom headers for requests
   */
  headers?: Record<string, string>;
}

// Factory function to create provider instance
function createCustom(options: CustomProviderSettings = {}): CustomProvider {
  const createChatModel = (
    modelId: string,
    settings: CustomChatSettings = {},
  ) =>
    new CustomChatLanguageModel(modelId, settings, {
      provider: 'custom',
      baseURL:
        withoutTrailingSlash(options.baseURL) ?? 'https://api.custom.ai/v1',
      headers: () => ({
        Authorization: `Bearer ${loadApiKey({
          apiKey: options.apiKey,
          environmentVariableName: 'CUSTOM_API_KEY',
          description: 'Custom Provider',
        })}`,
        ...options.headers,
      }),
      generateId: options.generateId ?? generateId,
    });

  const provider = function (modelId: string, settings?: CustomChatSettings) {
    if (new.target) {
      throw new Error(
        'The model factory function cannot be called with the new keyword.',
      );
    }

    return createChatModel(modelId, settings);
  };

  provider.languageModel = createChatModel;

  return provider as CustomProvider;
}

// Export default provider instance
const custom = createCustom();

import { LanguageModelV2, LanguageModelV2CallOptions } from '@ai-sdk/provider';
import { postJsonToApi } from '@ai-sdk/provider-utils';

class CustomChatLanguageModel implements LanguageModelV2 {
  readonly specificationVersion = 'V2';
  readonly provider: string;
  readonly modelId: string;

  constructor(
    modelId: string,
    settings: CustomChatSettings,
    config: CustomChatConfig,
  ) {
    this.provider = config.provider;
    this.modelId = modelId;
    // Initialize with settings and config
  }

  // Convert AI SDK prompt to provider format
  private getArgs(options: LanguageModelV2CallOptions) {
    const warnings: LanguageModelV2CallWarning[] = [];

    // Map messages to provider format
    const messages = this.convertToProviderMessages(options.prompt);

    // Handle tools if provided
    const tools = options.tools
      ? this.prepareTools(options.tools, options.toolChoice)
      : undefined;

    // Build request body
    const body = {
      model: this.modelId,
      messages,
      temperature: options.temperature,
      max_tokens: options.maxOutputTokens,
      stop: options.stopSequences,
      tools,
      // ... other parameters
    };

    return { args: body, warnings };
  }

  async doGenerate(options: LanguageModelV2CallOptions) {
    const { args, warnings } = this.getArgs(options);

    // Make API call
    const response = await postJsonToApi({
      url: `${this.config.baseURL}/chat/completions`,
      headers: this.config.headers(),
      body: args,
      abortSignal: options.abortSignal,
    });

    // Convert provider response to AI SDK format
    const content: LanguageModelV2Content[] = [];

    // Extract text content
    if (response.choices[0].message.content) {
      content.push({
        type: 'text',
        text: response.choices[0].message.content,
      });
    }

    // Extract tool calls
    if (response.choices[0].message.tool_calls) {
      for (const toolCall of response.choices[0].message.tool_calls) {
        content.push({
          type: 'tool-call',
          toolCallType: 'function',
          toolCallId: toolCall.id,
          toolName: toolCall.function.name,
          args: JSON.stringify(toolCall.function.arguments),
        });
      }
    }

    return {
      content,
      finishReason: this.mapFinishReason(response.choices[0].finish_reason),
      usage: {
        inputTokens: response.usage?.prompt_tokens,
        outputTokens: response.usage?.completion_tokens,
        totalTokens: response.usage?.total_tokens,
      },
      request: { body: args },
      response: { body: response },
      warnings,
    };
  }

  async doStream(options: LanguageModelV2CallOptions) {
    const { args, warnings } = this.getArgs(options);

    // Create streaming response
    const response = await fetch(`${this.config.baseURL}/chat/completions`, {
      method: 'POST',
      headers: {
        ...this.config.headers(),
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({ ...args, stream: true }),
      signal: options.abortSignal,
    });

    // Transform stream to AI SDK format
    const stream = response
      .body!.pipeThrough(new TextDecoderStream())
      .pipeThrough(this.createParser())
      .pipeThrough(this.createTransformer(warnings));

    return { stream, warnings };
  }

  // Supported URL patterns for native file handling
  get supportedUrls() {
    return {
      'image/*': [/^https:\/\/example\.com\/images\/.*/],
    };
  }
}

private convertToProviderMessages(prompt: LanguageModelV2Prompt) {
  return prompt.map((message) => {
    switch (message.role) {
      case 'system':
        return { role: 'system', content: message.content };

      case 'user':
        return {
          role: 'user',
          content: message.content.map((part) => {
            switch (part.type) {
              case 'text':
                return { type: 'text', text: part.text };
              case 'file':
                return {
                  type: 'image_url',
                  image_url: {
                    url: this.convertFileToUrl(part.data),
                  },
                };
              default:
                throw new Error(`Unsupported part type: ${part.type}`);
            }
          }),
        };

      case 'assistant':
        // Handle assistant messages with text, tool calls, etc.
        return this.convertAssistantMessage(message);

      case 'tool':
        // Handle tool results
        return this.convertToolMessage(message);

      default:
        throw new Error(`Unsupported message role: ${message.role}`);
    }
  });
}

private createTransformer(warnings: LanguageModelV2CallWarning[]) {
  let isFirstChunk = true;

  return new TransformStream<ParsedChunk, LanguageModelV2StreamPart>({
    async transform(chunk, controller) {
      // Send warnings with first chunk
      if (isFirstChunk) {
        controller.enqueue({ type: 'stream-start', warnings });
        isFirstChunk = false;
      }

      // Handle different chunk types
      if (chunk.choices?.[0]?.delta?.content) {
        controller.enqueue({
          type: 'text',
          text: chunk.choices[0].delta.content,
        });
      }

      if (chunk.choices?.[0]?.delta?.tool_calls) {
        for (const toolCall of chunk.choices[0].delta.tool_calls) {
          controller.enqueue({
            type: 'tool-call-delta',
            toolCallType: 'function',
            toolCallId: toolCall.id,
            toolName: toolCall.function.name,
            argsTextDelta: toolCall.function.arguments,
          });
        }
      }

      // Handle finish reason
      if (chunk.choices?.[0]?.finish_reason) {
        controller.enqueue({
          type: 'finish',
          finishReason: this.mapFinishReason(chunk.choices[0].finish_reason),
          usage: {
            inputTokens: chunk.usage?.prompt_tokens,
            outputTokens: chunk.usage?.completion_tokens,
            totalTokens: chunk.usage?.total_tokens,
          },
        });
      }
    },
  });
}

import {
  APICallError,
  InvalidResponseDataError,
  TooManyRequestsError,
} from '@ai-sdk/provider';

private handleError(error: unknown): never {
  if (error instanceof Response) {
    const status = error.status;

    if (status === 429) {
      throw new TooManyRequestsError({
        cause: error,
        retryAfter: this.getRetryAfter(error),
      });
    }

    throw new APICallError({
      statusCode: status,
      statusText: error.statusText,
      cause: error,
      isRetryable: status >= 500 && status < 600,
    });
  }

  throw error;
}