Skip to content

⚡ refactor: Replace tiktoken with ai-tokenizer#12175

Merged
danny-avila merged 11 commits into
devfrom
refactor/ai-tokenizer
Mar 11, 2026
Merged

⚡ refactor: Replace tiktoken with ai-tokenizer#12175
danny-avila merged 11 commits into
devfrom
refactor/ai-tokenizer

Conversation

@danny-avila

@danny-avila danny-avila commented Mar 11, 2026

Copy link
Copy Markdown
Owner

Summary

Replaced the WASM-based tiktoken dependency with ai-tokenizer, a pure-JS tokenizer that is 5–7x faster, removes the native binary requirement, and enables first-class Claude model tokenization support.

  • Rewrote the Tokenizer singleton in packages/api/src/utils/tokenizer.ts to use ai-tokenizer's count() API, backed by a two-entry encodingData map covering o200k_base and claude encodings.
  • Removed freeAndResetAllEncoders, resetTokenizersIfNecessary, and tokenizerCallsCount — memory management lifecycle methods that were required by tiktoken's WASM allocator but are unnecessary for a pure-JS tokenizer.
  • Added nested try/catch in getTokenCount so a double failure during error recovery logs and returns 0 rather than propagating an uncaught exception into buildMessages.
  • Added Claude model detection to AgentClient.getEncoding(), selecting the claude encoding for any model whose name includes "claude" and defaulting to o200k_base for all others.
  • Exported EncodingName as a public type from @librechat/api via packages/api/src/utils/index.ts for typed downstream consumers.
  • Updated countTokens() to use o200k_base instead of the legacy cl100k_base encoding, and updated its JSDoc to reflect the encoding used.
  • Removed the tiktokenModels export from tokens.ts, confirmed as dead code with no callers anywhere in the codebase.
  • Removed stale jest.mock('tiktoken') blocks from DALLE3.spec.js, DALLE3-proxy.spec.js, and samlStrategy.spec.js.
  • Updated all tests in tokenizer.spec.ts and text.spec.ts to use the new count()-based API, replacing tiktoken-specific assertions (encode/free/reset) with encoding-agnostic equivalents.

Change Type

  • New feature (non-breaking change which adds functionality)

Testing

Unit tests were rewritten alongside the implementation and cover the following:

  • getTokenizer returns a valid, cached AiTokenizer instance for both o200k_base and claude encodings, and defaults correctly when no encoding is specified.
  • getTokenCount returns a positive count for both encodings against a real text fixture ("Hello, world!").
  • Error recovery: getTokenizer is mocked to return a broken .count on the first call; the test asserts that getTokenCount logs the error, evicts the cache entry, and returns a valid non-zero count from the recreated tokenizer.
  • Singleton identity is asserted via dynamic re-import.
  • Integration benchmarks in text.spec.ts validate that processTextWithTokenLimit achieves ≥70% reduction in tokenizer call count versus the binary-search baseline, using the real ai-tokenizer implementation against 15k, 50k, and 120k token text fixtures.

Test Configuration:

cd packages/api && npx jest tokenizer
cd packages/api && npx jest text.spec

Checklist

  • My code adheres to this project's style guidelines
  • I have performed a self-review of my own code
  • I have commented in any complex areas of my code
  • My changes do not introduce new warnings
  • I have written tests demonstrating that my changes are effective or that my feature works
  • Local unit tests pass with my changes

- Added ai-tokenizer version 1.0.6 to package.json and package-lock.json across multiple packages.
- Removed tiktoken version 1.0.15 from package.json and package-lock.json in the same locations, streamlining dependency management.
Copilot AI review requested due to automatic review settings March 11, 2026 01:48

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR aims to swap tokenization dependencies in the API packages by adding ai-tokenizer and removing tiktoken, presumably to standardize/streamline token counting across the codebase.

Changes:

  • Add ai-tokenizer@^1.0.6 to api and packages/api.
  • Remove tiktoken@^1.0.15 from those packages.
  • Update root package-lock.json accordingly (adds ai-tokenizer, removes tiktoken entry).

Reviewed changes

Copilot reviewed 2 out of 3 changed files in this pull request and generated 2 comments.

File Description
packages/api/package.json Adds ai-tokenizer and removes tiktoken dependency for packages/api.
api/package.json Adds ai-tokenizer and removes tiktoken dependency for api.
package-lock.json Locks ai-tokenizer@1.0.6 and removes tiktoken from the resolved dependency tree.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

