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 VLLM
9
 * @param {string[]} texts - The array of texts to get the vectors for
10
 * @param {string} apiUrl - The API URL
11
 * @param {string} model - The model to use
12
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
13
 * @returns {Promise<number[][]>} - The array of vectors for the texts
14
 */
15
export async function getVllmBatchVector(texts, apiUrl, model, directories) {
16
    const url = new URL(urlJoin(trimV1(apiUrl), '/v1/embeddings'));
17

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

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

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

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

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

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

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

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