File size: 3,646 Bytes
c120a1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import fetch from 'node-fetch';
import { getGoogleApiConfig } from '../endpoints/google.js';
/**
* Gets the vector for the given text from Google AI Studio
* @param {string[]} texts - The array of texts to get the vector for
* @param {string} model - The model to use for embedding
* @param {import('express').Request} request - The request object to get API key and URL
* @returns {Promise<number[][]>} - The array of vectors for the texts
*/
export async function getMakerSuiteBatchVector(texts, model, request) {
const { url, headers, apiName } = await getGoogleApiConfig(request, model, 'batchEmbedContents');
const body = {
requests: texts.map(text => ({
model: `models/${model}`,
content: { parts: [{ text }] },
})),
};
const response = await fetch(url, {
body: JSON.stringify(body),
method: 'POST',
headers: headers,
});
if (!response.ok) {
const text = await response.text();
console.warn(`${apiName} batch request failed`, response.statusText, text);
throw new Error(`${apiName} batch request failed`);
}
/** @type {any} */
const data = await response.json();
if (!Array.isArray(data?.embeddings)) {
throw new Error(`${apiName} did not return an array`);
}
const embeddings = data.embeddings.map(embedding => embedding.values);
return embeddings;
}
/**
* Gets the vector for the given text from Google Vertex AI
* @param {string[]} texts - The array of texts to get the vector for
* @param {string} model - The model to use for embedding
* @param {import('express').Request} request - The request object to get API key and URL
* @returns {Promise<number[][]>} - The array of vectors for the texts
*/
export async function getVertexBatchVector(texts, model, request) {
const { url, headers, apiName } = await getGoogleApiConfig(request, model, 'predict');
const body = {
instances: texts.map(text => ({ content: text })),
};
const response = await fetch(url, {
body: JSON.stringify(body),
method: 'POST',
headers: headers,
});
if (!response.ok) {
const text = await response.text();
console.warn(`${apiName} batch request failed`, response.statusText, text);
throw new Error(`${apiName} batch request failed`);
}
/** @type {any} */
const data = await response.json();
if (!Array.isArray(data?.predictions)) {
throw new Error(`${apiName} did not return an array`);
}
const embeddings = data.predictions.map(p => p.embeddings.values);
return embeddings;
}
/**
* Gets the vector for the given text from Google AI Studio
* @param {string} text - The text to get the vector for
* @param {string} model - The model to use for embedding
* @param {import('express').Request} request - The request object to get API key and URL
* @returns {Promise<number[]>} - The vector for the text
*/
export async function getMakerSuiteVector(text, model, request) {
const [embedding] = await getMakerSuiteBatchVector([text], model, request);
return embedding;
}
/**
* Gets the vector for the given text from Google Vertex AI
* @param {string} text - The text to get the vector for
* @param {string} model - The model to use for embedding
* @param {import('express').Request} request - The request object to get API key and URL
* @returns {Promise<number[]>} - The vector for the text
*/
export async function getVertexVector(text, model, request) {
const [embedding] = await getVertexBatchVector([text], model, request);
return embedding;
}
|