Comment thread packages/api/package.json
Comment thread api/package.json
@danny-avila danny-avila changed the title feat: replace js-tiktoken with ai-tokenizer refactor: replace js-tiktoken with ai-tokenizer Mar 11, 2026
- Added support for 'claude' encoding in the AgentClient class to improve model compatibility.
- Updated Tokenizer class to utilize 'ai-tokenizer' for both 'o200k_base' and 'claude' encodings, replacing the previous 'tiktoken' dependency.
- Refactored tests to reflect changes in tokenizer behavior and ensure accurate token counting for both encoding types.
- Removed deprecated references to 'tiktoken' and adjusted related tests for improved clarity and functionality.
@danny-avila danny-avila changed the title refactor: replace js-tiktoken with ai-tokenizer refactor: replace tiktoken with ai-tokenizer Mar 11, 2026
- Eliminated mock implementations of 'tiktoken' from DALLE3-related test files to streamline test setup and align with recent dependency updates.
- Adjusted related test structures to ensure compatibility with the new tokenizer implementation.
- Introduced a new method `getEncoding` in the AgentClient class to handle the specific BPE tokenizer for Claude models, ensuring compatibility with the distinct encoding requirements.
- Updated documentation to clarify the encoding logic for Claude and other models.
…tClient

- Clarified the return type of the getEncoding method to specify that it can return an EncodingName or undefined, enhancing code readability and type safety.
- Exported the EncodingName type for broader usage.
- Renamed encodingMap to encodingData for clarity.
- Improved error handling in getTokenCount method to ensure recovery attempts are logged and return 0 on failure.
- Updated countTokens function documentation to specify the use of 'o200k_base' encoding.
@danny-avila danny-avila changed the title refactor: replace tiktoken with ai-tokenizer ⚡ refactor: Replace tiktoken with ai-tokenizer Mar 11, 2026
- Updated the getEncoding method documentation to clarify the default behavior for non-Anthropic Claude models.
- Exported the EncodingName type separately from the Tokenizer module for improved clarity and usage.
- Adjusted test cases to handle smaller text sizes, changing scenarios from ~120k tokens to ~20k tokens for both the real tokenizer and countTokens functions.
- Updated token limits in tests to reflect new constraints, ensuring tests accurately assess performance and call reduction.
- Enhanced console log messages for clarity regarding token counts and reductions in the updated scenarios.
- Moved Tokenizer and countTokens exports to the tokenizer module for better organization.
- Adjusted imports in memory.ts to reflect the new structure, ensuring consistent usage across the codebase.
- Updated memory.test.ts to mock the Tokenizer from the correct module path, enhancing test accuracy.
- Introduced an async `initEncoding` method to preload tokenizers, improving performance and accuracy in token counting.
- Updated `getTokenCount` to handle uninitialized tokenizers more gracefully, ensuring proper recovery and logging on errors.
- Removed deprecated synchronous tokenizer retrieval, streamlining the overall tokenizer management process.
- Added `beforeAll` hooks to initialize tokenizers for 'o200k_base' and 'claude' encodings before running tests, ensuring proper setup.
- Updated tests to validate the loading of encodings and the correctness of token counts for both 'o200k_base' and 'claude'.
- Improved test structure to deduplicate concurrent initialization calls, enhancing performance and reliability.
@danny-avila danny-avila merged commit 9a5d7ea into dev Mar 11, 2026
9 checks passed
@danny-avila danny-avila deleted the refactor/ai-tokenizer branch March 11, 2026 03:14
mstlaur1 pushed a commit to mstlaur1/LibreChat that referenced this pull request Mar 25, 2026
* chore: Update dependencies by adding ai-tokenizer and removing tiktoken

- Added ai-tokenizer version 1.0.6 to package.json and package-lock.json across multiple packages.
- Removed tiktoken version 1.0.15 from package.json and package-lock.json in the same locations, streamlining dependency management.

* refactor: replace js-tiktoken with ai-tokenizer

- Added support for 'claude' encoding in the AgentClient class to improve model compatibility.
- Updated Tokenizer class to utilize 'ai-tokenizer' for both 'o200k_base' and 'claude' encodings, replacing the previous 'tiktoken' dependency.
- Refactored tests to reflect changes in tokenizer behavior and ensure accurate token counting for both encoding types.
- Removed deprecated references to 'tiktoken' and adjusted related tests for improved clarity and functionality.

* chore: remove tiktoken mocks from DALLE3 tests

- Eliminated mock implementations of 'tiktoken' from DALLE3-related test files to streamline test setup and align with recent dependency updates.
- Adjusted related test structures to ensure compatibility with the new tokenizer implementation.

* chore: Add distinct encoding support for Anthropic Claude models

- Introduced a new method `getEncoding` in the AgentClient class to handle the specific BPE tokenizer for Claude models, ensuring compatibility with the distinct encoding requirements.
- Updated documentation to clarify the encoding logic for Claude and other models.

* docs: Update return type documentation for getEncoding method in AgentClient

- Clarified the return type of the getEncoding method to specify that it can return an EncodingName or undefined, enhancing code readability and type safety.

* refactor: Tokenizer class and error handling

- Exported the EncodingName type for broader usage.
- Renamed encodingMap to encodingData for clarity.
- Improved error handling in getTokenCount method to ensure recovery attempts are logged and return 0 on failure.
- Updated countTokens function documentation to specify the use of 'o200k_base' encoding.

* refactor: Simplify encoding documentation and export type

- Updated the getEncoding method documentation to clarify the default behavior for non-Anthropic Claude models.
- Exported the EncodingName type separately from the Tokenizer module for improved clarity and usage.

