Add token usage indicator

aiserver.py

@@ -3608,6 +3608,37 @@ def get_message(msg):
         emit('from_server', {'cmd': 'set_debug', 'data': msg['data']}, broadcast=True)
         if vars.debug:
             send_debug()
+    elif(msg['cmd'] == 'getfieldbudget'):
+        unencoded = msg["data"]["unencoded"]
+        field = msg["data"]["field"]
+
+        # Tokenizer may be undefined here when a model has not been chosen.
+        if "tokenizer" not in globals():
+            # We don't have a tokenizer, just return nulls.
+            emit(
+                'from_server',
+                {'cmd': 'showfieldbudget', 'data': {"length": None, "max": None, "field": field}},
+                broadcast=True
+            )
+            return
+
+        header_length = len(tokenizer._koboldai_header)
+        max_tokens = vars.max_length - header_length - vars.sp_length - vars.genamt
+
+        if not unencoded:
+            # Unencoded is empty, just return 0
+            emit(
+                'from_server',
+                {'cmd': 'showfieldbudget', 'data': {"length": 0, "max": max_tokens, "field": field}},
+                broadcast=True
+            )
+        else:
+            tokens_length = len(tokenizer.encode(unencoded))
+            emit(
+                'from_server',
+                {'cmd': 'showfieldbudget', 'data': {"length": tokens_length, "max": max_tokens, "field": field}},
+                broadcast=True
+            )
 
 #==================================================================#
 #  Send userscripts list to client

static/application.js

@@ -512,6 +512,16 @@ function addWiLine(ob) {
 		$(".wisortable-excluded-dynamic").removeClass("wisortable-excluded-dynamic");
 		$(this).parent().css("max-height", "").find(".wicomment").find(".form-control").css("max-height", "");
 	});
+
+	for (const wientry of document.getElementsByClassName("wientry")) {
+		// If we are uninitialized, skip.
+		if ($(wientry).closest(".wilistitem-uninitialized").length) continue;
+
+		// add() will not add if the class is already present
+		wientry.classList.add("tokens-counted");
+	}
+
+	registerTokenCounters();
 }
 
 function addWiFolder(uid, ob) {
@@ -835,6 +845,7 @@ function exitMemoryMode() {
 	button_actmem.html("Memory");
 	show([button_actback, button_actfwd, button_actretry, button_actwi]);
 	input_text.val("");
+	updateInputBudget(input_text[0]);
 	// Hide Author's Note field
 	anote_menu.slideUp("fast");
 }
@@ -2139,6 +2150,31 @@ function interpolateRGB(color0, color1, t) {
 	]
 }
 
+function updateInputBudget(inputElement) {
+	socket.send({"cmd": "getfieldbudget", "data": {"unencoded": inputElement.value, "field": inputElement.id}});
+}
+
+function registerTokenCounters() {
+	// Add token counters to all input containers with the class of "tokens-counted",
+	// if a token counter is not already a child of said container.
+	for (const el of document.getElementsByClassName("tokens-counted")) {
+		if (el.getElementsByClassName("input-token-usage").length) continue;
+
+		let span = document.createElement("span");
+		span.classList.add("input-token-usage");
+		span.innerText = "?/? Tokens";
+		el.appendChild(span);
+
+		let inputElement = el.querySelector("input, textarea");
+
+		inputElement.addEventListener("input", function() {
+			updateInputBudget(this);
+		});
+		
+		updateInputBudget(inputElement);
+	}
+}
+
 //=================================================================//
 //  READY/RUNTIME
 //=================================================================//
@@ -2481,6 +2517,7 @@ $(document).ready(function(){
 				memorytext = msg.data;
 				input_text.val(msg.data);
 			}
+			updateInputBudget(input_text[0]);
 		} else if(msg.cmd == "setmemory") {
 			memorytext = msg.data;
 			if(memorymode) {
@@ -2602,6 +2639,7 @@ $(document).ready(function(){
 		} else if(msg.cmd == "setanote") {
 			// Set contents of Author's Note field
 			anote_input.val(msg.data);
+			updateInputBudget(anote_input[0]);
 		} else if(msg.cmd == "setanotetemplate") {
 			// Set contents of Author's Note Template field
 			$("#anotetemplate").val(msg.data);
@@ -2913,6 +2951,12 @@ $(document).ready(function(){
 				opt.innerHTML = engine[1];
 				$("#oaimodel")[0].appendChild(opt);
 			}
+		} else if(msg.cmd == 'showfieldbudget') {
+			let inputElement = document.getElementById(msg.data.field);
+			let tokenBudgetElement = inputElement.parentNode.getElementsByClassName("input-token-usage")[0];
+			let tokenLength = msg.data.length ?? "?";
+			let tokenMax = msg.data.max ?? "?";
+			tokenBudgetElement.innerText = `${tokenLength}/${tokenMax} Tokens`;
 		}
 	});
 	
@@ -3381,6 +3425,11 @@ $(document).ready(function(){
 
 		if (handled) ev.preventDefault();
 	});
+
+	registerTokenCounters();
+
+	updateInputBudget(input_text[0]);
+
 });
 
 

static/custom.css

@@ -1695,3 +1695,30 @@ body.connected .popupfooter, .popupfooter.always-available {
 	overflow-x: auto;
 	white-space: nowrap;
 }
+
+.tokens-counted {
+	position: relative;
+}
+
+.input-token-usage {
+	color: white;
+	position: absolute;
+	font-size: 10px;
+	bottom: 2px;
+	right: 5px;
+
+	-webkit-user-select: none;
+	-moz-user-select: none;
+	-ms-user-select: none;
+	user-select: none;
+}
+
+/* Override needed here due to the 10px right padding on inputrowleft; add 10 px. */
+#inputrowleft > .input-token-usage {
+	right: 15px;
+	bottom: 1px;
+}
+
+.wientry > .input-token-usage {
+	bottom: 8px;
+}

templates/index.html

@@ -157,7 +157,7 @@
 				<div id="inputrowmode">
 					<button type="button" class="btn btn-secondary hidden" id="btnmode">Mode:<br/><b id="btnmode_label">Story</b></button>
 				</div>
-				<div id="inputrowleft">
+				<div id="inputrowleft" class="tokens-counted">
 					<textarea class="form-control" id="input_text" placeholder="Enter text here"></textarea>
 				</div>
 				<div id="inputrowright">
@@ -170,7 +170,7 @@
 						<div class="anotelabel no-padding">
 							Author's Note
 						</div>
-						<div class="anotefield">
+						<div class="anotefield tokens-counted">
 							<textarea class="form-control" placeholder="Author's Note" id="anoteinput"></textarea>
 						</div>
 					</div>