Use a WASM-based tiktoken implementation (#2859)

* Use a WASM-based tiktoken implementation

* Clean up imports

esbuild.js

@@ -29,15 +29,23 @@ const copyWasmFiles = {
 	name: "copy-wasm-files",
 	setup(build) {
 		build.onEnd(() => {
-			// tree sitter
-			const sourceDir = path.join(__dirname, "node_modules", "web-tree-sitter")
-			const targetDir = path.join(__dirname, "dist")
-
-			// Copy tree-sitter.wasm
-			fs.copyFileSync(path.join(sourceDir, "tree-sitter.wasm"), path.join(targetDir, "tree-sitter.wasm"))
-
-			// Copy language-specific WASM files
-			const languageWasmDir = path.join(__dirname, "node_modules", "tree-sitter-wasms", "out")
+			const nodeModulesDir = path.join(__dirname, "node_modules")
+			const distDir = path.join(__dirname, "dist")
+
+			// tiktoken
+			fs.copyFileSync(
+				path.join(nodeModulesDir, "tiktoken", "tiktoken_bg.wasm"),
+				path.join(distDir, "tiktoken_bg.wasm"),
+			)
+
+			// tree-sitter WASM
+			fs.copyFileSync(
+				path.join(nodeModulesDir, "web-tree-sitter", "tree-sitter.wasm"),
+				path.join(distDir, "tree-sitter.wasm"),
+			)
+
+			// language-specific tree-sitter WASMs
+			const languageWasmDir = path.join(nodeModulesDir, "tree-sitter-wasms", "out")
 			const languages = [
 				"typescript",
 				"tsx",
@@ -57,7 +65,7 @@ const copyWasmFiles = {
 
 			languages.forEach((lang) => {
 				const filename = `tree-sitter-${lang}.wasm`
-				fs.copyFileSync(path.join(languageWasmDir, filename), path.join(targetDir, filename))
+				fs.copyFileSync(path.join(languageWasmDir, filename), path.join(distDir, filename))
 			})
 		})
 	},

package-lock.json

@@ -37,7 +37,6 @@
 				"get-folder-size": "^5.0.0",
 				"i18next": "^24.2.2",
 				"isbinaryfile": "^5.0.2",
-				"js-tiktoken": "^1.0.19",
 				"mammoth": "^1.8.0",
 				"monaco-vscode-textmate-theme-converter": "^0.1.7",
 				"node-cache": "^5.1.2",
@@ -59,6 +58,7 @@
 				"string-similarity": "^4.0.4",
 				"strip-ansi": "^7.1.0",
 				"strip-bom": "^5.0.0",
+				"tiktoken": "^1.0.21",
 				"tmp": "^0.2.3",
 				"tree-sitter-wasms": "^0.1.11",
 				"turndown": "^7.2.0",
@@ -15495,14 +15495,6 @@
 				"node": ">=1.0.0"
 			}
 		},
-		"node_modules/js-tiktoken": {
-			"version": "1.0.19",
-			"resolved": "https://registry.npmjs.org/js-tiktoken/-/js-tiktoken-1.0.19.tgz",
-			"integrity": "sha512-XC63YQeEcS47Y53gg950xiZ4IWmkfMe4p2V9OSaBt26q+p47WHn18izuXzSclCI73B7yGqtfRsT6jcZQI0y08g==",
-			"dependencies": {
-				"base64-js": "^1.5.1"
-			}
-		},
 		"node_modules/js-tokens": {
 			"version": "4.0.0",
 			"resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz",
@@ -20048,6 +20040,12 @@
 				"xtend": "~4.0.1"
 			}
 		},
+		"node_modules/tiktoken": {
+			"version": "1.0.21",
+			"resolved": "https://registry.npmjs.org/tiktoken/-/tiktoken-1.0.21.tgz",
+			"integrity": "sha512-/kqtlepLMptX0OgbYD9aMYbM7EFrMZCL7EoHM8Psmg2FuhXoo/bH64KqOiZGGwa6oS9TPdSEDKBnV2LuB8+5vQ==",
+			"license": "MIT"
+		},
 		"node_modules/tinyexec": {
 			"version": "0.3.2",
 			"resolved": "https://registry.npmjs.org/tinyexec/-/tinyexec-0.3.2.tgz",

package.json

@@ -429,7 +429,6 @@
 		"get-folder-size": "^5.0.0",
 		"i18next": "^24.2.2",
 		"isbinaryfile": "^5.0.2",
-		"js-tiktoken": "^1.0.19",
 		"mammoth": "^1.8.0",
 		"monaco-vscode-textmate-theme-converter": "^0.1.7",
 		"node-cache": "^5.1.2",
@@ -451,6 +450,7 @@
 		"string-similarity": "^4.0.4",
 		"strip-ansi": "^7.1.0",
 		"strip-bom": "^5.0.0",
+		"tiktoken": "^1.0.21",
 		"tmp": "^0.2.3",
 		"tree-sitter-wasms": "^0.1.11",
 		"turndown": "^7.2.0",

src/api/providers/base-provider.ts

@@ -2,8 +2,8 @@ import { Anthropic } from "@anthropic-ai/sdk"
 import { ApiHandler } from ".."
 import { ModelInfo } from "../../shared/api"
 import { ApiStream } from "../transform/stream"
-import { Tiktoken } from "js-tiktoken/lite"
-import o200kBase from "js-tiktoken/ranks/o200k_base"
+import { Tiktoken } from "tiktoken/lite"
+import o200kBase from "tiktoken/encoders/o200k_base"
 
 // Reuse the fudge factor used in the original code
 const TOKEN_FUDGE_FACTOR = 1.5
@@ -34,7 +34,7 @@ export abstract class BaseProvider implements ApiHandler {
 
 		// Lazily create and cache the encoder if it doesn't exist
 		if (!this.encoder) {
-			this.encoder = new Tiktoken(o200kBase)
+			this.encoder = new Tiktoken(o200kBase.bpe_ranks, o200kBase.special_tokens, o200kBase.pat_str)
 		}
 
 		// Process each content block using the cached encoder
@@ -42,6 +42,7 @@ export abstract class BaseProvider implements ApiHandler {
 			if (block.type === "text") {
 				// Use tiktoken for text token counting
 				const text = block.text || ""
+
 				if (text.length > 0) {
 					const tokens = this.encoder.encode(text)
 					totalTokens += tokens.length
@@ -49,6 +50,7 @@ export abstract class BaseProvider implements ApiHandler {
 			} else if (block.type === "image") {
 				// For images, calculate based on data size
 				const imageSource = block.source
+
 				if (imageSource && typeof imageSource === "object" && "data" in imageSource) {
 					const base64Data = imageSource.data as string
 					totalTokens += Math.ceil(Math.sqrt(base64Data.length))

src/core/sliding-window/__tests__/sliding-window.test.ts

@@ -3,10 +3,13 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 
 import { ModelInfo } from "../../../shared/api"
-import { ApiHandler } from "../../../api"
 import { BaseProvider } from "../../../api/providers/base-provider"
-import { TOKEN_BUFFER_PERCENTAGE } from "../index"
-import { estimateTokenCount, truncateConversation, truncateConversationIfNeeded } from "../index"
+import {
+	TOKEN_BUFFER_PERCENTAGE,
+	estimateTokenCount,
+	truncateConversation,
+	truncateConversationIfNeeded,
+} from "../index"
 
 // Create a mock ApiHandler for testing
 class MockApiHandler extends BaseProvider {