@nomyx/assistant/src/ai/providers/VertexProvider/tools.ts

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

import { ChatMessage } from '../../../types/chat';
import { GenericToolSchema, StandardizedToolCall, ParameterDefinition } from '../../../types/tool';
import { VertexToolSchema, VertexModel } from './types';

export function convertToolSchema(schema: GenericToolSchema, modelType: VertexModel): VertexToolSchema {
  if (modelType === VertexModel.GEMINI_PRO || modelType === VertexModel.GEMINI_PRO_VISION) {
    return convertGeminiToolSchema(schema);
  } else if (modelType === VertexModel.CLAUDE_3_SONNET || modelType === VertexModel.CLAUDE_3_OPUS) {
    return convertClaudeToolSchema(schema);
  } else {
    throw new Error(`Unsupported model type for tool schema conversion: ${modelType}`);
  }
}

function convertGeminiToolSchema(schema: GenericToolSchema): VertexToolSchema {
  return {
    name: schema.name,
    description: schema.description,
    parameters: {
      type: 'object',
      properties: Object.entries(schema.parameters.properties).reduce((acc, [key, value]) => {
        acc[key] = {
          type: value.type,
          description: value.description || '',
          ...(value.enum ? { enum: value.enum } : {})
        };
        return acc;
      }, {} as { [key: string]: { type: string; description: string; enum?: string[] } }),
      required: schema.parameters.required || []
    }
  };
}

function convertClaudeToolSchema(schema: GenericToolSchema): VertexToolSchema {
  return {
    name: schema.name,
    description: schema.description,
    parameters: {
      type: 'object',
      properties: Object.entries(schema.parameters.properties).reduce((acc, [key, value]) => {
        acc[key] = {
          type: value.type,
          description: value.description || '',
          ...(value.enum ? { enum: value.enum } : {})
        };
        return acc;
      }, {} as { [key: string]: { type: string; description: string; enum?: string[] } }),
      required: schema.parameters.required || []
    }
  };
}

export function convertToolCall(toolCall: any, modelType: VertexModel): StandardizedToolCall {
  if (modelType === VertexModel.GEMINI_PRO || modelType === VertexModel.GEMINI_PRO_VISION) {
    return convertGeminiToolCall(toolCall);
  } else if (modelType === VertexModel.CLAUDE_3_SONNET || modelType === VertexModel.CLAUDE_3_OPUS) {
    return convertClaudeToolCall(toolCall);
  } else {
    throw new Error(`Unsupported model type for tool call conversion: ${modelType}`);
  }
}

function convertGeminiToolCall(toolCall: any): StandardizedToolCall {
  return {
    name: toolCall.name,
    arguments: JSON.parse(toolCall.args)
  };
}

function convertClaudeToolCall(toolCall: any): StandardizedToolCall {
  return {
    name: toolCall.name,
    arguments: toolCall.input
  };
}

export function convertMessages(messages: ChatMessage[], modelType: VertexModel): any[] {
  if (modelType === VertexModel.GEMINI_PRO || modelType === VertexModel.GEMINI_PRO_VISION) {
    return convertGeminiMessages(messages);
  } else if (modelType === VertexModel.CLAUDE_3_SONNET || modelType === VertexModel.CLAUDE_3_OPUS) {
    return convertClaudeMessages(messages);
  } else {
    throw new Error(`Unsupported model type for message conversion: ${modelType}`);
  }
}

function convertGeminiMessages(messages: ChatMessage[]): any[] {
  return messages.map(msg => ({
    author: msg.role === 'user' ? 'user' : 'bot',
    content: msg.content
  }));
}

function convertClaudeMessages(messages: ChatMessage[]): any[] {
  return messages.map(msg => ({
    role: msg.role,
    content: msg.content
  }));
}

export function extractToolCalls(content: any, modelType: VertexModel): StandardizedToolCall[] {
  if (modelType === VertexModel.GEMINI_PRO || modelType === VertexModel.GEMINI_PRO_VISION) {
    return extractGeminiToolCalls(content);
  } else if (modelType === VertexModel.CLAUDE_3_SONNET || modelType === VertexModel.CLAUDE_3_OPUS) {
    return extractClaudeToolCalls(content);
  } else {
    throw new Error(`Unsupported model type for extracting tool calls: ${modelType}`);
  }
}

function extractGeminiToolCalls(content: string): StandardizedToolCall[] {
  // Implement Gemini-specific tool call extraction logic
  // This is a placeholder and should be updated based on Gemini's actual format
  return [];
}

function extractClaudeToolCalls(content: any[]): StandardizedToolCall[] {
  return content
    .filter(item => item.type === 'tool_use')
    .map(item => convertClaudeToolCall(item.tool_use));
}