Add check for "embeddings" module.

This commit is contained in:
Cohee 2024-01-24 16:51:57 +02:00
parent 62170afd03
commit 0ad756c923
2 changed files with 41 additions and 16 deletions

View File

@ -1,5 +1,5 @@
import { eventSource, event_types, extension_prompt_types, getCurrentChatId, getRequestHeaders, is_send_press, saveSettingsDebounced, setExtensionPrompt, substituteParams } from '../../../script.js'; 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 { ModuleWorkerWrapper, extension_settings, getContext, modules, renderExtensionTemplate } from '../../extensions.js';
import { collapseNewlines } from '../../power-user.js'; import { collapseNewlines } from '../../power-user.js';
import { SECRET_KEYS, secret_state } from '../../secrets.js'; import { SECRET_KEYS, secret_state } from '../../secrets.js';
import { debounce, getStringHash as calculateHash, waitUntilCondition, onlyUnique, splitRecursive } from '../../utils.js'; import { debounce, getStringHash as calculateHash, waitUntilCondition, onlyUnique, splitRecursive } from '../../utils.js';
@ -152,8 +152,25 @@ async function synchronizeChat(batchSize = 5) {
return newVectorItems.length - batchSize; return newVectorItems.length - batchSize;
} catch (error) { } catch (error) {
/**
* Gets the error message for a given cause
* @param {string} cause Error cause key
* @returns {string} Error message
*/
function getErrorMessage(cause) {
switch (cause) {
case 'api_key_missing':
return 'API key missing. Save it in the "API Connections" panel.';
case 'extras_module_missing':
return 'Extras API must provide an "embeddings" module.';
default:
return 'Check server console for more details';
}
}
console.error('Vectors: Failed to synchronize chat', error); 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';
const message = getErrorMessage(error.cause);
toastr.error(message, 'Vectorization failed'); toastr.error(message, 'Vectorization failed');
return -1; return -1;
} finally { } finally {
@ -411,6 +428,18 @@ async function getSavedHashes(collectionId) {
return hashes; return hashes;
} }
/**
* Add headers for the Extras API source.
* @param {object} headers Headers object
*/
function addExtrasHeaders(headers) {
console.log(`Vector source is extras, populating API URL: ${extension_settings.apiUrl}`);
Object.assign(headers, {
'X-Extras-Url': extension_settings.apiUrl,
'X-Extras-Key': extension_settings.apiKey,
});
}
/** /**
* 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
@ -424,13 +453,13 @@ async function insertVectorItems(collectionId, items) {
throw new Error('Vectors: API key missing', { cause: 'api_key_missing' }); throw new Error('Vectors: API key missing', { cause: 'api_key_missing' });
} }
if (settings.source === 'extras' && !modules.includes('embeddings')) {
throw new Error('Vectors: Embeddings module missing', { cause: 'extras_module_missing' });
}
const headers = getRequestHeaders(); const headers = getRequestHeaders();
if (settings.source === 'extras') { if (settings.source === 'extras') {
console.log(`Vector source is extras, populating API URL: ${extension_settings.apiUrl}`); addExtrasHeaders(headers);
Object.assign(headers, {
'X-Extras-Url': extension_settings.apiUrl,
'X-Extras-Key': extension_settings.apiKey
});
} }
const response = await fetch('/api/vector/insert', { const response = await fetch('/api/vector/insert', {
@ -479,11 +508,7 @@ async function deleteVectorItems(collectionId, hashes) {
async function queryCollection(collectionId, searchText, topK) { async function queryCollection(collectionId, searchText, topK) {
const headers = getRequestHeaders(); const headers = getRequestHeaders();
if (settings.source === 'extras') { if (settings.source === 'extras') {
console.log(`Vector source is extras, populating API URL: ${extension_settings.apiUrl}`); addExtrasHeaders(headers);
Object.assign(headers, {
'X-Extras-Url': extension_settings.apiUrl,
'X-Extras-Key': extension_settings.apiKey
});
} }
const response = await fetch('/api/vector/query', { const response = await fetch('/api/vector/query', {

View File

@ -4,7 +4,7 @@ const fetch = require('node-fetch').default;
* Gets the vector for the given text from SillyTavern-extras * Gets the vector for the given text from SillyTavern-extras
* @param {string[]} texts - The array of texts to get the vectors for * @param {string[]} texts - The array of texts to get the vectors for
* @param {string} apiUrl - The Extras API URL * @param {string} apiUrl - The Extras API URL
* @param {string} - The Extras API key, or empty string if API key not enabled * @param {string} apiKey - The Extras API key, or empty string if API key not enabled
* @returns {Promise<number[][]>} - The array of vectors for the texts * @returns {Promise<number[][]>} - The array of vectors for the texts
*/ */
async function getExtrasBatchVector(texts, apiUrl, apiKey) { async function getExtrasBatchVector(texts, apiUrl, apiKey) {
@ -15,7 +15,7 @@ async function getExtrasBatchVector(texts, apiUrl, apiKey) {
* Gets the vector for the given text from SillyTavern-extras * Gets the vector for the given text from SillyTavern-extras
* @param {string} text - The text to get the vector for * @param {string} text - The text to get the vector for
* @param {string} apiUrl - The Extras API URL * @param {string} apiUrl - The Extras API URL
* @param {string} - The Extras API key, or empty string if API key not enabled * @param {string} apiKey - The Extras API key, or empty string if API key not enabled
* @returns {Promise<number[]>} - The vector for the text * @returns {Promise<number[]>} - The vector for the text
*/ */
async function getExtrasVector(text, apiUrl, apiKey) { async function getExtrasVector(text, apiUrl, apiKey) {
@ -26,8 +26,8 @@ async function getExtrasVector(text, apiUrl, apiKey) {
* Gets the vector for the given text from SillyTavern-extras * Gets the vector for the given text from SillyTavern-extras
* @param {string|string[]} text - The text or texts to get the vector(s) for * @param {string|string[]} text - The text or texts to get the vector(s) for
* @param {string} apiUrl - The Extras API URL * @param {string} apiUrl - The Extras API URL
* @param {string} - The Extras API key, or empty string if API key not enabled * @param {string} apiKey - The Extras API key, or empty string if API key not enabled *
* @returns {Promise<number[]>|Promise<number[][]>} - The vector for a single text, or the array of vectors for multiple texts * @returns {Promise<Array>} - The vector for a single text if input is string, or the array of vectors for multiple texts if input is string[]
*/ */
async function getExtrasVectorImpl(text, apiUrl, apiKey) { async function getExtrasVectorImpl(text, apiUrl, apiKey) {
let url; let url;