Use a WASM-based tiktoken implementation (RooCodeInc#2859)

* Use a WASM-based tiktoken implementation

* Clean up imports

Author: cte

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

@@ -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",

package.json

@@ -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,21 +34,23 @@ 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
 		for (const block of content) {
 			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
 				}
 			} 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/api/providers/base-provider.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 {