diff --git a/src/openai-vectors.js b/src/openai-vectors.js index 40c54ae2f..3b19a4c96 100644 --- a/src/openai-vectors.js +++ b/src/openai-vectors.js @@ -1,21 +1,41 @@ const fetch = require('node-fetch').default; 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. * @param {string} text - The text to get the vector for + * @param {string} source - The source of the vector * @returns {Promise} - The vector for the text */ async function getOpenAIVector(text, source) { - const isMistral = source === 'mistral'; - const key = readSecret(isMistral ? SECRET_KEYS.MISTRALAI : SECRET_KEYS.OPENAI); + const config = SOURCES[source]; - if (!key) { - console.log('No OpenAI key found'); - throw new Error('No OpenAI key found'); + if (!config) { + console.log('Unknown source', source); + 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`, { method: 'POST', headers: { @@ -24,22 +44,22 @@ async function getOpenAIVector(text, source) { }, body: JSON.stringify({ input: text, - model: isMistral ? 'mistral-embed' : 'text-embedding-ada-002', + model: config.model, }), }); if (!response.ok) { const text = await response.text(); - console.log('OpenAI request failed', response.statusText, text); - throw new Error('OpenAI request failed'); + console.log('API request failed', response.statusText, text); + throw new Error('API request failed'); } const data = await response.json(); const vector = data?.data[0]?.embedding; if (!Array.isArray(vector)) { - console.log('OpenAI response was not an array'); - throw new Error('OpenAI response was not an array'); + console.log('API response was not an array'); + throw new Error('API response was not an array'); } return vector;