Local poe support

poe.js

@@ -64,6 +64,11 @@ class Client {
     bot_names = [];
     ws = null;
     ws_connected = false;
+    auto_reconnect = false;
+
+    constructor(auto_reconnect = false) {
+        this.auto_reconnect = auto_reconnect;
+    }
 
     async init(token, proxy = null) {
         this.proxy = proxy;
@@ -210,6 +215,10 @@ class Client {
             this.on_message(this.ws, message);
         });
 
+        this.ws.on('close', () => {
+            this.ws_connected = false;
+        });
+
         this.ws.on('error', (error) => {
             this.on_ws_error(this.ws, error);
         });
@@ -237,7 +246,10 @@ class Client {
     on_ws_error(ws, error) {
         logger.warn(`Websocket returned error: ${error}`);
         this.disconnect_ws();
-        this.connect_ws();
+
+        if (this.auto_reconnect) {
+            this.connect_ws();
+        }
     }
 
     on_message(ws, msg) {

public/script.js

@@ -124,6 +124,7 @@ export {
     is_send_press,
     api_server_textgenerationwebui,
     count_view_mes,
+    max_context,
     default_avatar,
     system_message_types,
     talkativeness_default,

public/scripts/extensions/poe/index.js

@@ -1,12 +1,18 @@
-import { getContext, getApiUrl } from "../../extensions.js";
+import { getContext } from "../../extensions.js";
 import { delay } from "../../utils.js";
-import { showSwipeButtons, hideSwipeButtons, getExtensionPrompt, extension_prompt_types } from "../../../script.js";
+import {
+    showSwipeButtons,
+    hideSwipeButtons,
+    getExtensionPrompt,
+    extension_prompt_types,
+    token as csrf_token,
+    max_context,
+} from "../../../script.js";
 
 const TOKEN_KEY = 'extensions_poe_token';
 const BOT_KEY = 'extensions_poe_bot';
 const RESPONSE_KEY = 'extensions_poe_response';
 const MESSAGE_KEY = 'extensions_poe_message';
-const MAX_CONTEXT_KEY = 'extensions_poe_max_context';
 
 const DEFAULT_MAX_CONTEXT = 2000;
 const MAX_RETRIES_FOR_ACTIVATION = 5;
@@ -30,7 +36,6 @@ let token;
 let bot;
 let jailbreak_response = DEFAULT_JAILBREAK_RESPONSE;
 let jailbreak_message = DEFAULT_JAILBREAK_MESSAGE;
-let max_context = DEFAULT_MAX_CONTEXT;
 let jailbroken = false;
 let got_reply = false;
 
@@ -39,10 +44,8 @@ function loadSettings() {
     bot = localStorage.getItem(BOT_KEY);
     jailbreak_response = localStorage.getItem(RESPONSE_KEY) ?? DEFAULT_JAILBREAK_RESPONSE;
     jailbreak_message = localStorage.getItem(MESSAGE_KEY) ?? DEFAULT_JAILBREAK_MESSAGE;
-    max_context = Number(localStorage.getItem(MAX_CONTEXT_KEY) ?? DEFAULT_MAX_CONTEXT);
     $('#poe_activation_response').val(jailbreak_response);
     $('#poe_activation_message').val(jailbreak_message);
-    $('#poe_max_context').val(max_context);
     $('#poe_token').val(token ?? '');
     selectBot();
 
@@ -63,7 +66,6 @@ function saveSettings() {
     localStorage.setItem(BOT_KEY, bot);
     localStorage.setItem(RESPONSE_KEY, jailbreak_response);
     localStorage.setItem(MESSAGE_KEY, jailbreak_message);
-    localStorage.setItem(MAX_CONTEXT_KEY, max_context);
 }
 
 function onTokenInput() {
@@ -110,6 +112,7 @@ async function generate(type, chat2, storyString, mesExamplesArray, promptBias,
         let messageSendString = '';
 
         const allMessages = [...chat2, ...mesExamplesArray];
+        const maxContext = Math.min(Number(max_context), DEFAULT_MAX_CONTEXT);
     
         for (let index = 0; index < allMessages.length; ++index) {
             const item = allMessages[index];
@@ -117,7 +120,7 @@ async function generate(type, chat2, storyString, mesExamplesArray, promptBias,
             const promptLength = context.encode(prompt + messageSendString + item + activator + extensionPrompt).length;
             await delay(1);
     
-            if (promptLength >= Number(max_context)) {
+            if (promptLength >= maxContext) {
                 break;
             }
     
@@ -148,13 +151,9 @@ async function purgeConversation(count = -1) {
         count,
     });
 
-    const apiUrl = new URL(getApiUrl());
+    const response = await fetch('/purge_poe', {
-    apiUrl.pathname = '/api/poe/purge';
-
-
-    const response = await fetch(apiUrl, {
         headers: {
-            'Bypass-Tunnel-Reminder': 'bypass',
+            'X-CSRF-Token': csrf_token,
             'Content-Type': 'application/json',
         },
         body: body,
@@ -171,12 +170,9 @@ async function sendMessage(prompt) {
         token,
     });
 
-    const apiUrl = new URL(getApiUrl());
+    const response = await fetch('/generate_poe', {
-    apiUrl.pathname = '/api/poe/generate';
-
-    const response = await fetch(apiUrl, {
         headers: {
-            'Bypass-Tunnel-Reminder': 'bypass',
+            'X-CSRF-Token': csrf_token,
             'Content-Type': 'application/json',
         },
         body: body,
@@ -199,12 +195,9 @@ async function sendMessage(prompt) {
 
 async function onConnectClick() {
     const body = JSON.stringify({ token: token });
-    const apiUrl = new URL(getApiUrl());
+    const response = await fetch('/status_poe', {
-    apiUrl.pathname = '/api/poe/status';
-
-    const response = await fetch(apiUrl, {
         headers: {
-            'Bypass-Tunnel-Reminder': 'bypass',
+            'X-CSRF-Token': csrf_token,
             'Content-Type': 'application/json',
         },
         body: body,
@@ -248,11 +241,6 @@ function onMessageInput() {
     saveSettings();
 }
 
-function onMaxContextInput() {
-    max_context = Number($(this).val());
-    saveSettings();
-}
-
 $('document').ready(function () {
     function addExtensionControls() {
         const settingsHtml = `
@@ -271,9 +259,7 @@ $('document').ready(function () {
         <label for="poe_activation_message">Jailbreak activation message</label>
         <textarea id="poe_activation_message" rows="3"></textarea>
         <label for="poe_activation_response">Jailbreak activation response</label>
-        <input id="poe_activation_response" class="text_pole" type="text />
+        <input id="poe_activation_response" class="text_pole" type="text" />
-        <label for="poe_max_context">Max context (in tokens)</label>
-        <input id="poe_max_context" class="text_pole" type="number" min="0" max="4096" />
         <input id="poe_connect" class="menu_button" type="button" value="Connect" />
         <div>
             <label for="poe_bots">List of bots:</label>
@@ -289,7 +275,6 @@ $('document').ready(function () {
         $('#poe_connect').on('click', onConnectClick);
         $('#poe_activation_response').on('input', onResponseInput);
         $('#poe_activation_message').on('input', onMessageInput);
-        $('#poe_max_context').on('input', onMaxContextInput);
     }
 
     addExtensionControls();

public/scripts/extensions/poe/manifest.json

@@ -1,9 +1,7 @@
 {
     "display_name": "poe.com generation",
     "loading_order": 0,
-    "requires": [
+    "requires": [],
-        "poe"
-    ],
     "optional": [],
     "js": "index.js",
     "css": "style.css"

server.js

@@ -35,6 +35,8 @@ const tiktoken = require('@dqbd/tiktoken');
 var Client = require('node-rest-client').Client;
 var client = new Client();
 
+let poe = require('./poe');
+
 var api_server = "http://0.0.0.0:5000";
 var api_novelai = "https://api.novelai.net";
 let api_openai = "https://api.openai.com/v1";
@@ -1614,6 +1616,62 @@ app.post('/deletegroup', jsonParser, async (request, response) => {
     return response.send({ ok: true });
 });
 
+const POE_DEFAULT_BOT = 'a2';
+
+async function getPoeClient(token) {
+    let client = new poe.Client();
+    await client.init(token);
+    return client;
+}
+
+app.post('/status_poe', jsonParser, async (request, response) => {
+    if (!request.body.token) {
+        return response.sendStatus(400);
+    }
+
+    const client = await getPoeClient(request.body.token);
+    const botNames = client.get_bot_names();
+    client.disconnect_ws();
+    return response.send({'bot_names': botNames});
+});
+
+app.post('/purge_poe', jsonParser, async (request, response) => {
+    if (!request.body.token) {
+        return response.sendStatus(400);
+    }
+
+    const token = request.body.token;
+    const bot = request.body.bot ?? POE_DEFAULT_BOT;
+    const count = request.body.count ?? -1;
+
+    const client = await getPoeClient(token);
+    await client.purge_conversation(bot, count)
+    client.disconnect_ws();
+
+    return response.send({"ok" : true});
+});
+
+app.post('/generate_poe', jsonParser, async (request, response) => {
+    if (!request.body.token || !request.body.prompt) {
+        return response.sendStatus(400);
+    }
+    
+    const token = request.body.token;
+    const prompt = request.body.prompt;
+    const bot = request.body.bot ?? POE_DEFAULT_BOT;
+    
+    const client = await getPoeClient(token);
+    
+    let reply;
+    for await (const mes of client.send_message(bot, prompt)) {
+        reply = mes.text;
+    }
+
+    client.disconnect_ws();
+
+    return response.send({'reply': reply});
+});
+
 function getThumbnailFolder(type) {
     let thumbnailFolder;