Refactor openai vectors

This commit is contained in:
Cohee 2023-12-17 02:56:47 +02:00
parent b1f07eb989
commit 2d8a62d059
1 changed files with 31 additions and 11 deletions

View File

@ -1,21 +1,41 @@
const fetch = require('node-fetch').default; const fetch = require('node-fetch').default;
const { SECRET_KEYS, readSecret } = require('./endpoints/secrets'); const { SECRET_KEYS, readSecret } = require('./endpoints/secrets');
const SOURCES = {
'mistral': {
secretKey: SECRET_KEYS.MISTRAL,
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 from an OpenAI compatible endpoint. * Gets the vector for the given text from an OpenAI compatible endpoint.
* @param {string} text - The text to get the vector for * @param {string} text - The text to get the vector for
* @param {string} source - The source of the vector
* @returns {Promise<number[]>} - The vector for the text * @returns {Promise<number[]>} - The vector for the text
*/ */
async function getOpenAIVector(text, source) { async function getOpenAIVector(text, source) {
const isMistral = source === 'mistral'; const config = SOURCES[source];
const key = readSecret(isMistral ? SECRET_KEYS.MISTRALAI : SECRET_KEYS.OPENAI);
if (!key) { if (!config) {
console.log('No OpenAI key found'); console.log('Unknown source', source);
throw new Error('No OpenAI key found'); throw new Error('Unknown source');
} }
const url = isMistral ? 'api.mistral.ai' : 'api.openai.com'; const key = readSecret(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`, { const response = await fetch(`https://${url}/v1/embeddings`, {
method: 'POST', method: 'POST',
headers: { headers: {
@ -24,22 +44,22 @@ async function getOpenAIVector(text, source) {
}, },
body: JSON.stringify({ body: JSON.stringify({
input: text, input: text,
model: isMistral ? 'mistral-embed' : 'text-embedding-ada-002', model: config.model,
}), }),
}); });
if (!response.ok) { if (!response.ok) {
const text = await response.text(); const text = await response.text();
console.log('OpenAI request failed', response.statusText, text); console.log('API request failed', response.statusText, text);
throw new Error('OpenAI request failed'); throw new Error('API request failed');
} }
const data = await response.json(); const data = await response.json();
const vector = data?.data[0]?.embedding; const vector = data?.data[0]?.embedding;
if (!Array.isArray(vector)) { if (!Array.isArray(vector)) {
console.log('OpenAI response was not an array'); console.log('API response was not an array');
throw new Error('OpenAI response was not an array'); throw new Error('API response was not an array');
} }
return vector; return vector;