@nomyx/assistant/src/ai/providers/VertexProvider/index.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 { ProviderCapabilities, EmbeddingProvider } from '../../../types/provider';
import { Tool, GenericToolSchema, StandardizedToolCall } from '../../../types/tool';
import { ILogger } from '../../../types/common';
import { BaseProvider } from '../BaseProvider';
import { VertexProviderConfig, VertexModel } from './types';
import { VertexConfig } from './config';
import { VertexCache } from './cache';
import { chat } from './chat';
import { streamChat } from './streaming';
import { embed, embedBatch } from './embedding';
import { convertToolSchema, convertToolCall } from './tools';
import { handleVertexError, VertexProviderError } from './errors';

export class VertexProvider extends BaseProvider implements EmbeddingProvider {
  private config: VertexConfig;
  private cache: VertexCache;

  constructor(config: VertexProviderConfig, logger: ILogger) {
    super(logger);
    this.config = new VertexConfig(config, logger);
    this.cache = new VertexCache();
  }

  async chat(messages: ChatMessage[], options: ChatOptions, tools?: Tool[]): Promise<ProviderResponse> {
    try {
      const cacheKey = { messages, options, tools };
      if (this.cache.has(cacheKey)) {
        this.logger.debug('Cache hit for chat request');
        return this.cache.get(cacheKey)!;
      }

      const processedMessages = await this.applyMiddleware(messages);
      const response = await chat(this.config, processedMessages, options, tools);
      const processedResponse = await this.applyMiddlewareToResponse(response);
      if (processedResponse.toolCalls) {
        processedResponse.toolCalls = await this.applyMiddlewareToToolCalls(processedResponse.toolCalls);
      }

      this.cache.set(cacheKey, processedResponse);
      return processedResponse;
    } catch (error: unknown) {
      this.logger.error('Error in VertexProvider chat', { error: error instanceof Error ? error.message : String(error) });
      throw handleVertexError(error);
    }
  }

  async *streamChat(messages: ChatMessage[], options: ChatOptions, tools?: Tool[]): AsyncIterableIterator<ProviderResponse> {
    try {
      const processedMessages = await this.applyMiddleware(messages);
      for await (const response of streamChat(this.config, processedMessages, options, tools)) {
        const processedResponse = await this.applyMiddlewareToResponse(response);
        if (processedResponse.toolCalls) {
          processedResponse.toolCalls = await this.applyMiddlewareToToolCalls(processedResponse.toolCalls);
        }
        yield processedResponse;
      }
    } catch (error: unknown) {
      this.logger.error('Error in VertexProvider streamChat', { error: error instanceof Error ? error.message : String(error) });
      throw handleVertexError(error);
    }
  }

  getCapabilities(): ProviderCapabilities {
    return {
      maxTokens: 2048, // This may vary depending on the specific model
      supportsFunctionCalling: true,
      supportsStreaming: true,
      supportedModels: Object.values(VertexModel),
      maxSimultaneousCalls: 1,
      supportsSemanticCaching: false,
    };
  }

  async embed(text: string): Promise<number[]> {
    try {
      return await embed(this.config, text);
    } catch (error: unknown) {
      this.logger.error('Error in VertexProvider embed', { error: error instanceof Error ? error.message : String(error) });
      throw handleVertexError(error);
    }
  }

  async embedBatch(texts: string[]): Promise<number[][]> {
    try {
      return await embedBatch(this.config, texts);
    } catch (error: unknown) {
      this.logger.error('Error in VertexProvider embedBatch', { error: error instanceof Error ? error.message : String(error) });
      throw handleVertexError(error);
    }
  }

  convertToolSchema(schema: GenericToolSchema): any {
    return convertToolSchema(schema, this.config.getModelType());
  }

  convertToolCall(call: any): StandardizedToolCall {
    return convertToolCall(call, this.config.getModelType());
  }

  clearCache(): void {
    this.cache.clear();
    this.logger.debug('VertexProvider cache cleared');
  }
}

export * from './types';
export { VertexProviderError } from './errors';