Form new chat context

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 () => {