Add new OpenAI embedding models for selection

This commit is contained in:
Cohee
2024-03-12 21:41:30 +02:00
parent 44a7dd3d74
commit 5efdecc825
3 changed files with 65 additions and 13 deletions

View File

@ -13,6 +13,7 @@ const settings = {
source: 'transformers',
include_wi: false,
togetherai_model: 'togethercomputer/m2-bert-80M-32k-retrieval',
openai_model: 'text-embedding-ada-002',
// For chats
enabled_chats: false,
@ -429,6 +430,24 @@ async function getSavedHashes(collectionId) {
return hashes;
}
function getVectorHeaders() {
const headers = getRequestHeaders();
switch (settings.source) {
case 'extras':
addExtrasHeaders(headers);
break;
case 'togetherai':
addTogetherAiHeaders(headers);
break;
case 'openai':
addOpenAiHeaders(headers);
break;
default:
break;
}
return headers;
}
/**
* Add headers for the Extras API source.
* @param {object} headers Headers object
@ -442,7 +461,7 @@ function addExtrasHeaders(headers) {
}
/**
* Add headers for the Extras API source.
* Add headers for the TogetherAI API source.
* @param {object} headers Headers object
*/
function addTogetherAiHeaders(headers) {
@ -451,6 +470,16 @@ function addTogetherAiHeaders(headers) {
});
}
/**
* Add headers for the OpenAI API source.
* @param {object} headers Header object
*/
function addOpenAiHeaders(headers) {
Object.assign(headers, {
'X-OpenAI-Model': extension_settings.vectors.openai_model,
});
}
/**
* Inserts vector items into a collection
* @param {string} collectionId - The collection to insert into
@ -470,12 +499,7 @@ async function insertVectorItems(collectionId, items) {
throw new Error('Vectors: Embeddings module missing', { cause: 'extras_module_missing' });
}
const headers = getRequestHeaders();
if (settings.source === 'extras') {
addExtrasHeaders(headers);
} else if (settings.source === 'togetherai') {
addTogetherAiHeaders(headers);
}
const headers = getVectorHeaders();
const response = await fetch('/api/vector/insert', {
method: 'POST',
@ -521,12 +545,7 @@ async function deleteVectorItems(collectionId, hashes) {
* @returns {Promise<{ hashes: number[], metadata: object[]}>} - Hashes of the results
*/
async function queryCollection(collectionId, searchText, topK) {
const headers = getRequestHeaders();
if (settings.source === 'extras') {
addExtrasHeaders(headers);
} else if (settings.source === 'togetherai') {
addTogetherAiHeaders(headers);
}
const headers = getVectorHeaders();
const response = await fetch('/api/vector/query', {
method: 'POST',
@ -581,6 +600,7 @@ function toggleSettings() {
$('#vectors_files_settings').toggle(!!settings.enabled_files);
$('#vectors_chats_settings').toggle(!!settings.enabled_chats);
$('#together_vectorsModel').toggle(settings.source === 'togetherai');
$('#openai_vectorsModel').toggle(settings.source === 'openai');
$('#nomicai_apiKey').toggle(settings.source === 'nomicai');
}
@ -645,6 +665,7 @@ jQuery(async () => {
saveSettingsDebounced();
toggleSettings();
});
$('#vectors_modelWarning').hide();
$('#vectors_enabled_files').prop('checked', settings.enabled_files).on('input', () => {
settings.enabled_files = $('#vectors_enabled_files').prop('checked');
Object.assign(extension_settings.vectors, settings);
@ -665,10 +686,17 @@ jQuery(async () => {
saveSettingsDebounced();
});
$('#vectors_togetherai_model').val(settings.togetherai_model).on('change', () => {
$('#vectors_modelWarning').show();
settings.togetherai_model = String($('#vectors_togetherai_model').val());
Object.assign(extension_settings.vectors, settings);
saveSettingsDebounced();
});
$('#vectors_openai_model').val(settings.openai_model).on('change', () => {
$('#vectors_modelWarning').show();
settings.openai_model = String($('#vectors_openai_model').val());
Object.assign(extension_settings.vectors, settings);
saveSettingsDebounced();
});
$('#vectors_template').val(settings.template).on('input', () => {
settings.template = String($('#vectors_template').val());
Object.assign(extension_settings.vectors, settings);