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

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

import axios from 'axios';
import { VertexConfig } from './config';

interface VertexEmbeddingResponse {
  predictions: Array<{
    embeddings: {
      values: number[];
      statistics: {
        truncated: boolean;
        token_count: number;
      };
    };
  }>;
}

export async function embed(config: VertexConfig, text: string): Promise<number[]> {
  const url = `https://${config.location}-aiplatform.googleapis.com/v1/projects/${config.projectId}/locations/${config.location}/publishers/google/models/textembedding-gecko:predict`;

  const requestBody = {
    instances: [{ content: text }],
  };

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

  if (response.data.predictions && response.data.predictions.length > 0) {
    return response.data.predictions[0].embeddings.values;
  }

  throw new Error('Failed to get embedding from Vertex AI');
}

export async function embedBatch(config: VertexConfig, texts: string[]): Promise<number[][]> {
  const url = `https://${config.location}-aiplatform.googleapis.com/v1/projects/${config.projectId}/locations/${config.location}/publishers/google/models/textembedding-gecko:predict`;

  const requestBody = {
    instances: texts.map(text => ({ content: text })),
  };

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

  if (response.data.predictions && response.data.predictions.length > 0) {
    return response.data.predictions.map(prediction => prediction.embeddings.values);
  }

  throw new Error('Failed to get embeddings from Vertex AI');
}

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';
}