From b49097814293b11b81b5758c8459455e96a5487e Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sat, 2 Mar 2024 23:16:18 +0200
Subject: [PATCH] Refactor vector models code

---
 public/scripts/extensions/vectors/index.js    | 22 +++++--------------
 .../scripts/extensions/vectors/settings.html  |  6 ++---
 src/openai-vectors.js                         |  2 +-
 3 files changed, 10 insertions(+), 20 deletions(-)

diff --git a/public/scripts/extensions/vectors/index.js b/public/scripts/extensions/vectors/index.js
index ef13de22a..840e92687 100644
--- a/public/scripts/extensions/vectors/index.js
+++ b/public/scripts/extensions/vectors/index.js
@@ -12,7 +12,7 @@ const settings = {
     // For both
     source: 'transformers',
     include_wi: false,
-    model: 'togethercomputer/m2-bert-80M-32k-retrieval',
+    togetherai_model: 'togethercomputer/m2-bert-80M-32k-retrieval',
 
     // For chats
     enabled_chats: false,
@@ -447,11 +447,10 @@ function addExtrasHeaders(headers) {
  */
 function addTogetherAiHeaders(headers) {
     Object.assign(headers, {
-        'X-Togetherai-Model': extension_settings.vectors.model,
+        'X-Togetherai-Model': extension_settings.vectors.togetherai_model,
     });
 }
 
-
 /**
  * Inserts vector items into a collection
  * @param {string} collectionId - The collection to insert into
@@ -580,6 +579,7 @@ async function purgeVectorIndex(collectionId) {
 function toggleSettings() {
     $('#vectors_files_settings').toggle(!!settings.enabled_files);
     $('#vectors_chats_settings').toggle(!!settings.enabled_chats);
+    $('#together_vectorsModel').toggle(settings.source === 'togetherai');
 }
 
 async function onPurgeClick() {
@@ -634,12 +634,6 @@ jQuery(async () => {
 
     Object.assign(settings, extension_settings.vectors);
 
-    if (settings.source === 'togetherai') {
-        $('#vectorsModel').show();
-    } else {
-        $('#vectorsModel').hide();
-    }
-
     // Migrate from TensorFlow to Transformers
     settings.source = settings.source !== 'local' ? settings.source : 'transformers';
     $('#extensions_settings2').append(renderExtensionTemplate(MODULE_NAME, 'settings'));
@@ -657,17 +651,13 @@ jQuery(async () => {
     });
     $('#vectors_source').val(settings.source).on('change', () => {
         settings.source = String($('#vectors_source').val());
-        if (settings.source === 'togetherai') {
-            $('#vectorsModel').show();
-        } else {
-            $('#vectorsModel').hide();
-        }
         Object.assign(extension_settings.vectors, settings);
         saveSettingsDebounced();
+        toggleSettings();
     });
 
-    $('#vectors_model').val(settings.model).on('change', () => {
-        settings.model = String($('#vectors_model').val());
+    $('#vectors_togetherai_model').val(settings.togetherai_model).on('change', () => {
+        settings.togetherai_model = String($('#vectors_togetherai_model').val());
         Object.assign(extension_settings.vectors, settings);
         saveSettingsDebounced();
     });
diff --git a/public/scripts/extensions/vectors/settings.html b/public/scripts/extensions/vectors/settings.html
index fb26b14ff..5900d7145 100644
--- a/public/scripts/extensions/vectors/settings.html
+++ b/public/scripts/extensions/vectors/settings.html
@@ -18,11 +18,11 @@
                     <option value="togetherai">TogetherAI</option>
                 </select>
             </div>
-            <div class="flex-container flexFlowColumn" id="vectorsModel">
-                <label for="vectors_model">
+            <div class="flex-container flexFlowColumn" id="together_vectorsModel">
+                <label for="vectors_togetherai_model">
                     Vectorization Model
                 </label>
-                <select id="vectors_model" class="text_pole">
+                <select id="vectors_togetherai_model" class="text_pole">
                     <option value="togethercomputer/m2-bert-80M-32k-retrieval">M2-BERT-Retrieval-32k</option>
                     <option value="togethercomputer/m2-bert-80M-8k-retrieval">M2-BERT-Retrieval-8k</option>
                     <option value="togethercomputer/m2-bert-80M-2k-retrieval">M2-BERT-Retrieval-2K</option>
diff --git a/src/openai-vectors.js b/src/openai-vectors.js
index 222bc31aa..6ac486415 100644
--- a/src/openai-vectors.js
+++ b/src/openai-vectors.js
@@ -50,7 +50,7 @@ async function getOpenAIBatchVector(texts, source, model = '') {
         },
         body: JSON.stringify({
             input: texts,
-            model: source === 'togetherai' && model !== '' ? model : config.model,
+            model: model || config.model,
         }),
     });