1
import fetch from 'node-fetch';
2
import { SECRET_KEYS, readSecret } from '../endpoints/secrets.js';
3

4
/**
5
 * Gets the vector for the given text batch from an OpenAI compatible endpoint.
6
 * @param {string[]} texts - The array of texts to get the vector for
7
 * @param {boolean} isQuery - If the text is a query for embedding search
8
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
9
 * @param {string} model - The model to use for the embedding
10
 * @returns {Promise<number[][]>} - The array of vectors for the texts
11
 */
12
export async function getCohereBatchVector(texts, isQuery, directories, model) {
13
    const key = readSecret(directories, SECRET_KEYS.COHERE);
14

15
    if (!key) {
16
        console.warn('No API key found');
17
        throw new Error('No API key found');
18
    }
19

20
    const response = await fetch('https://api.cohere.ai/v2/embed', {
21
        method: 'POST',
22
        headers: {
23
            'Content-Type': 'application/json',
24
            Authorization: `Bearer ${key}`,
25
        },
26
        body: JSON.stringify({
27
            texts: texts,
28
            model: model,
29
            embedding_types: ['float'],
30
            input_type: isQuery ? 'search_query' : 'search_document',
31
            truncate: 'END',
32
        }),
33
    });
34

35
    if (!response.ok) {
36
        const text = await response.text();
37
        console.warn('API request failed', response.statusText, text);
38
        throw new Error('API request failed');
39
    }
40

41
    /** @type {any} */
42
    const data = await response.json();
43
    if (!Array.isArray(data?.embeddings?.float)) {
44
        console.warn('API response was not an array');
45
        throw new Error('API response was not an array');
46
    }
47

48
    return data.embeddings.float;
49
}
50

51
/**
52
 * Gets the vector for the given text from an OpenAI compatible endpoint.
53
 * @param {string} text - The text to get the vector for
54
 * @param {boolean} isQuery - If the text is a query for embedding search
55
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
56
 * @param {string} model - The model to use for the embedding
57
 * @returns {Promise<number[]>} - The vector for the text
58
 */
59
export async function getCohereVector(text, isQuery, directories, model) {
60
    const vectors = await getCohereBatchVector([text], isQuery, directories, model);
61
    return vectors[0];
62
}
63