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

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

import axios from 'axios';
import { ChatMessage, ChatOptions, ProviderResponse } from '../../../types/chat';
import { Tool } from '../../../types/tool';
import { VertexConfig } from './config';
import { VertexChatRequestBody, VertexChatResponse, VertexModel } from './types';
import { convertToolSchema, convertMessages, extractToolCalls } from './tools';

export async function chat(
  config: VertexConfig,
  messages: ChatMessage[],
  options: ChatOptions,
  tools?: Tool[]
): Promise<ProviderResponse> {
  const url = config.getApiUrl();
  const modelType = config.getModelType();
  
  const requestBody: VertexChatRequestBody = {
    instances: [
      {
        messages: convertMessages(messages, modelType),
      },
    ],
    parameters: {
      temperature: options.temperature || 0.7,
      maxOutputTokens: options.maxTokens || 1024,
      topP: 0.95,
      topK: 40,
    },
  };

  if (tools && tools.length > 0) {
    requestBody.instances[0].context = `You have access to the following tools: ${tools.map(tool => tool.name).join(', ')}. Use them when necessary.`;
  }

  const response = await axios.post<VertexChatResponse>(url, requestBody, {
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${await getAccessToken(config)}`,
    },
  });

  return createProviderResponse(response.data, modelType);
}

function createProviderResponse(data: VertexChatResponse, modelType: VertexModel): ProviderResponse {
  const response: ProviderResponse = {
    content: '',
    toolCalls: [],
    usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, // Vertex AI doesn't provide token usage
  };

  if (data.predictions && data.predictions.length > 0) {
    const prediction = data.predictions[0];
    if (prediction.candidates && prediction.candidates.length > 0) {
      response.content = prediction.candidates[0].content;
    }
  }

  response.toolCalls = extractToolCalls(response.content, modelType);

  return response;
}

async function getAccessToken(config: VertexConfig): Promise<string> {
  // Implement logic to get access token from Google Cloud credentials
  // This might involve using the Google Auth Library or a similar mechanism
  // For now, we'll return a placeholder
  return 'placeholder_access_token';
}