straico-custom-provider/src/provider/straico-chat-language-model.ts

Straico custom provider for Vercel AI SDK

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

// Define interfaces locally since they're not exported from the module
interface ContentPart {
  type: string;
  [key: string]: any;
}

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

interface LanguageModelV1 {
  readonly specificationVersion: string;
  readonly provider: string;
  readonly modelId: string;
  readonly defaultObjectGenerationMode?: string;
  doGenerate(parts: ContentPart[], options: any): Promise<ContentPart[]>;
  doStream(parts: ContentPart[], options: any): Promise<AsyncIterable<ContentPart>>;
}

interface GenerateOptions {
  signal?: AbortSignal;
}

interface StreamOptions {
  signal?: AbortSignal;
}

interface ModelSettings {
  [key: string]: any;
}

// Helper function to check if a part is a TextContentPart
function isTextContentPart(part: ContentPart): part is TextContentPart {
  return part.type === 'text' && typeof part.text === 'string';
}

import type { 
  StraicoChatModelId, 
  StraicoChatSettings 
} from './straico-chat-settings';

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

// Define the request arguments type
interface StraicoChatRequestArgs {
  models: string[];
  message: string;
  file_urls?: string[];
  youtube_urls?: string[];
  max_tokens?: number;
  temperature?: number;
}

export class StraicoChatLanguageModel implements LanguageModelV1 {
  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: GenerateOptions & ModelSettings
  ): 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 APICallError({
          message: `Straico API returned an error: ${response.statusText}`,
          url: `${this.baseURL}/prompt/completion`,
          requestBodyValues: args,
          statusCode: response.status,
          cause: new Error(response.statusText),
          isRetryable: response.status >= 500 || response.status === 429,
        });
      }

      const data = await response.json();

      if (!data.success) {
        throw new InvalidResponseDataError({
          data: data,
          message: '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 InvalidResponseDataError({
          data: data,
          message: `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 APICallError || error instanceof InvalidResponseDataError) {
        throw error;
      }

      throw new APICallError({
        message: `Error calling Straico API: ${(error as Error).message}`,
        url: `${this.baseURL}/prompt/completion`,
        requestBodyValues: args,
        cause: error as Error,
        isRetryable: true,
      });
    }
  }

  async doStream(
    parts: ContentPart[],
    options: StreamOptions & ModelSettings
  ): 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: GenerateOptions | StreamOptions & ModelSettings
  ): 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;
  }
}