Refactor openai vectors
This commit is contained in:
parent
b1f07eb989
commit
2d8a62d059
|
@ -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<number[]>} - 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;
|
||||
|
|
Loading…
Reference in New Issue