added mistral vector support (off the back of oai's)

This commit is contained in:
based
2023-12-16 08:37:39 +10:00
parent c517483141
commit 5dd2e8cd88
4 changed files with 12 additions and 7 deletions

View File

@ -2,19 +2,21 @@ const fetch = require('node-fetch').default;
const { SECRET_KEYS, readSecret } = require('./endpoints/secrets');
/**
* Gets the vector for the given text from OpenAI ada model
* Gets the vector for the given text from an OpenAI compatible endpoint.
* @param {string} text - The text to get the vector for
* @returns {Promise<number[]>} - The vector for the text
*/
async function getOpenAIVector(text) {
const key = readSecret(SECRET_KEYS.OPENAI);
async function getOpenAIVector(text, source) {
const isMistral = source === 'mistral';
const key = readSecret(isMistral ? SECRET_KEYS.MISTRALAI : SECRET_KEYS.OPENAI);
if (!key) {
console.log('No OpenAI key found');
throw new Error('No OpenAI key found');
}
const response = await fetch('https://api.openai.com/v1/embeddings', {
const url = isMistral ? 'api.mistral.ai' : 'api.openai.com';
const response = await fetch(`https://${url}/v1/embeddings`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
@ -22,7 +24,7 @@ async function getOpenAIVector(text) {
},
body: JSON.stringify({
input: text,
model: 'text-embedding-ada-002',
model: isMistral ? 'mistral-embed' : 'text-embedding-ada-002',
}),
});