* test: Update text processing tests for token limits

- Adjusted test cases to handle smaller text sizes, changing scenarios from ~120k tokens to ~20k tokens for both the real tokenizer and countTokens functions.
- Updated token limits in tests to reflect new constraints, ensuring tests accurately assess performance and call reduction.
- Enhanced console log messages for clarity regarding token counts and reductions in the updated scenarios.

* refactor: Update Tokenizer imports and exports

- Moved Tokenizer and countTokens exports to the tokenizer module for better organization.
- Adjusted imports in memory.ts to reflect the new structure, ensuring consistent usage across the codebase.
- Updated memory.test.ts to mock the Tokenizer from the correct module path, enhancing test accuracy.

* refactor: Tokenizer initialization and error handling

- Introduced an async `initEncoding` method to preload tokenizers, improving performance and accuracy in token counting.
- Updated `getTokenCount` to handle uninitialized tokenizers more gracefully, ensuring proper recovery and logging on errors.
- Removed deprecated synchronous tokenizer retrieval, streamlining the overall tokenizer management process.

* test: Enhance tokenizer tests with initialization and encoding checks

- Added `beforeAll` hooks to initialize tokenizers for 'o200k_base' and 'claude' encodings before running tests, ensuring proper setup.
- Updated tests to validate the loading of encodings and the correctness of token counts for both 'o200k_base' and 'claude'.
- Improved test structure to deduplicate concurrent initialization calls, enhancing performance and reliability.
jcbartle pushed a commit to jcbartle/LibreChat that referenced this pull request May 11, 2026
* chore: Update dependencies by adding ai-tokenizer and removing tiktoken

- Added ai-tokenizer version 1.0.6 to package.json and package-lock.json across multiple packages.
- Removed tiktoken version 1.0.15 from package.json and package-lock.json in the same locations, streamlining dependency management.

* refactor: replace js-tiktoken with ai-tokenizer

- Added support for 'claude' encoding in the AgentClient class to improve model compatibility.
- Updated Tokenizer class to utilize 'ai-tokenizer' for both 'o200k_base' and 'claude' encodings, replacing the previous 'tiktoken' dependency.
- Refactored tests to reflect changes in tokenizer behavior and ensure accurate token counting for both encoding types.
- Removed deprecated references to 'tiktoken' and adjusted related tests for improved clarity and functionality.

* chore: remove tiktoken mocks from DALLE3 tests

- Eliminated mock implementations of 'tiktoken' from DALLE3-related test files to streamline test setup and align with recent dependency updates.
- Adjusted related test structures to ensure compatibility with the new tokenizer implementation.

* chore: Add distinct encoding support for Anthropic Claude models

- Introduced a new method `getEncoding` in the AgentClient class to handle the specific BPE tokenizer for Claude models, ensuring compatibility with the distinct encoding requirements.
- Updated documentation to clarify the encoding logic for Claude and other models.

* docs: Update return type documentation for getEncoding method in AgentClient

- Clarified the return type of the getEncoding method to specify that it can return an EncodingName or undefined, enhancing code readability and type safety.

* refactor: Tokenizer class and error handling

- Exported the EncodingName type for broader usage.
- Renamed encodingMap to encodingData for clarity.
- Improved error handling in getTokenCount method to ensure recovery attempts are logged and return 0 on failure.
- Updated countTokens function documentation to specify the use of 'o200k_base' encoding.

* refactor: Simplify encoding documentation and export type

- Updated the getEncoding method documentation to clarify the default behavior for non-Anthropic Claude models.
- Exported the EncodingName type separately from the Tokenizer module for improved clarity and usage.

* test: Update text processing tests for token limits

- Adjusted test cases to handle smaller text sizes, changing scenarios from ~120k tokens to ~20k tokens for both the real tokenizer and countTokens functions.
- Updated token limits in tests to reflect new constraints, ensuring tests accurately assess performance and call reduction.
- Enhanced console log messages for clarity regarding token counts and reductions in the updated scenarios.

* refactor: Update Tokenizer imports and exports

- Moved Tokenizer and countTokens exports to the tokenizer module for better organization.
- Adjusted imports in memory.ts to reflect the new structure, ensuring consistent usage across the codebase.
- Updated memory.test.ts to mock the Tokenizer from the correct module path, enhancing test accuracy.

* refactor: Tokenizer initialization and error handling

- Introduced an async `initEncoding` method to preload tokenizers, improving performance and accuracy in token counting.
- Updated `getTokenCount` to handle uninitialized tokenizers more gracefully, ensuring proper recovery and logging on errors.
- Removed deprecated synchronous tokenizer retrieval, streamlining the overall tokenizer management process.

* test: Enhance tokenizer tests with initialization and encoding checks

- Added `beforeAll` hooks to initialize tokenizers for 'o200k_base' and 'claude' encodings before running tests, ensuring proper setup.
- Updated tests to validate the loading of encodings and the correctness of token counts for both 'o200k_base' and 'claude'.
- Improved test structure to deduplicate concurrent initialization calls, enhancing performance and reliability.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants