From e1b514ef5ab1c79f9f83a6e8b2d7031e2a6d6f0d Mon Sep 17 00:00:00 2001
From: SillyLossy <titan0599@gmail.com>
Date: Wed, 12 Apr 2023 02:26:09 +0300
Subject: [PATCH] Add poe streaming

---
 public/index.html     |  9 +++++++++
 public/script.js      | 16 ++++++++++++----
 public/scripts/poe.js | 34 +++++++++++++++++++++++++++++++---
 server.js             | 37 +++++++++++++++++++++++++++----------
 4 files changed, 79 insertions(+), 17 deletions(-)

diff --git a/public/index.html b/public/index.html
index 849d1f817..27475003c 100644
--- a/public/index.html
+++ b/public/index.html
@@ -347,6 +347,15 @@
                             </div>
                         </div>
                         <div id="range_block_poe">
+                            <div class="range-block">
+                                <label class="checkbox_label" for="poe_streaming">
+                                    <input id="poe_streaming" type="checkbox" />
+                                    Streaming
+                                </label>
+                                <div class="range-block-counter justifyLeft">
+                                    Displays the message text as soon as it is generated
+                                </div>
+                            </div>
                             <div class="range-block">
                                 <label for="poe_auto_purge" class="checkbox_label">
                                     <input id="poe_auto_purge" type="checkbox">
