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

5
const SOURCES = {
6
    'togetherai': {
7
        secretKey: SECRET_KEYS.TOGETHERAI,
8
        url: 'https://api.together.xyz/v1',
9
        model: 'togethercomputer/m2-bert-80M-32k-retrieval',
10
        headers: {},
11
        processBody: () => {},
12
    },
13
    'mistral': {
14
        secretKey: SECRET_KEYS.MISTRALAI,
15
        url: 'https://api.mistral.ai/v1',
16
        model: 'mistral-embed',
17
        headers: {},
18
        processBody: () => {},
19
    },
20
    'openai': {
21
        secretKey: SECRET_KEYS.OPENAI,
22
        url: 'https://api.openai.com/v1',
23
        model: 'text-embedding-ada-002',
24
        headers: {},
25
        processBody: () => {},
26
    },
27
    'electronhub': {
28
        secretKey: SECRET_KEYS.ELECTRONHUB,
29
        url: 'https://api.electronhub.ai/v1',
30
        model: 'text-embedding-3-small',
31
        headers: {},
32
        processBody: () => {},
33
    },
34
    'openrouter': {
35
        secretKey: SECRET_KEYS.OPENROUTER,
36
        url: 'https://openrouter.ai/api/v1',
37
        model: 'openai/text-embedding-3-large',
38
        headers: { ...OPENROUTER_HEADERS },
39
        processBody: () => {},
40
    },
41
    'chutes': {
42
        secretKey: SECRET_KEYS.CHUTES,
43
        url: 'https://{{MODEL}}.chutes.ai/v1',
44
        model: 'chutes-qwen-qwen3-embedding-8b',
45
        headers: {},
46
        processBody: (body) => {
47
            body.model = null;
48
        },
49
    },
50
    'nanogpt': {
51
        secretKey: SECRET_KEYS.NANOGPT,
52
        url: 'https://nano-gpt.com/api/v1',
53
        model: 'text-embedding-3-small',
54
        headers: {},
55
        processBody: () => {},
56
    },
57
    'siliconflow': {
58
        secretKey: SECRET_KEYS.SILICONFLOW,
59
        url: 'https://api.siliconflow.com/v1',
60
        model: 'Qwen/Qwen3-Embedding-0.6B',
61
        headers: {},
62
        processBody: () => {},
63
    },
64
    'workers_ai': {
65
        secretKey: SECRET_KEYS.WORKERS_AI,
66
        url: '', // Constructed at runtime from account ID via urlOverride
67
        model: '@cf/baai/bge-m3',
68
        headers: {},
69
        processBody: () => {},
70
    },
71
};
72

73
/**
74
 * Gets the vector for the given text batch from an OpenAI compatible endpoint.
75
 * @param {string[]} texts - The array of texts to get the vector for
76
 * @param {string} source - The source of the vector
77
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
78
 * @param {string} model - The model to use for the embedding
79
 * @param {string|null} urlOverride - Optional URL override for the API endpoint
80
 * @returns {Promise<number[][]>} - The array of vectors for the texts
81
 */
82
export async function getOpenAIBatchVector(texts, source, directories, model = '', urlOverride = null) {
83
    const config = SOURCES[source];
84

85
    if (!config) {
86
        console.error('Unknown source', source);
87
        throw new Error('Unknown source');
88
    }
89

90
    const key = readSecret(directories, config.secretKey);
91

92
    if (!key) {
93
        console.warn('No API key found');
94
        throw new Error('No API key found');
95
    }
96

97
    const modelName = model || config.model;
98
    const url = urlOverride || config.url?.replace('{{MODEL}}', modelName);
99

100
    if (!url) {
101
        throw new Error(`No API URL configured for source ${source}`);
102
    }
103

104
    const body = {
105
        input: texts,
106
        model: modelName,
107
    };
108

109
    if (typeof config.processBody === 'function') {
110
        config.processBody(body);
111
    }
112

113
    const response = await fetch(`${url}/embeddings`, {
114
        method: 'POST',
115
        headers: {
116
            'Content-Type': 'application/json',
117
            'Authorization': `Bearer ${key}`,
118
            ...config.headers,
119
        },
120
        body: JSON.stringify(body),
121
    });
122

123
    if (!response.ok) {
124
        const text = await response.text();
125
        console.warn('API request failed', response.statusText, text);
126
        throw new Error('API request failed');
127
    }
128

129
    /** @type {any} */
130
    const data = await response.json();
131

132
    if (!Array.isArray(data?.data)) {
133
        console.warn('API response was not an array');
134
        throw new Error('API response was not an array');
135
    }
136

137
    // Sort data by x.index to ensure the order is correct
138
    data.data.sort((a, b) => a.index - b.index);
139

140
    const vectors = data.data.map(x => x.embedding);
141
    return vectors;
142
}
143

144
/**
145
 * Gets the vector for the given text from an OpenAI compatible endpoint.
146
 * @param {string} text - The text to get the vector for
147
 * @param {string} source - The source of the vector
148
 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
149
 * @param {string} model - The model to use for the embedding
150
 * @param {string|null} urlOverride - Optional URL override for the API endpoint
151
 * @returns {Promise<number[]>} - The vector for the text
152
 */
153
export async function getOpenAIVector(text, source, directories, model = '', urlOverride = null) {
154
    const vectors = await getOpenAIBatchVector([text], source, directories, model, urlOverride);
155
    return vectors[0];
156
}