From 5efdecc825801760de68546821fa2b373985cf64 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Tue, 12 Mar 2024 21:41:30 +0200
Subject: [PATCH] Add new OpenAI embedding models for selection

---
 public/scripts/extensions/vectors/index.js    | 54 ++++++++++++++-----
 .../scripts/extensions/vectors/settings.html  | 18 +++++++
 src/endpoints/vectors.js                      |  6 +++
 3 files changed, 65 insertions(+), 13 deletions(-)

diff --git a/public/scripts/extensions/vectors/index.js b/public/scripts/extensions/vectors/index.js
index d2eab636a..3dda4a1f8 100644
--- a/public/scripts/extensions/vectors/index.js
+++ b/public/scripts/extensions/vectors/index.js
@@ -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);
diff --git a/public/scripts/extensions/vectors/settings.html b/public/scripts/extensions/vectors/settings.html
index 6bdb774be..c0aa0964a 100644
--- a/public/scripts/extensions/vectors/settings.html
+++ b/public/scripts/extensions/vectors/settings.html
@@ -19,6 +19,16 @@
                     <option value="nomicai">NomicAI</option>
                 </select>
             </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">
                 <label for="vectors_togetherai_model">
                     Vectorization Model
@@ -34,6 +44,14 @@
                     <option value="bert-base-uncased">Bert Base Uncased</option>
                 </select>
             </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">
                 <label for="api_key_nomicai">
                     <span>NomicAI API Key</span>
diff --git a/src/endpoints/vectors.js b/src/endpoints/vectors.js
index 4826d56dd..7bd38372d 100644
--- a/src/endpoints/vectors.js
+++ b/src/endpoints/vectors.js
@@ -175,6 +175,12 @@ function getSourceSettings(source, request) {
     if (source === 'togetherai') {
         let model = String(request.headers['x-togetherai-model']);
 
+        return {
+            model: model,
+        };
+    } else if (source === 'openai') {
+        let model = String(request.headers['x-openai-model']);
+
         return {
             model: model,
         };