straico-custom-provider/src/provider/straico-provider.ts

Straico custom provider for Vercel AI SDK

import {
  generateId,
  loadApiKey,
  withoutTrailingSlash,
} from '@ai-sdk/provider-utils';
import { StraicoChatLanguageModel } from './straico-chat-language-model';
import type { StraicoChatModelId, StraicoChatSettings } from './straico-chat-settings';

// Model factory function with additional methods and properties
export interface StraicoProvider {
  (
    modelId: StraicoChatModelId,
    settings?: StraicoChatSettings,
  ): StraicoChatLanguageModel;

  // Explicit method for targeting specific API
  chat(
    modelId: StraicoChatModelId,
    settings?: StraicoChatSettings,
  ): StraicoChatLanguageModel;
}

// Optional settings for the provider
export interface StraicoProviderSettings {
  /**
   * Use a different URL prefix for API calls, e.g. to use proxy servers.
   */
  baseURL?: string;

  /**
   * API key.
   */
  apiKey?: string;

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

  /**
   * Custom ID generator function.
   */
  generateId?: () => string;
}

// Provider factory function
export function createStraicioProvider(
  options: StraicoProviderSettings = {},
): StraicoProvider {
  const createModel = (
    modelId: StraicoChatModelId,
    settings: StraicoChatSettings = {},
  ) =>
    new StraicoChatLanguageModel(modelId, settings, {
      provider: 'straico.chat',
      baseURL:
        withoutTrailingSlash(options.baseURL) ?? 'https://api.straico.com/v1',
      headers: () => ({
        Authorization: `Bearer ${loadApiKey({
          apiKey: options.apiKey,
          environmentVariableName: 'STRAICO_API_KEY',
          description: 'Straico Provider',
        })}`,
        'Content-Type': 'application/json',
        ...options.headers,
      }),
      generateId: options.generateId ?? generateId,
    });

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

    return createModel(modelId, settings);
  };

  provider.chat = createModel;

  return provider;
}

/**
 * Default Straico provider instance.
 */
export const straico = createStraicioProvider();