1
import fetch from 'node-fetch';
2
import { getGoogleApiConfig } from '../endpoints/google.js';
3

4
/**
5
 * Gets the vector for the given text from Google AI Studio
6
 * @param {string[]} texts - The array of texts to get the vector for
7
 * @param {string} model - The model to use for embedding
8
 * @param {import('express').Request} request - The request object to get API key and URL
9
 * @returns {Promise<number[][]>} - The array of vectors for the texts
10
 */
11
export async function getMakerSuiteBatchVector(texts, model, request) {
12
    const { url, headers, apiName } = await getGoogleApiConfig(request, model, 'batchEmbedContents');
13

14
    const body = {
15
        requests: texts.map(text => ({
16
            model: `models/${model}`,
17
            content: { parts: [{ text }] },
18
        })),
19
    };
20

21
    const response = await fetch(url, {
22
        body: JSON.stringify(body),
23
        method: 'POST',
24
        headers: headers,
25
    });
26

27
    if (!response.ok) {
28
        const text = await response.text();
29
        console.warn(`${apiName} batch request failed`, response.statusText, text);
30
        throw new Error(`${apiName} batch request failed`);
31
    }
32

33
    /** @type {any} */
34
    const data = await response.json();
35
    if (!Array.isArray(data?.embeddings)) {
36
        throw new Error(`${apiName} did not return an array`);
37
    }
38

39
    const embeddings = data.embeddings.map(embedding => embedding.values);
40
    return embeddings;
41
}
42

43
/**
44
 * Gets the vector for the given text from Google Vertex AI
45
 * @param {string[]} texts - The array of texts to get the vector for
46
 * @param {string} model - The model to use for embedding
47
 * @param {import('express').Request} request - The request object to get API key and URL
48
 * @returns {Promise<number[][]>} - The array of vectors for the texts
49
 */
50
export async function getVertexBatchVector(texts, model, request) {
51
    const { url, headers, apiName } = await getGoogleApiConfig(request, model, 'predict');
52

53
    const body = {
54
        instances: texts.map(text => ({ content: text })),
55
    };
56

57
    const response = await fetch(url, {
58
        body: JSON.stringify(body),
59
        method: 'POST',
60
        headers: headers,
61
    });
62

63
    if (!response.ok) {
64
        const text = await response.text();
65
        console.warn(`${apiName} batch request failed`, response.statusText, text);
66
        throw new Error(`${apiName} batch request failed`);
67
    }
68

69
    /** @type {any} */
70
    const data = await response.json();
71
    if (!Array.isArray(data?.predictions)) {
72
        throw new Error(`${apiName} did not return an array`);
73
    }
74

75
    const embeddings = data.predictions.map(p => p.embeddings.values);
76
    return embeddings;
77
}
78

79
/**
80
 * Gets the vector for the given text from Google AI Studio
81
 * @param {string} text - The text to get the vector for
82
 * @param {string} model - The model to use for embedding
83
 * @param {import('express').Request} request - The request object to get API key and URL
84
 * @returns {Promise<number[]>} - The vector for the text
85
 */
86
export async function getMakerSuiteVector(text, model, request) {
87
    const [embedding] = await getMakerSuiteBatchVector([text], model, request);
88
    return embedding;
89
}
90

91
/**
92
 * Gets the vector for the given text from Google Vertex AI
93
 * @param {string} text - The text to get the vector for
94
 * @param {string} model - The model to use for embedding
95
 * @param {import('express').Request} request - The request object to get API key and URL
96
 * @returns {Promise<number[]>} - The vector for the text
97
 */
98
export async function getVertexVector(text, model, request) {
99
    const [embedding] = await getVertexBatchVector([text], model, request);
100
    return embedding;
101
}