From 06db54730877052b38538aa3be98057065ee07a8 Mon Sep 17 00:00:00 2001 From: Chris Estreich Date: Fri, 25 Apr 2025 10:32:35 -0700 Subject: [PATCH] Use a WASM-based tiktoken implementation (#2859) * Use a WASM-based tiktoken implementation * Clean up imports --- esbuild.js | 24 ++++++++++++------- package-lock.json | 16 ++++++------- package.json | 2 +- src/api/providers/base-provider.ts | 8 ++++--- .../__tests__/sliding-window.test.ts | 9 ++++--- 5 files changed, 35 insertions(+), 24 deletions(-) diff --git a/esbuild.js b/esbuild.js index 6fc0c24729..c6a555f5a5 100644 --- a/esbuild.js +++ b/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") + const nodeModulesDir = path.join(__dirname, "node_modules") + const distDir = path.join(__dirname, "dist") - // Copy tree-sitter.wasm - fs.copyFileSync(path.join(sourceDir, "tree-sitter.wasm"), path.join(targetDir, "tree-sitter.wasm")) + // tiktoken + fs.copyFileSync( + path.join(nodeModulesDir, "tiktoken", "tiktoken_bg.wasm"), + path.join(distDir, "tiktoken_bg.wasm"), + ) - // Copy language-specific WASM files - const languageWasmDir = path.join(__dirname, "node_modules", "tree-sitter-wasms", "out") + // 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)) }) }) }, diff --git a/package-lock.json b/package-lock.json index 5d679b8e2f..48413ebf90 100644 --- a/package-lock.json +++ b/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", diff --git a/package.json b/package.json index c864851294..d7c3de86ed 100644 --- a/package.json +++ b/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", diff --git a/src/api/providers/base-provider.ts b/src/api/providers/base-provider.ts index 34156e4adf..057b48661d 100644 --- a/src/api/providers/base-provider.ts +++ b/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)) diff --git a/src/core/sliding-window/__tests__/sliding-window.test.ts b/src/core/sliding-window/__tests__/sliding-window.test.ts index 532d00067a..fb7bd9c227 100644 --- a/src/core/sliding-window/__tests__/sliding-window.test.ts +++ b/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 {