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

Straico custom provider for Vercel AI SDK

/**
 * Simplified implementation of a Straico provider for the Vercel AI SDK
 * This doesn't rely on the @ai-sdk/provider package and is intended as a starting point
 */

// Define minimal environment interface
declare const process: {
  env: Record<string, string | undefined>;
};

export interface ContentPart {
  type: string;
  [key: string]: any;
}

export interface TextContentPart extends ContentPart {
  type: 'text';
  text: string;
}

export function isTextContentPart(part: ContentPart): part is TextContentPart {
  return part.type === 'text' && typeof part.text === 'string';
}

export type StraicoChatModelId = string;

export interface StraicoChatSettings {
  maxTokens?: number;
  temperature?: number;
  fileUrls?: string[];
  youtubeUrls?: string[];
}

interface StraicoChatRequestArgs {
  models: string[];
  message: string;
  file_urls?: string[];
  youtube_urls?: string[];
  max_tokens?: number;
  temperature?: number;
}

export interface StraicoChatModelOptions {
  provider: string;
  baseURL: string;
  headers: () => Record<string, string>;
  generateId: () => string;
}

export class StraicoChatLanguageModel {
  readonly specificationVersion = 'v1';
  readonly provider: string;
  readonly modelId: string;
  readonly defaultObjectGenerationMode = 'json';

  private readonly baseURL: string;
  private readonly getHeaders: () => Record<string, string>;
  private readonly generateId: () => string;
  private readonly defaultSettings: StraicoChatSettings;

  constructor(
    modelId: StraicoChatModelId,
    settings: StraicoChatSettings,
    options: StraicoChatModelOptions
  ) {
    this.modelId = modelId;
    this.provider = options.provider;
    this.baseURL = options.baseURL;
    this.getHeaders = options.headers;
    this.generateId = options.generateId;
    this.defaultSettings = settings;
  }

  async doGenerate(
    parts: ContentPart[],
    options: any = {}
  ): Promise<ContentPart[]> {
    const { signal } = options;
    const args = this.getArgs(parts, options);
    
    try {
      const response = await fetch(
        `${this.baseURL}/prompt/completion`,
        {
          method: 'POST',
          headers: this.getHeaders(),
          body: JSON.stringify(args),
          signal,
        }
      );

      if (!response.ok) {
        throw new Error(`Straico API returned an error: ${response.statusText}`);
      }

      const data = await response.json();

      if (!data.success) {
        throw new Error('Straico API returned an unsuccessful response');
      }

      // Extract the model completion for the specified modelId
      const modelResult = data.data.completions[this.modelId];
      if (!modelResult) {
        throw new Error(`No completion found for model ${this.modelId}`);
      }

      // Extract the content from the completion
      const content = modelResult.completion.choices[0].message.content;
      return [{ type: 'text', text: content }];
    } catch (error) {
      if (error instanceof Error) {
        throw error;
      }
      throw new Error(`Error calling Straico API: ${String(error)}`);
    }
  }

  async doStream(
    parts: ContentPart[],
    options: any = {}
  ): Promise<AsyncIterable<ContentPart>> {
    // Straico API doesn't support streaming natively
    // We'll simulate streaming by breaking up the complete response
    const completion = await this.doGenerate(parts, options);
    
    return this.simulateStream(completion);
  }

  private async *simulateStream(
    completionParts: ContentPart[]
  ): AsyncIterable<ContentPart> {
    for (const part of completionParts) {
      if (isTextContentPart(part)) {
        // Simulate streaming by yielding character by character
        // In a real implementation, you might want to yield chunks instead
        const text = part.text;
        let currentText = '';
        
        for (const char of text) {
          currentText += char;
          yield { type: 'text', text: currentText };
          // Small delay to simulate streaming
          await new Promise(resolve => setTimeout(resolve, 10));
        }
      } else {
        yield part;
      }
    }
  }

  private getArgs(parts: ContentPart[], options: any): StraicoChatRequestArgs {
    // Convert Vercel AI SDK format to Straico API format
    const mergedSettings = { ...this.defaultSettings, ...options };
    
    // Extract the message from content parts
    const message = parts
      .filter(isTextContentPart)
      .map(part => part.text)
      .join('\n');

    // Basic Straico API request structure
    const args: StraicoChatRequestArgs = {
      models: [this.modelId],
      message: message
    };

    // Add optional parameters if they exist
    if (mergedSettings.fileUrls && mergedSettings.fileUrls.length > 0) {
      args.file_urls = mergedSettings.fileUrls;
    }

    if (mergedSettings.youtubeUrls && mergedSettings.youtubeUrls.length > 0) {
      args.youtube_urls = mergedSettings.youtubeUrls;
    }

    if (mergedSettings.maxTokens) {
      args.max_tokens = mergedSettings.maxTokens;
    }

    if (mergedSettings.temperature !== undefined) {
      args.temperature = mergedSettings.temperature;
    }

    return args;
  }
}

// Helper function to load API key from options or environment
export function loadApiKey(options: {
  apiKey?: string;
  environmentVariableName: string;
  description: string;
}): string {
  const { apiKey, environmentVariableName, description } = options;
  
  if (apiKey) {
    return apiKey;
  }
  
  // For browser environments, we only use provided apiKey
  if (typeof window !== 'undefined') {
    throw new Error(
      `API key for ${description} not found. Please provide an apiKey option.`
    );
  }
  
  // For Node.js environments
  try {
    // @ts-ignore - Accessing process.env
    const envApiKey = process.env[environmentVariableName];
    if (!envApiKey) {
      throw new Error();
    }
    return envApiKey;
  } catch (e) {
    throw new Error(
      `API key for ${description} not found. Either provide an apiKey option or set the ${environmentVariableName} environment variable.`
    );
  }
}

// Helper function to generate a unique ID
export function generateId(): string {
  return `gen-${Math.random().toString(36).substring(2, 12)}`;
}

// Helper function to ensure URLs don't have trailing slashes
export function withoutTrailingSlash(url?: string): string | undefined {
  if (!url) return undefined;
  return url.endsWith('/') ? url.slice(0, -1) : url;
}

// 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 {
  baseURL?: string;
  apiKey?: string;
  headers?: Record<string, string>;
  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();