1
import fetch from 'node-fetch';
2
import urlJoin from 'url-join';
3
import { setAdditionalHeadersByType } from '../additional-headers.js';
4
import { TEXTGEN_TYPES } from '../constants.js';
5
import { trimV1 } from '../util.js';
6

7
/**
8
 * Gets the vector for the given text from LlamaCpp
9
 * @param {string[]} texts - The array of texts to get the vectors for
10
 * @param {string} apiUrl - The API URL
11
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
12
 * @returns {Promise<number[][]>} - The array of vectors for the texts
13
 */
14
export async function getLlamaCppBatchVector(texts, apiUrl, directories) {
15
    const url = new URL(urlJoin(trimV1(apiUrl), '/v1/embeddings'));
16

17
    const headers = {};
18
    setAdditionalHeadersByType(headers, TEXTGEN_TYPES.LLAMACPP, apiUrl, directories);
19

20
    const response = await fetch(url, {
21
        method: 'POST',
22
        headers: {
23
            'Content-Type': 'application/json',
24
            ...headers,
25
        },
26
        body: JSON.stringify({ input: texts }),
27
    });
28

29
    if (!response.ok) {
30
        const responseText = await response.text();
31
        throw new Error(`LlamaCpp: Failed to get vector for text: ${response.statusText} ${responseText}`);
32
    }
33

34
    /** @type {any} */
35
    const data = await response.json();
36

37
    if (!Array.isArray(data?.data)) {
38
        throw new Error('API response was not an array');
39
    }
40

41
    // Sort data by x.index to ensure the order is correct
42
    data.data.sort((a, b) => a.index - b.index);
43

44
    const vectors = data.data.map(x => x.embedding);
45
    return vectors;
46
}
47

48
/**
49
 * Gets the vector for the given text from LlamaCpp
50
 * @param {string} text - The text to get the vector for
51
 * @param {string} apiUrl - The API URL
52
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
53
 * @returns {Promise<number[]>} - The vector for the text
54
 */
55
export async function getLlamaCppVector(text, apiUrl, directories) {
56
    const vectors = await getLlamaCppBatchVector([text], apiUrl, directories);
57
    return vectors[0];
58
}