feat: add token counter and model search

Added /api/count_tokens endpoint to count tokens using model tokenizer
Implemented token counter in chat UI
Added model search functionality
Updated .gitignore to include .vscode/

.gitignore

@@ -5,4 +5,5 @@ build/
 __pycache__/
 .idea
 venv
-dist
+dist
+.vscode/

server.py

@@ -250,6 +250,25 @@ def api_generate():
         if verbose: print("->", result)
         return result
 
+@app.route("/api/count_tokens", methods=['POST'])
+def api_count_tokens():
+    global api_lock, verbose
+    if verbose: print("/api/count_tokens")
+    with api_lock:
+        data = request.get_json()
+        if verbose: print("<-", data)
+        model = get_loaded_model()
+        if model is None:
+            # If no model is loaded, return 0 tokens
+            result = { "result": "ok", "token_count": 0 }
+        else:
+            # Use the model's tokenizer to get actual token count
+            tokenizer = model.tokenizer
+            tokens = tokenizer.encode(data["text"])
+            result = { "result": "ok", "token_count": tokens.shape[-1] }
+        if verbose: print("->", result)
+        return json.dumps(result) + "\n"
+
 @app.route("/api/cancel_generate")
 def api_cancel_generate():
     global api_lock_cancel, verbose
@@ -467,4 +486,3 @@ if browser_start:
     print(f" -- Opening UI in default web browser")
 
 serve(app, host = host, port = port, threads = 8)
-

static/chat.css

@@ -131,6 +131,40 @@
     color: var(--textcolor-dim);
 }
 
+.token-counter {
+    position: absolute;
+    bottom: 18px;
+    right: 120px;
+    font-size: var(--font-size-small);
+    color: var(--textcolor-text);
+    font-family: var(--font-family);
+    user-select: none;
+    pointer-events: none;
+    background-color: var(--background-color-control);
+    padding: 2px 6px;
+    border-radius: 4px;
+    opacity: 0.8;
+    z-index: 1;
+    transition: opacity 0.2s ease, color 0.2s ease, background-color 0.2s ease;
+}
+
+.session-input:focus ~ .token-counter,
+.session-input:hover ~ .token-counter {
+    opacity: 1;
+    color: var(--textcolor-text);
+    background-color: var(--background-color-control);
+    filter: brightness(var(--select-brightness));
+}
+
+.session-input:hover ~ .token-counter {
+    filter: brightness(var(--hover-brightness));
+}
+
+
+.session-input-surround {
+    position: relative;  /* Ensure absolute positioning works correctly */
+}
+
 .session-block {
     padding-bottom: 10px;
     padding-right: 10px;

static/chat.js

@@ -234,8 +234,26 @@ class SessionView {
         this.items = new Map();
     }
 
+    async countTokens(text) {
+        // Get real token count from server
+        try {
+            const response = await fetch("/api/count_tokens", {
+                method: "POST",
+                headers: { "Content-Type": "application/json" },
+                body: JSON.stringify({ text: text })
+            });
+            const data = await response.json();
+            return data.token_count;
+        } catch (error) {
+            console.error('Error counting tokens:', error);
+            // Fallback to simple counting if API fails
+            return text.trim() === '' ? 0 : text.trim().split(/\s+/).length;
+        }
+    }
+
     createInputField() {
         let sdiv = util.newVFlex();
+        sdiv.style.position = 'relative'; // For absolute positioning of counter
 
         let div = document.createElement("textarea");
         div.className = "session-input";
@@ -251,7 +269,26 @@ class SessionView {
                 }
             }
         });
-        div.addEventListener('input', () => { this.inputFieldAutogrow(); });
+
+const debounce = (func, delay) => {
+    let timeoutId;
+    return (...args) => {
+        clearTimeout(timeoutId);
+        timeoutId = setTimeout(() => func.apply(this, args), delay);
+    };
+};
+
+div.addEventListener('input', debounce(async () => {
+    this.inputFieldAutogrow();
+    const tokens = await this.countTokens(div.value);
+    tokenCounter.textContent = tokens === 1 ? "1 token" : `${tokens} tokens`;
+}, 700)); // Delay in ms
+
+        // Create token counter after the textarea
+        let tokenCounter = util.newDiv(null, "token-counter");
+        tokenCounter.textContent = "0 token";
+        sdiv.appendChild(div);
+        sdiv.appendChild(tokenCounter);
 
         this.inputButton = new controls.Button("⏵ Chat", () => { this.submitInput() }, "session-input-button");
         this.cancelButton = new controls.Button("⏹ Stop", () => { this.cancelGen() }, "session-input-button");