diff --git a/public/script.js b/public/script.js
index c5df2cf0e..ec62a0c5f 100644
--- a/public/script.js
+++ b/public/script.js
@@ -1114,7 +1114,7 @@ function appendToStoryString(value, prefix) {
 }
 
 function isStreamingEnabled() {
-    return (main_api == 'openai' && oai_settings.stream_openai);
+    return (main_api == 'openai' && oai_settings.stream_openai) || (main_api == 'poe' && poe_settings.streaming);
 }
 
 class StreamingProcessor {
@@ -1811,8 +1811,6 @@ async function Generate(type, automatic_trigger, force_name2) {
 
                 if (isStreamingEnabled()) {
                     streamingProcessor.generator = await sendOpenAIRequest(prompt);
-                    await streamingProcessor.generate();
-                    streamingProcessor = null;
                 }
                 else {
                     sendOpenAIRequest(prompt).then(onSuccess).catch(onError);
@@ -1822,7 +1820,12 @@ async function Generate(type, automatic_trigger, force_name2) {
                 generateHorde(finalPromt, generate_data).then(onSuccess).catch(onError);
             }
             else if (main_api == 'poe') {
-                generatePoe(finalPromt).then(onSuccess).catch(onError);
+                if (isStreamingEnabled()) {
+                    streamingProcessor.generator = await generatePoe(finalPromt);
+                }
+                else {
+                    generatePoe(finalPromt).then(onSuccess).catch(onError);
+                }
             }
             else {
                 jQuery.ajax({
@@ -1840,6 +1843,11 @@ async function Generate(type, automatic_trigger, force_name2) {
                 }); //end of "if not data error"
             }
 
+            if (isStreamingEnabled()) {
+                await streamingProcessor.generate();
+                streamingProcessor = null;
+            }
+
             function onSuccess(data) {
                 tokens_already_generated += this_amount_gen;			// add new gen amt to any prev gen counter..
 
diff --git a/public/scripts/poe.js b/public/scripts/poe.js
index 820562b4a..6517e55f4 100644
--- a/public/scripts/poe.js
+++ b/public/scripts/poe.js
@@ -43,6 +43,7 @@ const poe_settings = {
     auto_jailbreak: true,
     character_nudge: true,
     auto_purge: true,
+    streaming: false,
 };
 
 let auto_jailbroken = false;
@@ -61,6 +62,7 @@ function loadPoeSettings(settings) {
     $('#poe_character_nudge').prop('checked', poe_settings.character_nudge);
     $('#poe_auto_jailbreak').prop('checked', poe_settings.auto_jailbreak);
     $('#poe_auto_purge').prop('checked', poe_settings.auto_purge);
+    $('#poe_streaming').prop('checked', poe_settings.streaming);
     $('#poe_token').val(poe_settings.token ?? '');
     selectBot();
 }
@@ -94,7 +96,7 @@ async function generatePoe(finalPrompt) {
 
     if (poe_settings.auto_jailbreak && !auto_jailbroken) {
         for (let retryNumber = 0; retryNumber < MAX_RETRIES_FOR_ACTIVATION; retryNumber++) {
-            const reply = await sendMessage(poe_settings.jailbreak_message);
+            const reply = await sendMessage(poe_settings.jailbreak_message, false);
 
             if (reply.toLowerCase().includes(poe_settings.jailbreak_response.toLowerCase())) {
                 auto_jailbroken = true;
@@ -115,7 +117,7 @@ async function generatePoe(finalPrompt) {
         finalPrompt += nudge;
     }
 
-    const reply = await sendMessage(finalPrompt);
+    const reply = await sendMessage(finalPrompt, true);
     got_reply = true;
     return reply;
 }
@@ -139,10 +141,11 @@ async function purgeConversation(count = -1) {
     return response.ok;
 }
 
-async function sendMessage(prompt) {
+async function sendMessage(prompt, withStreaming) {
     const body = JSON.stringify({
         bot: poe_settings.bot,
         token: poe_settings.token,
+        streaming: withStreaming && poe_settings.streaming,
         prompt,
     });
 
@@ -155,6 +158,25 @@ async function sendMessage(prompt) {
         method: 'POST',
     });
 
+    if (withStreaming && poe_settings.streaming) {
+        return async function* streamData() {
+            const decoder = new TextDecoder();
+            const reader = response.body.getReader();
+            let getMessage = '';
+            while (true) {
+                const { done, value } = await reader.read();
+                let response = decoder.decode(value);
+                getMessage += response;
+
+                if (done) {
+                    return;
+                }
+
+                yield getMessage;
+            }
+        }
+    }
+
     try {
         if (response.ok) {
             const data = await response.json();
@@ -261,6 +283,11 @@ function onCharacterNudgeMessageInput() {
     saveSettingsDebounced();
 }
 
+function onStreamingInput() {
+    poe_settings.streaming = !!$(this).prop('checked');
+    saveSettingsDebounced();
+}
+
 $('document').ready(function () {
     $('#poe_token').on('input', onTokenInput);
     $('#poe_bots').on('change', onBotChange);
@@ -271,4 +298,5 @@ $('document').ready(function () {
     $('#poe_auto_jailbreak').on('input', onAutoJailbreakInput);
     $('#poe_character_nudge').on('input', onCharacterNudgeInput);
     $('#poe_nudge_text').on('input', onCharacterNudgeMessageInput);
+    $('#poe_streaming').on('input', onStreamingInput);
 });
\ No newline at end of file
diff --git a/server.js b/server.js
index c2d3d9c6f..8caba2a14 100644
--- a/server.js
+++ b/server.js
@@ -1776,22 +1776,39 @@ app.post('/generate_poe', jsonParser, async (request, response) => {
     const token = request.body.token;
     const prompt = request.body.prompt;
     const bot = request.body.bot ?? POE_DEFAULT_BOT;
+    const streaming = request.body.streaming ?? false;
 
     try {
         const client = await getPoeClient(token);
 
-        let reply;
-        for await (const mes of client.send_message(bot, prompt)) {
-            reply = mes.text;
+        if (streaming) {
+            response.writeHead(200, {
+                'Transfer-Encoding': 'chunked',
+                'Cache-Control': 'no-transform',
+            });
+          
+            let reply = '';
+            for await (const mes of client.send_message(bot, prompt)) {
+                let newText = mes.text.substring(reply.length);
+                reply = mes.text;
+                response.write(newText);
+            }
+            console.log(reply);
+            client.disconnect_ws();
+            response.end();
+        }
+        else {
+            let reply;
+            for await (const mes of client.send_message(bot, prompt)) {
+                reply = mes.text;
+            }
+            console.log(reply);
+            client.disconnect_ws();
+            return response.send({ 'reply': reply });
         }
-
-        console.log(reply);
-
-        client.disconnect_ws();
-
-        return response.send({ 'reply': reply });
     }
-    catch {
+    catch (error) {
+        console.error(error);
         return response.sendStatus(500);
     }
 });