Spaces:
Runtime error
Runtime error
File size: 3,005 Bytes
1307964 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
const fetch = require('node-fetch').default;
const { SECRET_KEYS, readSecret } = require('../endpoints/secrets');
const SOURCES = {
'togetherai': {
secretKey: SECRET_KEYS.TOGETHERAI,
url: 'api.together.xyz',
model: 'togethercomputer/m2-bert-80M-32k-retrieval',
},
'mistral': {
secretKey: SECRET_KEYS.MISTRALAI,
url: 'api.mistral.ai',
model: 'mistral-embed',
},
'openai': {
secretKey: SECRET_KEYS.OPENAI,
url: 'api.openai.com',
model: 'text-embedding-ada-002',
},
};
/**
* Gets the vector for the given text batch from an OpenAI compatible endpoint.
* @param {string[]} texts - The array of texts to get the vector for
* @param {string} source - The source of the vector
* @param {import('../users').UserDirectoryList} directories - The directories object for the user
* @param {string} model - The model to use for the embedding
* @returns {Promise<number[][]>} - The array of vectors for the texts
*/
async function getOpenAIBatchVector(texts, source, directories, model = '') {
const config = SOURCES[source];
if (!config) {
console.log('Unknown source', source);
throw new Error('Unknown source');
}
const key = readSecret(directories, config.secretKey);
if (!key) {
console.log('No API key found');
throw new Error('No API key found');
}
const url = config.url;
const response = await fetch(`https://${url}/v1/embeddings`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${key}`,
},
body: JSON.stringify({
input: texts,
model: model || config.model,
}),
});
if (!response.ok) {
const text = await response.text();
console.log('API request failed', response.statusText, text);
throw new Error('API request failed');
}
const data = await response.json();
if (!Array.isArray(data?.data)) {
console.log('API response was not an array');
throw new Error('API response was not an array');
}
// Sort data by x.index to ensure the order is correct
data.data.sort((a, b) => a.index - b.index);
const vectors = data.data.map(x => x.embedding);
return vectors;
}
/**
* Gets the vector for the given text from an OpenAI compatible endpoint.
* @param {string} text - The text to get the vector for
* @param {string} source - The source of the vector
* @param {import('../users').UserDirectoryList} directories - The directories object for the user
* @param {string} model - The model to use for the embedding
* @returns {Promise<number[]>} - The vector for the text
*/
async function getOpenAIVector(text, source, directories, model = '') {
const vectors = await getOpenAIBatchVector([text], source, directories, model);
return vectors[0];
}
module.exports = {
getOpenAIVector,
getOpenAIBatchVector,
};
|