@nomyx/assistant/src/ai/AIProvider.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import { ChatMessage, ChatOptions, ProviderResponse } from "../types/chat";
import { AIProviderError, AIProviderMiddleware, AIProviderPlugin, ProviderCapabilities, IAIProvider } from "../types/provider";
import { GenericToolSchema, StandardizedToolCall } from "../types/tool";

export abstract class AIProvider implements IAIProvider {
  protected middleware: AIProviderMiddleware[] = [];
  protected plugins: AIProviderPlugin[] = [];

  abstract chat(messages: ChatMessage[], options: ChatOptions): Promise<ProviderResponse>;
  abstract streamChat(messages: ChatMessage[], options: ChatOptions): AsyncIterableIterator<ProviderResponse>;
  abstract convertToolSchema(schema: GenericToolSchema): unknown;
  abstract convertToolCall(call: unknown): StandardizedToolCall;
  abstract getCapabilities(): ProviderCapabilities;

  use(middleware: AIProviderMiddleware): void {
    this.middleware.push(middleware);
  }

  registerPlugin(plugin: AIProviderPlugin): void {
    this.plugins.push(plugin);
    plugin.initialize(this);
  }

  protected async applyMiddleware(messages: ChatMessage[]): Promise<ChatMessage[]> {
    let processedMessages = messages;
    for (const mw of this.middleware) {
      processedMessages = await mw.preProcess(processedMessages);
    }
    return processedMessages;
  }

  protected async applyMiddlewareToResponse(response: ProviderResponse): Promise<ProviderResponse> {
    let processedResponse = response;
    for (const mw of this.middleware.slice().reverse()) {
      processedResponse = await mw.postProcess(processedResponse);
    }
    return processedResponse;
  }

  protected async notifyPluginsBeforeChat(messages: ChatMessage[], options: ChatOptions): Promise<void> {
    for (const plugin of this.plugins) {
      if (plugin.onBeforeChat) {
        await plugin.onBeforeChat(messages, options);
      }
    }
  }

  protected async notifyPluginsAfterChat(response: ProviderResponse): Promise<void> {
    for (const plugin of this.plugins) {
      if (plugin.onAfterChat) {
        await plugin.onAfterChat(response);
      }
    }
  }

  protected async notifyPluginsOnError(error: Error): Promise<void> {
    const aiProviderError = new AIProviderError(
      error.message,
      'UNKNOWN_ERROR',
      'unknown',
      error,
      false
    );
    for (const plugin of this.plugins) {
      if (plugin.onError) {
        await plugin.onError(aiProviderError);
      }
    }
  }

  protected async notifyPluginsOnStreamStart(): Promise<void> {
    for (const plugin of this.plugins) {
      if (plugin.onStreamStart) {
        await plugin.onStreamStart();
      }
    }
  }

  protected async notifyPluginsOnStreamChunk(chunk: ProviderResponse): Promise<void> {
    for (const plugin of this.plugins) {
      if (plugin.onStreamChunk) {
        await plugin.onStreamChunk(chunk);
      }
    }
  }

  protected async notifyPluginsOnStreamEnd(): Promise<void> {
    for (const plugin of this.plugins) {
      if (plugin.onStreamEnd) {
        await plugin.onStreamEnd();
      }
    }
  }
}