From ac2fceeddfd309a0260da868688fccb3fbf1d511 Mon Sep 17 00:00:00 2001
From: Cohee <titan0599@gmail.com>
Date: Mon, 12 Jun 2023 02:35:54 +0300
Subject: [PATCH] Add token counter for author's notes

---
 .../extensions/floating-prompt/index.js       |  21 ++-
 .../extensions/floating-prompt/style.css      | 128 +++++++++---------
 2 files changed, 87 insertions(+), 62 deletions(-)

diff --git a/public/scripts/extensions/floating-prompt/index.js b/public/scripts/extensions/floating-prompt/index.js
index a84769e53..57b3d3cdb 100644
--- a/public/scripts/extensions/floating-prompt/index.js
+++ b/public/scripts/extensions/floating-prompt/index.js
@@ -1,4 +1,11 @@
-import { chat_metadata, saveSettingsDebounced, this_chid } from "../../../script.js";
+import {
+    chat_metadata,
+    eventSource,
+    event_types,
+    getTokenCount,
+    saveSettingsDebounced,
+    this_chid,
+} from "../../../script.js";
 import { selected_group } from "../../group-chats.js";
 import { ModuleWorkerWrapper, extension_settings, getContext, saveMetadataDebounced } from "../../extensions.js";
 import { registerSlashCommand } from "../../slash-commands.js";
@@ -66,6 +73,7 @@ function setNotePositionCommand(_, text) {
 
 async function onExtensionFloatingPromptInput() {
     chat_metadata[metadata_keys.prompt] = $(this).val();
+    $('#extension_floating_prompt_token_counter').text(getTokenCount(chat_metadata[metadata_keys.prompt]));
     saveMetadataDebounced();
 }
 
@@ -93,6 +101,7 @@ async function onExtensionFloatingPositionInput(e) {
 
 function onExtensionFloatingDefaultInput() {
     extension_settings.note.default = $(this).val();
+    $('#extension_floating_default_token_counter').text(getTokenCount(extension_settings.note.default));
     saveSettingsDebounced();
 }
 
@@ -173,6 +182,13 @@ function onANMenuItemClick() {
     }
 }
 
+function onChatChanged() {
+    const tokenCounter1 = chat_metadata[metadata_keys.prompt] ? getTokenCount(chat_metadata[metadata_keys.prompt]) : 0;
+    const tokenCounter2 = extension_settings.note.default ? getTokenCount(extension_settings.note.default) : 0;
+    $('#extension_floating_prompt_token_counter').text(tokenCounter1);
+    $('#extension_floating_default_token_counter').text(tokenCounter2);
+}
+
 (function () {
     function addExtensionsSettings() {
         const settingsHtml = `
@@ -195,6 +211,7 @@ function onANMenuItemClick() {
                     </small>
 
                     <textarea id="extension_floating_prompt" class="text_pole" rows="8" maxlength="10000"></textarea>
+                    <div class="extension_token_counter">Tokens: <span id="extension_floating_prompt_token_counter">0</small></div>
 
                     <div class="floating_prompt_radio_group">
                         <label>
@@ -228,6 +245,7 @@ function onANMenuItemClick() {
 
                         <textarea id="extension_floating_default" class="text_pole" rows="8" maxlength="10000"
                         placeholder="Example:\n[Scenario: wacky adventures; Genre: romantic comedy; Style: verbose, creative]"></textarea>
+                        <div class="extension_token_counter">Tokens: <span id="extension_floating_default_token_counter">0</small></div>
                     </div>
                 </div>
             </div>
@@ -265,4 +283,5 @@ function onANMenuItemClick() {
     registerSlashCommand('depth', setNoteDepthCommand, [], "<span class='monospace'>(number)</span> – sets an author's note depth for in-chat positioning", true, true);
     registerSlashCommand('freq', setNoteIntervalCommand, ['interval'], "<span class='monospace'>(number)</span> – sets an author's note insertion frequency", true, true);
     registerSlashCommand('pos', setNotePositionCommand, ['position'], "(<span class='monospace'>chat</span> or <span class='monospace'>scenario</span>) – sets an author's note position", true, true);
+    eventSource.on(event_types.CHAT_CHANGED, onChatChanged);
 })();
diff --git a/public/scripts/extensions/floating-prompt/style.css b/public/scripts/extensions/floating-prompt/style.css
index 8474ba1bf..229288c51 100644
--- a/public/scripts/extensions/floating-prompt/style.css
+++ b/public/scripts/extensions/floating-prompt/style.css
@@ -1,61 +1,67 @@
-#floatingPrompt {
-    overflow-y: auto;
-    max-width: 90svw;
-    max-height: 90svh;
-    min-width: 100px;
-    min-height: 100px;
-    border-radius: 10px;
-    border: 1px solid var(--white30a);
-    position: fixed;
-    padding: 10px;
-    padding-top: 25px;
-    display: none;
-    flex-direction: column;
-    box-shadow: 0 0 10px var(--black70a);
-    z-index: 3000;
-    left: 0;
-    top: 0;
-    margin: 0;
-    right: unset;
-    width: calc(((100svw - var(--sheldWidth)) / 2) - 1px);
-
-}
-
-.floating_prompt_radio_group {
-    display: flex;
-    flex-direction: column;
-}
-
-#extension_floating_counter {
-    font-weight: 600;
-    color: orange;
-}
-
-.floating_prompt_settings textarea {
-    font-size: calc(var(--mainFontSize) * 0.9);
-    line-height: 1.2;
-}
-
-#ANClose {
-    height: 15px;
-    aspect-ratio: 1 / 1;
-    font-size: 20px;
-    opacity: 0.5;
-    transition: all 250ms;
-}
-
-#ANClose:hover {
-    cursor: pointer;
-    opacity: 1;
-}
-
-.panelControlBar {
-    position: absolute;
-    right: 5px;
-    top: 5px;
-    margin-right: 5px;
-}
-
-#floatingPrompt .drag-grabber {
-    position: unset;
-}
\ No newline at end of file
+#floatingPrompt {
+    overflow-y: auto;
+    max-width: 90svw;
+    max-height: 90svh;
+    min-width: 100px;
+    min-height: 100px;
+    border-radius: 10px;
+    border: 1px solid var(--white30a);
+    position: fixed;
+    padding: 10px;
+    padding-top: 25px;
+    display: none;
+    flex-direction: column;
+    box-shadow: 0 0 10px var(--black70a);
+    z-index: 3000;
+    left: 0;
+    top: 0;
+    margin: 0;
+    right: unset;
+    width: calc(((100svw - var(--sheldWidth)) / 2) - 1px);
+
+}
+
+.floating_prompt_radio_group {
+    display: flex;
+    flex-direction: column;
+}
+
+#extension_floating_counter {
+    font-weight: 600;
+    color: orange;
+}
+
+.extension_token_counter {
+    font-size: calc(var(--mainFontSize) * 0.9);
+    width: 100%;
+    text-align: right;
+}
+
+.floating_prompt_settings textarea {
+    font-size: calc(var(--mainFontSize) * 0.9);
+    line-height: 1.2;
+}
+
+#ANClose {
+    height: 15px;
+    aspect-ratio: 1 / 1;
+    font-size: 20px;
+    opacity: 0.5;
+    transition: all 250ms;
+}
+
+#ANClose:hover {
+    cursor: pointer;
+    opacity: 1;
+}
+
+.panelControlBar {
+    position: absolute;
+    right: 5px;
+    top: 5px;
+    margin-right: 5px;
+}
+
+#floatingPrompt .drag-grabber {
+    position: unset;
+}