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

View File

@ -19,6 +19,16 @@
<option value="nomicai">NomicAI</option> <option value="nomicai">NomicAI</option>
</select> </select>
</div> </div>
<div class="flex-container flexFlowColumn" id="openai_vectorsModel">
<label for="vectors_openai_model">
Vectorization Model
</label>
<select id="vectors_openai_model" class="text_pole">
<option value="text-embedding-ada-002">text-embedding-ada-002</option>
<option value="text-embedding-3-small">text-embedding-3-small</option>
<option value="text-embedding-3-large">text-embedding-3-large</option>
</select>
</div>
<div class="flex-container flexFlowColumn" id="together_vectorsModel"> <div class="flex-container flexFlowColumn" id="together_vectorsModel">
<label for="vectors_togetherai_model"> <label for="vectors_togetherai_model">
Vectorization Model Vectorization Model
@ -34,6 +44,14 @@
<option value="bert-base-uncased">Bert Base Uncased</option> <option value="bert-base-uncased">Bert Base Uncased</option>
</select> </select>
</div> </div>
<small id="vectors_modelWarning">
<i class="fa-solid fa-exclamation-triangle"></i>
<span data-i18n="Vectors Model Warning">
It is recommended to purge vectors when changing the model mid-chat. Otherwise, it will lead to sub-par results.
</span>
</small>
<div class="flex-container flexFlowColumn" id="nomicai_apiKey"> <div class="flex-container flexFlowColumn" id="nomicai_apiKey">
<label for="api_key_nomicai"> <label for="api_key_nomicai">
<span>NomicAI API Key</span> <span>NomicAI API Key</span>

View File

@ -175,6 +175,12 @@ function getSourceSettings(source, request) {
if (source === 'togetherai') { if (source === 'togetherai') {
let model = String(request.headers['x-togetherai-model']); let model = String(request.headers['x-togetherai-model']);
return {
model: model,
};
} else if (source === 'openai') {
let model = String(request.headers['x-openai-model']);
return { return {
model: model, model: model,
}; };