replaced returnChunks with return arg to use existing helper functions

This commit is contained in:
zerofata
2025-01-11 17:42:19 +13:00
committed by GitHub
parent 7f5ce54b1f
commit fc5debc660

View File

@ -30,6 +30,8 @@ import { textgen_types, textgenerationwebui_settings } from '../../textgen-setti
import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js'; import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
import { SlashCommand } from '../../slash-commands/SlashCommand.js'; import { SlashCommand } from '../../slash-commands/SlashCommand.js';
import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js'; import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';
import { SlashCommandEnumValue, enumTypes } from '../../slash-commands/SlashCommandEnumValue.js';
import { slashCommandReturnHelper } from '../../slash-commands/SlashCommandReturnHelper.js';
import { callGenericPopup, POPUP_RESULT, POPUP_TYPE } from '../../popup.js'; import { callGenericPopup, POPUP_RESULT, POPUP_TYPE } from '../../popup.js';
import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js'; import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js';
@ -1619,25 +1621,31 @@ jQuery(async () => {
const attachments = source ? getDataBankAttachmentsForSource(source, false) : getDataBankAttachments(false); const attachments = source ? getDataBankAttachmentsForSource(source, false) : getDataBankAttachments(false);
const collectionIds = await ingestDataBankAttachments(String(source)); const collectionIds = await ingestDataBankAttachments(String(source));
const queryResults = await queryMultipleCollections(collectionIds, String(query), resultSize, threshold); const queryResults = await queryMultipleCollections(collectionIds, String(query), resultSize, threshold);
const returnChunks = String(args?.returnChunks).toLowerCase() === 'true'
// Get URLs
if (returnChunks) { const urls = Object
.keys(queryResults)
.map(x => attachments.find(y => getFileCollectionId(y.url) === x))
.filter(x => x)
.map(x => x.url);
// Gets the actual text content of chunks
const getChunksText = () => {
let textResult = ''; let textResult = '';
for (const collectionId in queryResults) { for (const collectionId in queryResults) {
const metadata = queryResults[collectionId].metadata?.filter(x => x.text)?.sort((a, b) => a.index - b.index)?.map(x => x.text)?.filter(onlyUnique) || []; const metadata = queryResults[collectionId].metadata?.filter(x => x.text)?.sort((a, b) => a.index - b.index)?.map(x => x.text)?.filter(onlyUnique) || [];
textResult += metadata.join('\n') + '\n\n'; textResult += metadata.join('\n') + '\n\n';
} }
return textResult; return textResult;
} else { };
// Map collection IDs to file URLs
const urls = Object if (args.return === 'chunks') {
.keys(queryResults) return getChunksText();
.map(x => attachments.find(y => getFileCollectionId(y.url) === x))
.filter(x => x)
.map(x => x.url);
return JSON.stringify(urls);
} }
// @ts-ignore
return slashCommandReturnHelper.doReturn(args.return ?? 'object', urls, { objectToStringFunc: list => list.join('\n') });
}, },
aliases: ['databank-search', 'data-bank-search'], aliases: ['databank-search', 'data-bank-search'],
helpString: 'Search the Data Bank for a specific query using vector similarity. Returns a list of file URLs with the most relevant content.', helpString: 'Search the Data Bank for a specific query using vector similarity. Returns a list of file URLs with the most relevant content.',
@ -1645,7 +1653,17 @@ jQuery(async () => {
new SlashCommandNamedArgument('threshold', 'Threshold for the similarity score in the [0, 1] range. Uses the global config value if not set.', ARGUMENT_TYPE.NUMBER, false, false, ''), new SlashCommandNamedArgument('threshold', 'Threshold for the similarity score in the [0, 1] range. Uses the global config value if not set.', ARGUMENT_TYPE.NUMBER, false, false, ''),
new SlashCommandNamedArgument('resultSize', 'Maximum number of query results to return.', ARGUMENT_TYPE.NUMBER, false, false, ''), new SlashCommandNamedArgument('resultSize', 'Maximum number of query results to return.', ARGUMENT_TYPE.NUMBER, false, false, ''),
new SlashCommandNamedArgument('source', 'Optional filter for the attachments by source.', ARGUMENT_TYPE.STRING, false, false, '', ['global', 'character', 'chat']), new SlashCommandNamedArgument('source', 'Optional filter for the attachments by source.', ARGUMENT_TYPE.STRING, false, false, '', ['global', 'character', 'chat']),
new SlashCommandNamedArgument('returnChunks', 'If true, returns the actual content chunks instead of URLs.', ARGUMENT_TYPE.STRING, false, false, '', ['true', 'false']), SlashCommandNamedArgument.fromProps({
name: 'return',
description: 'How you want the return value to be provided',
typeList: [ARGUMENT_TYPE.STRING],
defaultValue: 'object',
enumList: [
new SlashCommandEnumValue('chunks', 'Return the actual content chunks', enumTypes.enum, '{}'),
...slashCommandReturnHelper.enumList({ allowObject: true })
],
forceEnum: true,
})
], ],
unnamedArgumentList: [ unnamedArgumentList: [
new SlashCommandArgument('Query to search by.', ARGUMENT_TYPE.STRING, true, false), new SlashCommandArgument('Query to search by.', ARGUMENT_TYPE.STRING, true, false),