fix(rag): keep tiktoken token offsets aligned with the text for multi-byte characters

OS
2d ago
+57-2

PR Checklist

  • The commit message follows our guidelines
  • Tests for the changes have been added
  • Docs have been added / updated (not needed, no API change)
  • Changesets have been added

What is the current behavior?

createTikTokenizer().tokenize (packages/rag/src/utils/tokenizer.ts) decodes every token on its own to work out its length. A character that spans several tokens, which covers most CJK text and emoji, decodes to U+FFFD for each of those tokens, and U+FFFD is one UTF-16 unit long whatever the character really is. The offsets then drift ahead of the text.

TokenChunker and SentenceChunker cut the text with those offsets, so for Chinese, Japanese or emoji text:

  • the last token ends past the end of the text (520 on a 460-character Chinese string),
  • chunks hold more tokens than maxTokens,
  • chunks can end between the two halves of a surrogate pair.

Measured on the published @voltagent/rag 1.0.2 (same as main) with new TokenChunker().chunk(text, { maxTokens: 30, overlap: 0 }) and the tokens of each chunk counted again with countTokens:

textchunksover maxTokenswith a lone surrogatemost tokens in one chunk
Chinese, 460 chars1714036
Japanese, 520 chars1915039
emoji, 360 chars1812844
ASCII, 1800 chars140030

What is the new behavior?

The tokens of a character that is not finished yet stay at its start as zero-width tokens, and the whole character goes to the token that completes it. The number of tokens is unchanged, offsets are contiguous and end at text.length, and chunks end on whole characters. With the same inputs the most tokens in one chunk is 30 (Chinese), 31 (Japanese) and 32 (emoji), and no chunk has a lone surrogate. The extra one or two come from a character that straddles a chunk boundary: it goes whole to one chunk, and encoding it again on its own can cost a token or two more than that chunk's share.

A chunk whose only token is part of a character can now be empty (for example maxTokens: 1 on emoji); I left that as is.

Tests: offsets for Japanese, emoji and mixed scripts (contiguous, start at 0, end at text.length, slices rejoin to the text) and TokenChunker on Chinese and emoji (chunks rejoin to the text, content matches start/end, no lone surrogates, at most maxTokens + 2 tokens). The five new tests fail on main and pass with the change; the rest of the @voltagent/rag suite passes (52 tests), and biome is clean. tsc --noEmit reports the same 3 errors in late-chunker.spec.ts with and without this change.

Notes for reviewers

Not related to #1470, which fixes the character ranges SentenceChunker reports; this changes the offsets the tokenizer itself returns.


Summary by cubic

Fixes createTikTokenizer().tokenize so token offsets stay aligned with the text for multi-byte characters (most CJK text and emoji), which previously drifted and caused TokenChunker and SentenceChunker to produce chunks that exceeded maxTokens or split surrogate pairs.

  • Decoding each token on its own yielded U+FFFD per token, one UTF-16 unit long regardless of the actual character, so offsets ran ahead of the text.
  • Tokens of an unfinished character now stay at its start and the whole character goes to the token that completes it; token counts are unchanged and chunks end on whole characters.
  • A chunk whose only content is the partial tokens of one character can come back empty (for example maxTokens: 1 on emoji), left as is.
  • Adds tests for Japanese, emoji, and mixed scripts that fail on main.

Summary by CodeRabbit

  • Bug Fixes
    • Improved text chunking for Chinese, Japanese, accented text, emoji, and mixed scripts. Chunk boundaries now stay aligned with complete characters, preventing chunks from splitting multi-byte characters.
    • Improved token offsets so text ranges remain contiguous and cover the original input, including when characters require multiple tokens. Chunk sizes remain close to the configured limit, with allowance for a character that spans a boundary.