diff --git a/public/scripts/extensions/infinity-context/index.js b/public/scripts/extensions/infinity-context/index.js index 0858e4a18..eae7ad4ca 100644 --- a/public/scripts/extensions/infinity-context/index.js +++ b/public/scripts/extensions/infinity-context/index.js @@ -1,5 +1,6 @@ -import { chat, saveSettingsDebounced, getCurrentChatId } from "../../../script.js"; -import { getApiUrl, extension_settings, getContext } from "../../extensions.js"; +import { saveSettingsDebounced, getCurrentChatId } from "../../../script.js"; +import { getApiUrl, extension_settings } from "../../extensions.js"; +import { splitRecursive } from "../../utils.js"; export { MODULE_NAME, chromadb_interceptGeneration }; const MODULE_NAME = 'chromadb'; @@ -14,6 +15,11 @@ const defaultSettings = { n_results_min: 1, n_results_max: 100, n_results_step: 1, + + split_length: 384, + split_length_min: 64, + split_length_max: 4096, + split_length_step: 64, }; const postHeaders = { @@ -28,6 +34,7 @@ async function loadSettings() { $('#chromadb_keep_context').val(extension_settings.chromadb.keep_context).trigger('input'); $('#chromadb_n_results').val(extension_settings.chromadb.n_results).trigger('input'); + $('#chromadb_split_length').val(extension_settings.chromadb.split_length).trigger('input'); } function onKeepContextInput() { @@ -42,16 +49,32 @@ function onNResultsInput() { saveSettingsDebounced(); } -async function moduleWorker() { - // ??? +function onSplitLengthInput() { + extension_settings.chromadb.split_length = Number($('#chromadb_split_length').val()); + $('#chromadb_split_length_value').text(extension_settings.chromadb.split_length); + saveSettingsDebounced(); } async function addMessages(chat_id, messages) { const url = new URL(getApiUrl()); url.pathname = '/api/chroma'; - const transformedMessages = messages.map((m, id) => ({ - id: id, + const messagesDeepCopy = JSON.parse(JSON.stringify(messages)); + const splittedMessages = []; + + let id = 0; + messagesDeepCopy.forEach(m => { + const split = splitRecursive(m.mes, extension_settings.chromadb.split_length); + splittedMessages.push(...split.map(text => ({ + ...m, + mes: text, + send_date: id, + id: `msg-${id++}`, + }))); + }); + + const transformedMessages = splittedMessages.map((m) => ({ + id: m.id, role: m.is_user ? 'user' : 'assistant', content: m.mes, date: m.send_date, @@ -92,15 +115,11 @@ async function queryMessages(chat_id, query) { return []; } -setInterval(moduleWorker, UPDATE_INTERVAL); - -window.chromadb_interceptGeneration = async () => { - const context = getContext(); +window.chromadb_interceptGeneration = async (chat) => { const currentChatId = getCurrentChatId(); if (currentChatId) { const messagesToStore = chat.slice(0, -extension_settings.chromadb.keep_context); - const messagesToKeep = chat.slice(-extension_settings.chromadb.keep_context); if (messagesToStore.length > 0) { await addMessages(currentChatId, messagesToStore); @@ -113,15 +132,11 @@ window.chromadb_interceptGeneration = async () => { queriedMessages.sort((a, b) => a.date - b.date); - const newChat = [ - ...queriedMessages.map(m => JSON.parse(m.meta)), - ...messagesToKeep, - ]; + const newChat = queriedMessages.map(m => JSON.parse(m.meta)); console.log(newChat); - // TODO replace current chat temporarily somehow... - + chat.splice(0, messagesToStore.length, newChat); } } } @@ -139,12 +154,15 @@ jQuery(async () => { + + `; $('#extensions_settings').append(settingsHtml); $('#chromadb_keep_context').on('input', onKeepContextInput); $('#chromadb_n_results').on('input', onNResultsInput); + $('#chromadb_split_length').on('input', onSplitLengthInput); await loadSettings(); }); \ No newline at end of file diff --git a/public/scripts/utils.js b/public/scripts/utils.js index 77763b536..158f95471 100644 --- a/public/scripts/utils.js +++ b/public/scripts/utils.js @@ -231,4 +231,35 @@ export function countOccurrences(string, character) { export function isOdd(number) { return number % 2 !== 0; +} + +/** Split string to parts no more than length in size */ +export function splitRecursive(input, length, delimitiers = ['\n\n', '\n', ' ', '']) { + const delim = delimitiers[0] ?? ''; + const parts = input.split(delim); + + const flatParts = parts.flatMap(p => { + if (p.length < length) return p; + return splitRecursive(input, length, delimitiers.slice(1)); + }); + + // Merge short chunks + const result = []; + let currentChunk = ''; + for (let i = 0; i < flatParts.length;) { + currentChunk = flatParts[i]; + let j = i + 1; + while (j < flatParts.length) { + const nextChunk = flatParts[j]; + if (currentChunk.length + nextChunk.length + delim.length <= length) { + currentChunk += delim + nextChunk; + } else { + break; + } + j++; + } + i = j; + result.push(currentChunk); + } + return result; } \ No newline at end of file