1
import fetch from 'node-fetch';
2
import { setAdditionalHeadersByType } from '../additional-headers.js';
3
import { TEXTGEN_TYPES } from '../constants.js';
4

5
/**
6
 * Gets the vector for the given text from Ollama
7
 * @param {string[]} texts - The array of texts to get the vectors for
8
 * @param {string} apiUrl - The API URL
9
 * @param {string} model - The model to use
10
 * @param {boolean} keep - Keep the model loaded in memory
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 getOllamaBatchVector(texts, apiUrl, model, keep, directories) {
15
    const url = new URL(apiUrl);
16
    url.pathname = '/api/embed';
17

18
    const headers = {};
19
    setAdditionalHeadersByType(headers, TEXTGEN_TYPES.OLLAMA, 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({
28
            input: texts,
29
            model: model,
30
            keep_alive: keep ? -1 : undefined,
31
            truncate: true,
32
        }),
33
    });
34

35
    if (!response.ok) {
36
        const responseText = await response.text();
37
        throw new Error(`Ollama: Failed to get batch vectors: ${response.statusText} ${responseText}`);
38
    }
39

40
    /** @type {any} */
41
    const data = await response.json();
42

43
    if (!Array.isArray(data?.embeddings)) {
44
        throw new Error('API response was not an array');
45
    }
46

47
    return data.embeddings;
48
}
49

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