Add NomicAI for vectorization (#1922)

* Crudely add NomicAi for vectorization

* Move NomicAI to its own endpoint, properly handle API key

* Adjust clear button html

* Remove leftover nomicai http header code

* Revert changes to openai-vectors.js

* Fix UI issues

* Revert change to settings, fix UI

---------

Co-authored-by: Cohee <18619528+Cohee1207@users.noreply.github.com>

public/scripts/extensions/vectors/index.js

@@ -1,7 +1,7 @@
 import { eventSource, event_types, extension_prompt_types, getCurrentChatId, getRequestHeaders, is_send_press, saveSettingsDebounced, setExtensionPrompt, substituteParams } from '../../../script.js';
 import { ModuleWorkerWrapper, extension_settings, getContext, modules, renderExtensionTemplate } from '../../extensions.js';
 import { collapseNewlines } from '../../power-user.js';
-import { SECRET_KEYS, secret_state } from '../../secrets.js';
+import { SECRET_KEYS, secret_state, writeSecret } from '../../secrets.js';
 import { debounce, getStringHash as calculateHash, waitUntilCondition, onlyUnique, splitRecursive } from '../../utils.js';
 
 const MODULE_NAME = 'vectors';
@@ -461,7 +461,8 @@ async function insertVectorItems(collectionId, items) {
     if (settings.source === 'openai' && !secret_state[SECRET_KEYS.OPENAI] ||
         settings.source === 'palm' && !secret_state[SECRET_KEYS.MAKERSUITE] ||
         settings.source === 'mistral' && !secret_state[SECRET_KEYS.MISTRALAI] ||
-        settings.source === 'togetherai' && !secret_state[SECRET_KEYS.TOGETHERAI]) {
+        settings.source === 'togetherai' && !secret_state[SECRET_KEYS.TOGETHERAI] ||
+        settings.source === 'nomicai' && !secret_state[SECRET_KEYS.NOMICAI]) {
         throw new Error('Vectors: API key missing', { cause: 'api_key_missing' });
     }
 
@@ -580,6 +581,7 @@ function toggleSettings() {
     $('#vectors_files_settings').toggle(!!settings.enabled_files);
     $('#vectors_chats_settings').toggle(!!settings.enabled_chats);
     $('#together_vectorsModel').toggle(settings.source === 'togetherai');
+    $('#nomicai_apiKey').toggle(settings.source === 'nomicai');
 }
 
 async function onPurgeClick() {
@@ -655,7 +657,13 @@ jQuery(async () => {
         saveSettingsDebounced();
         toggleSettings();
     });
-
+    $('#api_key_nomicai').on('change', () => {
+        const nomicKey = String($('#api_key_nomicai').val()).trim();
+        if (nomicKey.length) {
+            writeSecret(SECRET_KEYS.NOMICAI, nomicKey);
+        }
+        saveSettingsDebounced();
+    });
     $('#vectors_togetherai_model').val(settings.togetherai_model).on('change', () => {
         settings.togetherai_model = String($('#vectors_togetherai_model').val());
         Object.assign(extension_settings.vectors, settings);
@@ -726,6 +734,10 @@ jQuery(async () => {
         saveSettingsDebounced();
     });
 
+    const validSecret = !!secret_state[SECRET_KEYS.NOMICAI];
+    const placeholder = validSecret ? '✔️ Key saved' : '❌ Missing key';
+    $('#api_key_nomicai').attr('placeholder', placeholder);
+
     toggleSettings();
     eventSource.on(event_types.MESSAGE_DELETED, onChatEvent);
     eventSource.on(event_types.MESSAGE_EDITED, onChatEvent);

public/scripts/extensions/vectors/settings.html

@@ -16,6 +16,7 @@
                     <option value="palm">Google MakerSuite (PaLM)</option>
                     <option value="mistral">MistralAI</option>
                     <option value="togetherai">TogetherAI</option>
+                    <option value="nomicai">NomicAI</option>
                 </select>
             </div>
             <div class="flex-container flexFlowColumn" id="together_vectorsModel">
@@ -33,6 +34,19 @@
                     <option value="bert-base-uncased">Bert Base Uncased</option>
                 </select>
             </div>
+            <div class="flex-container flexFlowColumn" id="nomicai_apiKey">
+                <label for="api_key_nomicai">
+                    <span>NomicAI API Key</span>
+                </label>
+                <div class="flex-container">
+                    <input id="api_key_nomicai" name="api_key_nomicai" class="text_pole flex1 wide100p" maxlength="500" size="35" type="text" autocomplete="off">
+                    <div title="Clear your API key" class="menu_button fa-solid fa-circle-xmark clear-api-key" data-key="api_key_nomicai">
+                    </div>
+                </div>
+                <div data-for="api_key_nomicai" class="neutral_warning" data-i18n="For privacy reasons, your API key will be hidden after you reload the page.">
+                    For privacy reasons, your API key will be hidden after you reload the page.
+                </div>
+            </div>
 
             <div class="flex-container flexFlowColumn" title="How many last messages will be matched for relevance.">
                 <label for="vectors_query">

public/scripts/secrets.js

@@ -20,6 +20,7 @@ export const SECRET_KEYS = {
     DREAMGEN: 'api_key_dreamgen',
     CUSTOM: 'api_key_custom',
     OOBA: 'api_key_ooba',
+    NOMICAI: 'api_key_nomicai',
 };
 
 const INPUT_MAP = {
@@ -41,6 +42,7 @@ const INPUT_MAP = {
     [SECRET_KEYS.OOBA]: '#api_key_ooba',
     [SECRET_KEYS.INFERMATICAI]: '#api_key_infermaticai',
     [SECRET_KEYS.DREAMGEN]: '#api_key_dreamgen',
+    [SECRET_KEYS.NOMICAI]: '#api_key_nomicai',
 };
 
 async function clearSecret() {
@@ -48,7 +50,7 @@ async function clearSecret() {
     await writeSecret(key, '');
     secret_state[key] = false;
     updateSecretDisplay();
-    $(INPUT_MAP[key]).val('');
+    $(INPUT_MAP[key]).val('').trigger('input');
     $('#main_api').trigger('change');
 }
 

src/endpoints/secrets.js

@@ -32,6 +32,7 @@ const SECRET_KEYS = {
     OOBA: 'api_key_ooba',
     INFERMATICAI: 'api_key_infermaticai',
     DREAMGEN: 'api_key_dreamgen',
+    NOMICAI: 'api_key_nomicai',
 };
 
 // These are the keys that are safe to expose, even if allowKeysExposure is false

src/endpoints/vectors.js

@@ -5,7 +5,7 @@ const sanitize = require('sanitize-filename');
 const { jsonParser } = require('../express-common');
 
 // Don't forget to add new sources to the SOURCES array
-const SOURCES = ['transformers', 'mistral', 'openai', 'extras', 'palm', 'togetherai'];
+const SOURCES = ['transformers', 'mistral', 'openai', 'extras', 'palm', 'togetherai', 'nomicai'];
 
 /**
  * Gets the vector for the given text from the given source.
@@ -16,6 +16,8 @@ const SOURCES = ['transformers', 'mistral', 'openai', 'extras', 'palm', 'togethe
  */
 async function getVector(source, sourceSettings, text) {
     switch (source) {
+        case 'nomicai':
+            return require('../nomicai-vectors').getNomicAIVector(text, source);
         case 'togetherai':
         case 'mistral':
         case 'openai':
@@ -45,6 +47,9 @@ async function getBatchVector(source, sourceSettings, texts) {
     let results = [];
     for (let batch of batches) {
         switch (source) {
+            case 'nomicai':
+                results.push(...await require('../nomicai-vectors').getNomicAIBatchVector(batch, source));
+                break;
             case 'togetherai':
             case 'mistral':
             case 'openai':

src/nomicai-vectors.js

@@ -0,0 +1,76 @@
+const fetch = require('node-fetch').default;
+const { SECRET_KEYS, readSecret } = require('./endpoints/secrets');
+
+const SOURCES = {
+    'nomicai': {
+        secretKey: SECRET_KEYS.NOMICAI,
+        url: 'api-atlas.nomic.ai/v1/embedding/text',
+        model: 'nomic-embed-text-v1.5',
+    }
+};
+
+/**
+ * Gets the vector for the given text batch from an OpenAI compatible endpoint.
+ * @param {string[]} texts - The array of texts to get the vector for
+ * @param {string} source - The source of the vector
+ * @returns {Promise<number[][]>} - The array of vectors for the texts
+ */
+async function getNomicAIBatchVector(texts, source) {
+    const config = SOURCES[source];
+
+    if (!config) {
+        console.log('Unknown source', source);
+        throw new Error('Unknown source');
+    }
+
+    const key = readSecret(config.secretKey);
+
+    if (!key) {
+        console.log('No API key found');
+        throw new Error('No API key found');
+    }
+
+    const url = config.url;
+    let response;
+    response = await fetch(`https://${url}`, {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/json',
+            Authorization: `Bearer ${key}`,
+        },
+        body: JSON.stringify({
+            texts: texts,
+            model: config.model,
+        }),
+    });
+
+    if (!response.ok) {
+        const text = await response.text();
+        console.log('API request failed', response.statusText, text);
+        throw new Error('API request failed');
+    }
+
+    const data = await response.json();
+    if (!Array.isArray(data?.embeddings)) {
+        console.log('API response was not an array');
+        throw new Error('API response was not an array');
+    }
+
+    return data.embeddings;
+}
+
+/**
+ * Gets the vector for the given text from an OpenAI compatible endpoint.
+ * @param {string} text - The text to get the vector for
+ * @param {string} source - The source of the vector
+ * @returns {Promise<number[]>} - The vector for the text
+ */
+async function getNomicAIVector(text, source) {
+    const vectors = await getNomicAIBatchVector([text], source);
+    return vectors[0];
+}
+
+module.exports = {
+    getNomicAIVector,
+    getNomicAIBatchVector,
+};