Implement collection scopes for vector models (#2846) * Implement collection scopes for vector models * Update makersuite-vectors.js to use Gemini API text-embedding-004 model * Add model scope for Google vectors * Fix purge log * Refactor header setting * Fix typo * Only display UI warning if scopes disabled * Remove i18n attribute --------- Co-authored-by: ceruleandeep <83318388+ceruleandeep@users.noreply.github.com>

9ef33852553975b0cc132e97454cab87cf7a134a

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
Parents: 7eb7204dc747194042ca292558a67fed0cfa4ac0
4 files changed, +166 -182Ignore whitespace
default/config.yaml+3 -0
@@ -108,6 +108,9 @@ enableExtensionsAutoUpdate: true
108108# Additional model tokenizers can be downloaded on demand.
109109# Disabling will fallback to another locally available tokenizer.
110110enableDownloadableTokenizers: true
111+# Vector storage settings
112+vectors:
113+ enableModelScopes: false
111114# Extension settings
112115extras:
113116 # Disables automatic model download from HuggingFace
public/scripts/extensions/vectors/index.js+55 -92
@@ -718,7 +718,7 @@ async function getQueryText(chat, initiator) {
718718async function getSavedHashes(collectionId) {
719719 const response = await fetch('/api/vector/list', {
720720 method: 'POST',
721721 headers: getRequestHeadersgetVectorHeaders(),
722722 body: JSON.stringify({
723723 collectionId: collectionId,
724724 source: settings.source,
@@ -737,25 +737,43 @@ function getVectorHeaders() {
737737 const headers = getRequestHeaders();
738738 switch (settings.source) {
739739 case 'extras':
740- addExtrasHeaders(headers);
740+ Object.assign(headers, {
741+ 'X-Extras-Url': extension_settings.apiUrl,
742+ 'X-Extras-Key': extension_settings.apiKey,
743+ });
741744 break;
742745 case 'togetherai':
743- addTogetherAiHeaders(headers);
746+ Object.assign(headers, {
747+ 'X-Togetherai-Model': extension_settings.vectors.togetherai_model,
748+ });
744749 break;
745750 case 'openai':
746- addOpenAiHeaders(headers);
751+ Object.assign(headers, {
752+ 'X-OpenAI-Model': extension_settings.vectors.openai_model,
753+ });
747754 break;
748755 case 'cohere':
749- addCohereHeaders(headers);
756+ Object.assign(headers, {
757+ 'X-Cohere-Model': extension_settings.vectors.cohere_model,
758+ });
750759 break;
751760 case 'ollama':
752- addOllamaHeaders(headers);
761+ Object.assign(headers, {
762+ 'X-Ollama-Model': extension_settings.vectors.ollama_model,
763+ 'X-Ollama-URL': textgenerationwebui_settings.server_urls[textgen_types.OLLAMA],
764+ 'X-Ollama-Keep': !!extension_settings.vectors.ollama_keep,
765+ });
753766 break;
754767 case 'llamacpp':
755- addLlamaCppHeaders(headers);
768+ Object.assign(headers, {
769+ 'X-LlamaCpp-URL': textgenerationwebui_settings.server_urls[textgen_types.LLAMACPP],
770+ });
756771 break;
757772 case 'vllm':
758- addVllmHeaders(headers);
773+ Object.assign(headers, {
774+ 'X-Vllm-URL': textgenerationwebui_settings.server_urls[textgen_types.VLLM],
775+ 'X-Vllm-Model': extension_settings.vectors.vllm_model,
776+ });
759777 break;
760778 default:
761779 break;
@@ -764,81 +782,6 @@ function getVectorHeaders() {
764782}
765783
766784/**
767- * Add headers for the Extras API source.
768- * @param {object} headers Headers object
769- */
770-function addExtrasHeaders(headers) {
771- console.log(`Vector source is extras, populating API URL: ${extension_settings.apiUrl}`);
772- Object.assign(headers, {
773- 'X-Extras-Url': extension_settings.apiUrl,
774- 'X-Extras-Key': extension_settings.apiKey,
775- });
776-}
777-
778-/**
779- * Add headers for the TogetherAI API source.
780- * @param {object} headers Headers object
781- */
782-function addTogetherAiHeaders(headers) {
783- Object.assign(headers, {
784- 'X-Togetherai-Model': extension_settings.vectors.togetherai_model,
785- });
786-}
787-
788-/**
789- * Add headers for the OpenAI API source.
790- * @param {object} headers Header object
791- */
792-function addOpenAiHeaders(headers) {
793- Object.assign(headers, {
794- 'X-OpenAI-Model': extension_settings.vectors.openai_model,
795- });
796-}
797-
798-/**
799- * Add headers for the Cohere API source.
800- * @param {object} headers Header object
801- */
802-function addCohereHeaders(headers) {
803- Object.assign(headers, {
804- 'X-Cohere-Model': extension_settings.vectors.cohere_model,
805- });
806-}
807-
808-/**
809- * Add headers for the Ollama API source.
810- * @param {object} headers Header object
811- */
812-function addOllamaHeaders(headers) {
813- Object.assign(headers, {
814- 'X-Ollama-Model': extension_settings.vectors.ollama_model,
815- 'X-Ollama-URL': textgenerationwebui_settings.server_urls[textgen_types.OLLAMA],
816- 'X-Ollama-Keep': !!extension_settings.vectors.ollama_keep,
817- });
818-}
819-
820-/**
821- * Add headers for the LlamaCpp API source.
822- * @param {object} headers Header object
823- */
824-function addLlamaCppHeaders(headers) {
825- Object.assign(headers, {
826- 'X-LlamaCpp-URL': textgenerationwebui_settings.server_urls[textgen_types.LLAMACPP],
827- });
828-}
829-
830-/**
831- * Add headers for the VLLM API source.
832- * @param {object} headers Header object
833- */
834-function addVllmHeaders(headers) {
835- Object.assign(headers, {
836- 'X-Vllm-URL': textgenerationwebui_settings.server_urls[textgen_types.VLLM],
837- 'X-Vllm-Model': extension_settings.vectors.vllm_model,
838- });
839-}
840-
841-/**
842785 * Inserts vector items into a collection
843786 * @param {string} collectionId - The collection to insert into
844787 * @param {{ hash: number, text: string }[]} items - The items to insert
@@ -901,7 +844,7 @@ function throwIfSourceInvalid() {
901844async function deleteVectorItems(collectionId, hashes) {
902845 const response = await fetch('/api/vector/delete', {
903846 method: 'POST',
904847 headers: getRequestHeadersgetVectorHeaders(),
905848 body: JSON.stringify({
906849 collectionId: collectionId,
907850 hashes: hashes,
@@ -987,7 +930,7 @@ async function purgeFileVectorIndex(fileUrl) {
987930
988931 const response = await fetch('/api/vector/purge', {
989932 method: 'POST',
990933 headers: getRequestHeadersgetVectorHeaders(),
991934 body: JSON.stringify({
992935 collectionId: collectionId,
993936 }),
@@ -1016,7 +959,7 @@ async function purgeVectorIndex(collectionId) {
1016959
1017960 const response = await fetch('/api/vector/purge', {
1018961 method: 'POST',
1019962 headers: getRequestHeadersgetVectorHeaders(),
1020963 body: JSON.stringify({
1021964 collectionId: collectionId,
1022965 }),
@@ -1041,7 +984,7 @@ async function purgeAllVectorIndexes() {
1041984 try {
1042985 const response = await fetch('/api/vector/purge-all', {
1043986 method: 'POST',
1044987 headers: getRequestHeadersgetVectorHeaders(),
1045988 });
1046989
1047990 if (!response.ok) {
@@ -1056,6 +999,25 @@ async function purgeAllVectorIndexes() {
1056999 }
10571000}
10581001
1002+async function isModelScopesEnabled() {
1003+ try {
1004+ const response = await fetch('/api/vector/scopes-enabled', {
1005+ method: 'GET',
1006+ headers: getVectorHeaders(),
1007+ });
1008+
1009+ if (!response.ok) {
1010+ return false;
1011+ }
1012+
1013+ const data = await response.json();
1014+ return data?.enabled ?? false;
1015+ } catch (error) {
1016+ console.error('Vectors: Failed to check model scopes', error);
1017+ return false;
1018+ }
1019+}
1020+
10591021function toggleSettings() {
10601022 $('#vectors_files_settings').toggle(!!settings.enabled_files);
10611023 $('#vectors_chats_settings').toggle(!!settings.enabled_chats);
@@ -1320,6 +1282,7 @@ jQuery(async () => {
13201282 }
13211283
13221284 Object.assign(settings, extension_settings.vectors);
1285+ const scopesEnabled = await isModelScopesEnabled();
13231286
13241287 // Migrate from TensorFlow to Transformers
13251288 settings.source = settings.source !== 'local' ? settings.source : 'transformers';
@@ -1371,31 +1334,31 @@ jQuery(async () => {
13711334 saveSettingsDebounced();
13721335 });
13731336 $('#vectors_togetherai_model').val(settings.togetherai_model).on('change', () => {
13741337 !scopesEnabled && $('#vectors_modelWarning').show();
13751338 settings.togetherai_model = String($('#vectors_togetherai_model').val());
13761339 Object.assign(extension_settings.vectors, settings);
13771340 saveSettingsDebounced();
13781341 });
13791342 $('#vectors_openai_model').val(settings.openai_model).on('change', () => {
13801343 !scopesEnabled && $('#vectors_modelWarning').show();
13811344 settings.openai_model = String($('#vectors_openai_model').val());
13821345 Object.assign(extension_settings.vectors, settings);
13831346 saveSettingsDebounced();
13841347 });
13851348 $('#vectors_cohere_model').val(settings.cohere_model).on('change', () => {
13861349 !scopesEnabled && $('#vectors_modelWarning').show();
13871350 settings.cohere_model = String($('#vectors_cohere_model').val());
13881351 Object.assign(extension_settings.vectors, settings);
13891352 saveSettingsDebounced();
13901353 });
13911354 $('#vectors_ollama_model').val(settings.ollama_model).on('input', () => {
13921355 !scopesEnabled && $('#vectors_modelWarning').show();
13931356 settings.ollama_model = String($('#vectors_ollama_model').val());
13941357 Object.assign(extension_settings.vectors, settings);
13951358 saveSettingsDebounced();
13961359 });
13971360 $('#vectors_vllm_model').val(settings.vllm_model).on('input', () => {
13981361 !scopesEnabled && $('#vectors_modelWarning').show();
13991362 settings.vllm_model = String($('#vectors_vllm_model').val());
14001363 Object.assign(extension_settings.vectors, settings);
14011364 saveSettingsDebounced();
public/scripts/extensions/vectors/settings.html+3 -2
@@ -98,8 +98,9 @@
9898
9999 <small id="vectors_modelWarning">
100100 <i class="fa-solid fa-exclamation-triangle"></i>
101- <span data-i18n="Vectors Model Warning">
101+ <span>
102- It is recommended to purge vectors when changing the model mid-chat. Otherwise, it will lead to sub-par results.
102+ Set <code>vectors.enableModelScopes</code> to true in config.yaml to switch between vectorization models without needing to purge existing vectors.
103+ This option will soon be enabled by default.
103104 </span>
104105 </small>
105106
src/endpoints/vectors.js+105 -88
@@ -4,6 +4,7 @@ const fs = require('fs');
44const express = require('express');
55const sanitize = require('sanitize-filename');
66const { jsonParser } = require('../express-common');
7+const { getConfigValue, color } = require('../util');
78
89// Don't forget to add new sources to the SOURCES array
910const SOURCES = [
@@ -110,18 +111,94 @@ async function getBatchVector(source, sourceSettings, texts, isQuery, directorie
110111}
111112
112113/**
114+ * Extracts settings for the vectorization sources from the HTTP request headers.
115+ * @param {string} source - Which source to extract settings for.
116+ * @param {object} request - The HTTP request object.
117+ * @returns {object} - An object that can be used as `sourceSettings` in functions that take that parameter.
118+ */
119+function getSourceSettings(source, request) {
120+ switch (source) {
121+ case 'togetherai':
122+ return {
123+ model: String(request.headers['x-togetherai-model']),
124+ };
125+ case 'openai':
126+ return {
127+ model: String(request.headers['x-openai-model']),
128+ };
129+ case 'cohere':
130+ return {
131+ model: String(request.headers['x-cohere-model']),
132+ };
133+ case 'llamacpp':
134+ return {
135+ apiUrl: String(request.headers['x-llamacpp-url']),
136+ };
137+ case 'vllm':
138+ return {
139+ apiUrl: String(request.headers['x-vllm-url']),
140+ model: String(request.headers['x-vllm-model']),
141+ };
142+ case 'ollama':
143+ return {
144+ apiUrl: String(request.headers['x-ollama-url']),
145+ model: String(request.headers['x-ollama-model']),
146+ keep: Boolean(request.headers['x-ollama-keep']),
147+ };
148+ case 'extras':
149+ return {
150+ extrasUrl: String(request.headers['x-extras-url']),
151+ extrasKey: String(request.headers['x-extras-key']),
152+ };
153+ case 'local':
154+ return {
155+ model: getConfigValue('extras.embeddingModel', ''),
156+ };
157+ case 'palm':
158+ return {
159+ // TODO: Add support for multiple models
160+ model: 'text-embedding-004',
161+ };
162+ default:
163+ return {};
164+ }
165+}
166+
167+/**
168+ * Gets the model scope for the source.
169+ * @param {object} sourceSettings - The settings for the source
170+ * @returns {string} The model scope for the source
171+ */
172+function getModelScope(sourceSettings) {
173+ const scopesEnabled = getConfigValue('vectors.enableModelScopes', false);
174+ const warningShown = global.process.env.VECTORS_MODEL_SCOPE_WARNING_SHOWN === 'true';
175+
176+ if (!scopesEnabled && !warningShown) {
177+ console.log();
178+ console.warn(color.red('[DEPRECATION NOTICE]'), 'Model scopes for Vectore Storage are disabled, but will soon be required.');
179+ console.log(`To enable model scopes, set the ${color.cyan('vectors.enableModelScopes')} in config.yaml to ${color.green(true)}.`);
180+ console.log('This message won\'t be shown again in the current session.');
181+ console.log();
182+ global.process.env.VECTORS_MODEL_SCOPE_WARNING_SHOWN = 'true';
183+ }
184+
185+ return scopesEnabled ? (sourceSettings?.model || '') : '';
186+}
187+
188+/**
113189 * Gets the index for the vector collection
114190 * @param {import('../users').UserDirectoryList} directories - User directories
115191 * @param {string} collectionId - The collection ID
116192 * @param {string} source - The source of the vector
117193 * @param {booleanobject} createsourceSettings - WhetherThe tomodel createfor the index if it doesn't existsource
118194 * @returns {Promise<vectra.LocalIndex>} - The index for the collection
119195 */
120196async function getIndex(directories, collectionId, source, create = truesourceSettings) {
121- const pathToFile = path.join(directories.vectors, sanitize(source), sanitize(collectionId));
197+ const model = getModelScope(sourceSettings);
198+ const pathToFile = path.join(directories.vectors, sanitize(source), sanitize(collectionId), sanitize(model));
122199 const store = new vectra.LocalIndex(pathToFile);
123200
124201 if (create && !await store.isIndexCreated()) {
125202 await store.createIndex();
126203 }
127204
@@ -137,7 +214,7 @@ async function getIndex(directories, collectionId, source, create = true) {
137214 * @param {{ hash: number; text: string; index: number; }[]} items - The items to insert
138215 */
139216async function insertVectorItems(directories, collectionId, source, sourceSettings, items) {
140217 const store = await getIndex(directories, collectionId, source, sourceSettings);
141218
142219 await store.beginUpdate();
143220
@@ -157,10 +234,11 @@ async function insertVectorItems(directories, collectionId, source, sourceSettin
157234 * @param {import('../users').UserDirectoryList} directories - User directories
158235 * @param {string} collectionId - The collection ID
159236 * @param {string} source - The source of the vector
237+ * @param {Object} sourceSettings - Settings for the source, if it needs any
160238 * @returns {Promise<number[]>} - The hashes of the items in the collection
161239 */
162240async function getSavedHashes(directories, collectionId, source, sourceSettings) {
163241 const store = await getIndex(directories, collectionId, source, sourceSettings);
164242
165243 const items = await store.listItems();
166244 const hashes = items.map(x => Number(x.metadata.hash));
@@ -173,10 +251,11 @@ async function getSavedHashes(directories, collectionId, source) {
173251 * @param {import('../users').UserDirectoryList} directories - User directories
174252 * @param {string} collectionId - The collection ID
175253 * @param {string} source - The source of the vector
254+ * @param {Object} sourceSettings - Settings for the source, if it needs any
176255 * @param {number[]} hashes - The hashes of the items to delete
177256 */
178257async function deleteVectorItems(directories, collectionId, source, sourceSettings, hashes) {
179258 const store = await getIndex(directories, collectionId, source, sourceSettings);
180259 const items = await store.listItemsByMetadata({ hash: { '$in': hashes } });
181260
182261 await store.beginUpdate();
@@ -200,7 +279,7 @@ async function deleteVectorItems(directories, collectionId, source, hashes) {
200279 * @returns {Promise<{hashes: number[], metadata: object[]}>} - The metadata of the items that match the search text
201280 */
202281async function queryCollection(directories, collectionId, source, sourceSettings, searchText, topK, threshold) {
203282 const store = await getIndex(directories, collectionId, source, sourceSettings);
204283 const vector = await getVector(source, sourceSettings, searchText, true, directories);
205284
206285 const result = await store.queryItems(vector, topK);
@@ -226,7 +305,7 @@ async function multiQueryCollection(directories, collectionIds, source, sourceSe
226305 const results = [];
227306
228307 for (const collectionId of collectionIds) {
229308 const store = await getIndex(directories, collectionId, source, sourceSettings);
230309 const result = await store.queryItems(vector, topK);
231310 results.push(...result.map(result => ({ collectionId, result })));
232311 }
@@ -255,71 +334,6 @@ async function multiQueryCollection(directories, collectionIds, source, sourceSe
255334}
256335
257336/**
258- * Extracts settings for the vectorization sources from the HTTP request headers.
259- * @param {string} source - Which source to extract settings for.
260- * @param {object} request - The HTTP request object.
261- * @returns {object} - An object that can be used as `sourceSettings` in functions that take that parameter.
262- */
263-function getSourceSettings(source, request) {
264- if (source === 'togetherai') {
265- const model = String(request.headers['x-togetherai-model']);
266-
267- return {
268- model: model,
269- };
270- } else if (source === 'openai') {
271- const model = String(request.headers['x-openai-model']);
272-
273- return {
274- model: model,
275- };
276- } else if (source === 'cohere') {
277- const model = String(request.headers['x-cohere-model']);
278-
279- return {
280- model: model,
281- };
282- } else if (source === 'llamacpp') {
283- const apiUrl = String(request.headers['x-llamacpp-url']);
284-
285- return {
286- apiUrl: apiUrl,
287- };
288- } else if (source === 'vllm') {
289- const apiUrl = String(request.headers['x-vllm-url']);
290- const model = String(request.headers['x-vllm-model']);
291-
292- return {
293- apiUrl: apiUrl,
294- model: model,
295- };
296- } else if (source === 'ollama') {
297- const apiUrl = String(request.headers['x-ollama-url']);
298- const model = String(request.headers['x-ollama-model']);
299- const keep = Boolean(request.headers['x-ollama-keep']);
300-
301- return {
302- apiUrl: apiUrl,
303- model: model,
304- keep: keep,
305- };
306- } else {
307- // Extras API settings to connect to the Extras embeddings provider
308- let extrasUrl = '';
309- let extrasKey = '';
310- if (source === 'extras') {
311- extrasUrl = String(request.headers['x-extras-url']);
312- extrasKey = String(request.headers['x-extras-key']);
313- }
314-
315- return {
316- extrasUrl: extrasUrl,
317- extrasKey: extrasKey,
318- };
319- }
320-}
321-
322-/**
323337 * Performs a request to regenerate the index if it is corrupted.
324338 * @param {import('express').Request} req Express request object
325339 * @param {import('express').Response} res Express response object
@@ -330,9 +344,10 @@ async function regenerateCorruptedIndexErrorHandler(req, res, error) {
330344 if (error instanceof SyntaxError && !req.query.regenerated) {
331345 const collectionId = String(req.body.collectionId);
332346 const source = String(req.body.source) || 'transformers';
347+ const sourceSettings = getSourceSettings(source, req);
333348
334349 if (collectionId && source) {
335350 const index = await getIndex(req.user.directories, collectionId, source, falsesourceSettings);
336351 const exists = await index.isIndexCreated();
337352
338353 if (exists) {
@@ -350,6 +365,11 @@ async function regenerateCorruptedIndexErrorHandler(req, res, error) {
350365
351366const router = express.Router();
352367
368+router.get('/scopes-enabled', (_req, res) => {
369+ const scopesEnabled = getConfigValue('vectors.enableModelScopes', false);
370+ return res.json({ enabled: scopesEnabled });
371+});
372+
353373router.post('/query', jsonParser, async (req, res) => {
354374 try {
355375 if (!req.body.collectionId || !req.body.searchText) {
@@ -416,8 +436,9 @@ router.post('/list', jsonParser, async (req, res) => {
416436
417437 const collectionId = String(req.body.collectionId);
418438 const source = String(req.body.source) || 'transformers';
439+ const sourceSettings = getSourceSettings(source, req);
419440
420441 const hashes = await getSavedHashes(req.user.directories, collectionId, source, sourceSettings);
421442 return res.json(hashes);
422443 } catch (error) {
423444 return regenerateCorruptedIndexErrorHandler(req, res, error);
@@ -433,8 +454,9 @@ router.post('/delete', jsonParser, async (req, res) => {
433454 const collectionId = String(req.body.collectionId);
434455 const hashes = req.body.hashes.map(x => Number(x));
435456 const source = String(req.body.source) || 'transformers';
457+ const sourceSettings = getSourceSettings(source, req);
436458
437459 await deleteVectorItems(req.user.directories, collectionId, source, sourceSettings, hashes);
438460 return res.sendStatus(200);
439461 } catch (error) {
440462 return regenerateCorruptedIndexErrorHandler(req, res, error);
@@ -468,17 +490,12 @@ router.post('/purge', jsonParser, async (req, res) => {
468490 const collectionId = String(req.body.collectionId);
469491
470492 for (const source of SOURCES) {
471493 const indexsourcePath = await getIndexpath.join(req.user.directories, collectionId.vectors, sanitize(source), falsesanitize(collectionId));
472-
494+ if (!fs.existsSync(sourcePath)) {
473- const exists = await index.isIndexCreated();
474-
475- if (!exists) {
476495 continue;
477496 }
478-
497+ await fs.promises.rm(sourcePath, { recursive: true });
479- const path = index.folderPath;
498+ console.log(`Deleted vector index at ${sourcePath}`);
480- await index.deleteIndex();
481- console.log(`Deleted vector index at ${path}`);
482499 }
483500
484501 return res.sendStatus(200);