diff --git a/public/script.js b/public/script.js index e03bfd575..c5438d457 100644 --- a/public/script.js +++ b/public/script.js @@ -169,6 +169,7 @@ export { getStoppingStrings, getStatus, reloadMarkdownProcessor, + getCurrentChatId, chat, this_chid, selected_button, @@ -485,6 +486,15 @@ function reloadMarkdownProcessor(render_formulas = false) { return converter; } +function getCurrentChatId() { + if (selected_group) { + return groups.find(x => x.id == selected_group)?.chat_id; + } + else if (this_chid) { + return characters[this_chid].chat; + } +} + const CHARACTERS_PER_TOKEN_RATIO = 3.35; const talkativeness_default = 0.5; diff --git a/public/scripts/extensions/infinity-context/index.js b/public/scripts/extensions/infinity-context/index.js index 2f4371f0d..06b01a694 100644 --- a/public/scripts/extensions/infinity-context/index.js +++ b/public/scripts/extensions/infinity-context/index.js @@ -1,7 +1,5 @@ -import { - saveSettingsDebounced, -} from "../../../script.js"; -import { getApiUrl, getContext, extension_settings, defaultRequestArgs } from "../../extensions.js"; +import { chat, saveSettingsDebounced, getCurrentChatId } from "../../../script.js"; +import { getApiUrl, extension_settings, getContext } from "../../extensions.js"; export { MODULE_NAME, chromadb_interceptGeneration }; const MODULE_NAME = 'chromadb'; @@ -16,7 +14,12 @@ const defaultSettings = { n_results_min: 1, n_results_max: 100, n_results_step: 1, -} +}; + +const postHeaders = { + 'Content-Type': 'application/json', + 'Bypass-Tunnel-Reminder': 'bypass', +}; async function loadSettings() { if (Object.keys(extension_settings.chromadb).length === 0) { @@ -43,12 +46,89 @@ async function moduleWorker() { // ??? } +async function addMessages(chat_id, messages) { + const url = new URL(getApiUrl()); + url.pathname = '/api/chroma'; + + const transformedMessages = messages.map((m, id) => ({ + id: id, + role: m.is_user ? 'user' : 'assistant', + content: m.mes, + date: m.send_date, + })); + + const addMessagesResult = await fetch(url, { + method: 'POST', + headers: postHeaders, + body: JSON.stringify({ chat_id, messages: transformedMessages }), + }); + + if (addMessagesResult.ok) { + const addMessagesData = await addMessagesResult.json(); + + return addMessagesData; // { count: 1 } + } + + return { count: 0 }; +} + +async function queryMessages(chat_id, query) { + const url = new URL(getApiUrl()); + url.pathname = '/api/chroma/query'; + + const queryMessagesResult = await fetch(url, { + method: 'POST', + headers: postHeaders, + body: JSON.stringify({ chat_id, query, n_results: extension_settings.chromadb.n_results }), + }); + + if (queryMessagesResult.ok) { + const queryMessagesData = await queryMessagesResult.json(); + + return queryMessagesData; + } + + return []; +} + setInterval(moduleWorker, UPDATE_INTERVAL); window.chromadb_interceptGeneration = async () => { const context = getContext(); + const currentChatId = getCurrentChatId(); - // TODO substitute context + 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); + } + + const lastMessage = messagesToKeep[messagesToKeep.length - 1]; + + if (lastMessage) { + const queriedMessages = await queryMessages(currentChatId, lastMessage.mes); + + queriedMessages.sort((a, b) => a.date - b.date); + + const newChat = [ + ...queriedMessages.map(m => ({ + name: m.role === 'user' ? context.name1: context.name2, + is_user: m.role === 'user', + is_name: true, + send_date: m.date, + mes: m.content, + })), + ...messagesToKeep, + ]; + + console.log(newChat); + + // TODO replace current chat temporarily somehow... + + } + } } jQuery(async () => {