Implement TicToken library in frontend (stitionai#283)

ui/package.json

@@ -16,12 +16,14 @@
     "postcss-load-config": "^5.0.2",
     "svelte": "^4.2.7",
     "tailwindcss": "^3.3.6",
-    "vite": "^5.0.3"
+    "vite": "^5.0.3",
+    "vite-plugin-wasm": "^3.3.0"
   },
   "type": "module",
   "dependencies": {
-    "socket.io-client": "^4.7.5",
     "clsx": "^2.1.0",
+    "socket.io-client": "^4.7.5",
+    "tiktoken": "^1.0.13",
     "xterm": "^5.3.0",
     "xterm-addon-fit": "^0.8.0"
   }

ui/src/lib/components/MessageInput.svelte

@@ -1,6 +1,7 @@
 <script>
-  import { API_BASE_URL,  socket } from "$lib/api";
+  import { socket } from "$lib/api";
   import { agentState, messages } from "$lib/store";
+  import { calculateTokens } from "$lib/token";
 
   let isAgentActive = false;
 
@@ -13,7 +14,7 @@
     const projectName = localStorage.getItem("selectedProject");
     const selectedModel = localStorage.getItem("selectedModel");
     const serachEngine = localStorage.getItem("selectedSearchEngine");
-    
+
     if (!projectName) {
       alert("Please select a project first!");
       return;
@@ -30,41 +31,26 @@
           message: messageInput,
           base_model: selectedModel,
           project_name: projectName,
-          search_engine: serachEngine
+          search_engine: serachEngine,
         });
       } else {
         socket.emit("user-message", { 
           action: "continue",
           message: messageInput,
           base_model: selectedModel,
           project_name: projectName,
-          search_engine: serachEngine
-         });
+          search_engine: serachEngine,
+        });
       }
       messageInput = "";
     }
   }
 
-  function calculateTokens(event) {
+  function setTokenSize(event) {
     const prompt = event.target.value;
-    fetch(`${API_BASE_URL}/api/calculate-tokens`, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-      },
-      body: JSON.stringify({ prompt }),
-    })
-      .then((response) => response.json())
-      .then((data) => {
-        document.querySelector(".token-count").textContent =
-          `${data.token_usage} tokens`;
-      })
-      .catch((error) => {
-        console.error("Error:", error);
-      });
+    let tokens = calculateTokens(prompt);
+    document.querySelector(".token-count").textContent = `${tokens} tokens`;
   }
-
-
 </script>
 
 <div class="expandable-input relative">
@@ -73,7 +59,7 @@
     class="w-full p-2 border-2 rounded-lg pr-20"
     placeholder="Type your message..."
     bind:value={messageInput}
-    on:input={calculateTokens}
+    on:input={setTokenSize}
     on:keydown={(e) => {
       if (e.key === "Enter" && !e.shiftKey) {
         e.preventDefault();

ui/src/lib/token.js

@@ -0,0 +1,13 @@
+import { Tiktoken } from "tiktoken/lite";
+import cl100k_base from "tiktoken/encoders/cl100k_base.json";
+
+const encoding = new Tiktoken(
+  cl100k_base.bpe_ranks,
+  cl100k_base.special_tokens,
+  cl100k_base.pat_str
+);
+
+export function calculateTokens(text) {
+  const tokens = encoding.encode(text);
+  return tokens.length;
+}

ui/vite.config.js

@@ -1,9 +1,10 @@
-import { sveltekit } from '@sveltejs/kit/vite';
-import { defineConfig } from 'vite';
+import { sveltekit } from "@sveltejs/kit/vite";
+import { defineConfig } from "vite";
+import wasm from "vite-plugin-wasm";
 
 export default defineConfig({
-	plugins: [sveltekit()],
-	server: {
-	    port: 3000,
-	},
+  plugins: [sveltekit(), wasm()],
+  server: {
+    port: 3000,
+  },
 });