@@ -259,7 +296,6 @@ class SessionView {
         this.cancelButton.setHidden(true);
         this.inputButton.refresh();
         this.cancelButton.refresh();
-        sdiv.appendChild(div);
         sdiv.appendChild(this.inputButton.element);
         sdiv.appendChild(this.cancelButton.element);
         return sdiv;
@@ -341,6 +377,10 @@ class SessionView {
         this.sessionInput.value = "";
         this.inputFieldAutogrow();
         this.scrollToBottom();
+        
+        // Reset token counter
+        const tokenCounter = this.element.querySelector('.token-counter');
+        if (tokenCounter) tokenCounter.textContent = "0 token";
 
         if (!this.sessionID || this.sessionID == "new") {
             if (input && input != "") {

static/models.css

@@ -4,6 +4,31 @@
     flex-grow: 1;
 }
 
+.model-search-container {
+    display: flex;
+    padding: 10px;
+    padding-bottom: 5px;
+}
+
+.model-search-box {
+    flex-grow: 1;
+    background-color: var(--textbox-background);
+    color: var(--textcolor-text);
+    border-radius: 5px;
+    border: 1px solid var(--textbox-border);
+    padding: 5px;
+    font-size: var(--font-size-medium);
+}
+
+.model-search-box::placeholder {
+    color: var(--textcolor-dim);
+}
+
+.model-search-box:focus {
+    outline: none;
+    filter: brightness(var(--select-brightness));
+}
+
 .model-list {
     background-color: var(--background-color-body);
     display: flex;
@@ -12,7 +37,7 @@
     min-width: 250px;
     flex-grow: 0;
     padding: 0px;
-    padding-top: 10px;
+    padding-top: 0px;
     padding-left: 10px;
     padding-right: 10px;
     height: calc(100vh - 56px);
@@ -182,5 +207,8 @@
     justify-content: end;
 }
 
-
+.model-search-container .linkbutton {
-
+    margin-left: -25px;
+    padding-top: 5px;
+    z-index: 10;
+}

static/models.js

@@ -12,10 +12,12 @@ export class Models {
         this.page.appendChild(layout);
 
         let layout_l = util.newHFlex();
+        this.searchContainer = util.newDiv(null, "model-search-container");
         this.modelList = util.newDiv(null, "model-list");
         this.modelView = util.newDiv(null, "model-view");
         let panel = util.newDiv(null, "model-list-controls");
         layout.appendChild(layout_l);
+        layout_l.appendChild(this.searchContainer);
         layout_l.appendChild(this.modelList);
         layout_l.appendChild(panel);
         layout.appendChild(this.modelView);
@@ -23,11 +25,16 @@ export class Models {
         this.removeButton = new controls.LinkButton("✖ Remove model", "✖ Confirm", () => { this.removeModel(this.lastModelUUID); });
         panel.appendChild(this.removeButton.element);
 
+        this.searchBox = null;
+        this.searchState = "";
         this.items = new Map();
         this.labels = new Map();
         this.currentView = null;
 
+
         this.lastModelUUID = null;
+
+        this.createSearchBox();
     }
 
     onEnter() {
@@ -39,13 +46,35 @@ export class Models {
         });
     }
 
-    populateModelList(response) {
+   createSearchBox() {
+        this.searchBox = new controls.LabelTextboxButton(null, null, "model-search-box", "Search models...", this, "searchState", null, () => {}, null, "✖", () => {
+            this.searchState = "";
+            this.searchBox.tb.value = "";
+            this.populateModelList();
+        });
+        this.searchBox.tb.addEventListener("input", () => {
+            this.searchState = this.searchBox.tb.value;
+            this.populateModelList();
+        });
+        this.searchContainer.appendChild(this.searchBox.element);
+    }
+
+    populateModelList(response = null) {
+        if (response) {
+            this.modelData = response.models;
+        }
 
         this.modelList.innerHTML = "";
 
-        for (let model_uuid in response.models)
+        for (let model_uuid in this.modelData) {
-            if (response.models.hasOwnProperty(model_uuid))
+            if (this.modelData.hasOwnProperty(model_uuid)) {
-                this.addModel(response.models[model_uuid], model_uuid);
+                const name = this.modelData[model_uuid];
+                if (this.searchState && !name.toLowerCase().includes(this.searchState.toLowerCase())) {
+                    continue;
+                }
+                this.addModel(name, model_uuid);
+            }
+        }
 
         this.addModel("New model", "new");
         let m = this.lastModelUUID ? this.lastModelUUID : "new";