[FIXED] Add PaLM vectorization source (#1182)

* Add PaLM vectorization source

* FIX

* Add API key check. Fix comment

---------

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

public/scripts/extensions/vectors/index.js

@@ -1,6 +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, renderExtensionTemplate } from "../../extensions.js";
 import { collapseNewlines, power_user, ui_mode } from "../../power-user.js";
+import { SECRET_KEYS, secret_state } from "../../secrets.js";
 import { debounce, getStringHash as calculateHash, waitUntilCondition, onlyUnique } from "../../utils.js";
 
 const MODULE_NAME = 'vectors';
@@ -116,8 +117,9 @@ async function synchronizeChat(batchSize = 5) {
 
         return newVectorItems.length - batchSize;
     } catch (error) {
-        toastr.error('Check server console for more details', 'Vectorization failed');
         console.error('Vectors: Failed to synchronize chat', error);
+        const message = error.cause === 'api_key_missing' ? 'API key missing. Save it in the "API Connections" panel.' : 'Check server console for more details';
+        toastr.error(message, 'Vectorization failed');
         return -1;
     } finally {
         syncBlocked = false;
@@ -287,6 +289,11 @@ async function getSavedHashes(collectionId) {
  * @returns {Promise<void>}
  */
 async function insertVectorItems(collectionId, items) {
+    if (settings.source === 'openai' && !secret_state[SECRET_KEYS.OPENAI] ||
+        settings.source === 'palm' && !secret_state[SECRET_KEYS.PALM]) {
+        throw new Error('Vectors: API key missing', { cause: 'api_key_missing' });
+    }
+
     const response = await fetch('/api/vector/insert', {
         method: 'POST',
         headers: getRequestHeaders(),

public/scripts/extensions/vectors/settings.html

@@ -15,6 +15,7 @@
             <select id="vectors_source" class="select">
                 <option value="transformers">Local (Transformers)</option>
                 <option value="openai">OpenAI</option>
+                <option value="palm">Google MakerSuite (PaLM)</option>
             </select>
             <div id="vectors_advanced_settings" data-newbie-hidden>
                 <label for="vectors_template">

src/palm-vectors.js

@@ -0,0 +1,43 @@
+const fetch = require('node-fetch').default;
+const { SECRET_KEYS, readSecret } = require('./secrets');
+
+/**
+ * Gets the vector for the given text from PaLM gecko model
+ * @param {string} text - The text to get the vector for
+ * @returns {Promise<number[]>} - The vector for the text
+ */
+async function getPaLMVector(text) {
+    const key = readSecret(SECRET_KEYS.PALM);
+
+    if (!key) {
+        console.log('No PaLM key found');
+        throw new Error('No PaLM key found');
+    }
+
+    const response = await fetch(`https://generativelanguage.googleapis.com/v1beta2/models/embedding-gecko-001:embedText?key=${key}`, {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/json',
+        },
+        body: JSON.stringify({
+            text: text,
+        })
+    });
+
+    if (!response.ok) {
+        const text = await response.text();
+        console.log('PaLM request failed', response.statusText, text);
+        throw new Error('PaLM request failed');
+    }
+
+    const data = await response.json();
+
+    // Access the "value" dictionary
+    const vector = data.embedding.value;
+
+    return vector;
+}
+
+module.exports = {
+    getPaLMVector,
+};

src/vectors.js

@@ -15,6 +15,8 @@ async function getVector(source, text) {
             return require('./openai-vectors').getOpenAIVector(text);
         case 'transformers':
             return require('./embedding').getTransformersVector(text);
+        case 'palm':
+            return require('./palm-vectors').getPaLMVector(text);
     }
 
     throw new Error(`Unknown vector source ${source}`);