1
import fetch from 'node-fetch';
2

3
/**
4
 * Gets the vector for the given text from SillyTavern-extras
5
 * @param {string[]} texts - The array of texts to get the vectors for
6
 * @param {string} apiUrl - The Extras API URL
7
 * @param {string} apiKey - The Extras API key, or empty string if API key not enabled
8
 * @returns {Promise<number[][]>} - The array of vectors for the texts
9
 */
10
export async function getExtrasBatchVector(texts, apiUrl, apiKey) {
11
    return getExtrasVectorImpl(texts, apiUrl, apiKey);
12
}
13

14
/**
15
 * Gets the vector for the given text from SillyTavern-extras
16
 * @param {string} text - The text to get the vector for
17
 * @param {string} apiUrl - The Extras API URL
18
 * @param {string} apiKey - The Extras API key, or empty string if API key not enabled
19
 * @returns {Promise<number[]>} - The vector for the text
20
 */
21
export async function getExtrasVector(text, apiUrl, apiKey) {
22
    return getExtrasVectorImpl(text, apiUrl, apiKey);
23
}
24

25
/**
26
 * Gets the vector for the given text from SillyTavern-extras
27
 * @param {string|string[]} text - The text or texts to get the vector(s) for
28
 * @param {string} apiUrl - The Extras API URL
29
 * @param {string} apiKey - The Extras API key, or empty string if API key not enabled *
30
 * @returns {Promise<Array>} - The vector for a single text if input is string, or the array of vectors for multiple texts if input is string[]
31
 */
32
async function getExtrasVectorImpl(text, apiUrl, apiKey) {
33
    let url;
34
    try {
35
        url = new URL(apiUrl);
36
        url.pathname = '/api/embeddings/compute';
37
    } catch (error) {
38
        console.error('Failed to set up Extras API call:', error);
39
        console.debug('Extras API URL given was:', apiUrl);
40
        throw error;
41
    }
42

43
    const headers = {
44
        'Content-Type': 'application/json',
45
    };
46

47
    // Include the Extras API key, if enabled
48
    if (apiKey && apiKey.length > 0) {
49
        Object.assign(headers, {
50
            'Authorization': `Bearer ${apiKey}`,
51
        });
52
    }
53

54
    const response = await fetch(url, {
55
        method: 'POST',
56
        headers: headers,
57
        body: JSON.stringify({
58
            text: text,  // The backend accepts {string|string[]} for one or multiple text items, respectively.
59
        }),
60
    });
61

62
    if (!response.ok) {
63
        const text = await response.text();
64
        console.warn('Extras request failed', response.statusText, text);
65
        throw new Error('Extras request failed');
66
    }
67

68
    /** @type {any} */
69
    const data = await response.json();
70
    const vector = data.embedding;  // `embedding`: number[] (one text item), or number[][] (multiple text items).
71

72
    return vector;
73
}