diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ca03b0..650987a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## [Unreleased] +## [0.4.1] - 2026-02-26 + +- Fix tiktoken-rs stack overflow crash by chunking large inputs at whitespace boundaries before encoding + ## [0.4.0] - 2026-01-06 - Add Ruby 4.0 compatibility diff --git a/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb b/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb index f43789c..4bfb8ce 100644 --- a/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb +++ b/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb @@ -4,17 +4,22 @@ module DiscourseAi module Tokenizer # Wrapper for OpenAI tokenizer library for compatibility with Discourse AI API class OpenAiTokenizer < BasicTokenizer + # tiktoken-rs uses fancy-regex which can stack overflow on large inputs + # due to catastrophic backtracking (github.com/openai/tiktoken/issues/245). + # Chunking at whitespace boundaries prevents this while preserving accuracy. + SAFE_CHUNK_SIZE = 50_000 + class << self def tokenizer @tokenizer ||= Tiktoken.get_encoding("o200k_base") end def tokenize(text) - tokenizer.encode(text) + safe_encode(text) end def encode(text) - tokenizer.encode(text) + safe_encode(text) end def decode(token_ids) @@ -72,7 +77,34 @@ def below_limit?(text, limit, strict: false) # than can take more than 1 token per char return true if !strict && text.size < limit / 2 - tokenizer.encode(text).length < limit + safe_encode(text).length < limit + end + + private + + def safe_encode(text) + if !text.is_a?(String) || text.size <= SAFE_CHUNK_SIZE + return tokenizer.encode(text) + end + + tokens = [] + offset = 0 + while offset < text.size + chunk_end = offset + SAFE_CHUNK_SIZE + + if chunk_end < text.size + # Split at a whitespace boundary to preserve tokenization accuracy + break_point = text.rindex(/\s/, chunk_end) + chunk_end = break_point if break_point && break_point > offset + else + chunk_end = text.size + end + + tokens.concat(tokenizer.encode(text[offset...chunk_end])) + offset = chunk_end + end + + tokens end end end diff --git a/lib/discourse_ai/tokenizers/version.rb b/lib/discourse_ai/tokenizers/version.rb index 872d415..9b6d014 100644 --- a/lib/discourse_ai/tokenizers/version.rb +++ b/lib/discourse_ai/tokenizers/version.rb @@ -2,6 +2,6 @@ module DiscourseAi module Tokenizers - VERSION = "0.4" + VERSION = "0.4.1" end end diff --git a/spec/discourse_ai/tokenizers/safe_encode_spec.rb b/spec/discourse_ai/tokenizers/safe_encode_spec.rb new file mode 100644 index 0000000..3b48d24 --- /dev/null +++ b/spec/discourse_ai/tokenizers/safe_encode_spec.rb @@ -0,0 +1,113 @@ +# frozen_string_literal: true + +require "spec_helper" + +RSpec.describe DiscourseAi::Tokenizer::OpenAiTokenizer do + let(:chunk_size) { described_class::SAFE_CHUNK_SIZE } + + shared_examples "safe encoding" do |tokenizer_class| + describe "text under chunk size" do + it "encodes normally without chunking" do + text = "Hello world, this is a normal sentence." + raw_tokenizer = tokenizer_class.tokenizer + + expect(tokenizer_class.encode(text)).to eq(raw_tokenizer.encode(text)) + end + end + + describe "large text with whitespace" do + it "produces identical tokens to direct encoding" do + text = "The quick brown fox jumps over the lazy dog. " * 2000 + expect(text.size).to be > chunk_size + + direct_tokens = tokenizer_class.tokenizer.encode(text) + chunked_tokens = tokenizer_class.encode(text) + + expect(chunked_tokens).to eq(direct_tokens) + end + + it "works with tokenize" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + tokens = tokenizer_class.tokenize(text) + expect(tokens).to be_an(Array) + expect(tokens.length).to be > 0 + end + + it "works with size" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + direct_size = tokenizer_class.tokenizer.encode(text).length + expect(tokenizer_class.size(text)).to eq(direct_size) + end + end + + describe "large text without whitespace" do + it "encodes repeated characters without crashing" do + text = "M" * (chunk_size + 1000) + + tokens = tokenizer_class.encode(text) + expect(tokens).to be_an(Array) + expect(tokens.length).to be > 0 + end + + it "encodes long non-whitespace runs without crashing" do + text = ("a".."z").to_a.join * 3000 + expect(text.size).to be > chunk_size + + tokens = tokenizer_class.encode(text) + expect(tokens).to be_an(Array) + expect(tokens.length).to be > 0 + end + end + + describe "chunking splits at whitespace boundaries" do + it "splits at the last whitespace before chunk boundary" do + before_boundary = "a" * (chunk_size - 10) + after_boundary = "b" * 100 + text = "#{before_boundary} #{after_boundary}" + + direct_tokens = + tokenizer_class.tokenizer.encode(before_boundary) + + tokenizer_class.tokenizer.encode(" #{after_boundary}") + + expect(tokenizer_class.encode(text)).to eq(direct_tokens) + end + end + + describe "below_limit? with large text" do + it "returns correct result for text over chunk size" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + actual_token_count = tokenizer_class.tokenizer.encode(text).length + expect( + tokenizer_class.below_limit?(text, actual_token_count + 1) + ).to be true + expect(tokenizer_class.below_limit?(text, 1)).to be false + end + end + + describe "truncate with large text" do + it "truncates large text correctly" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + result = tokenizer_class.truncate(text, 100, strict: true) + expect(tokenizer_class.size(result)).to be <= 100 + expect(result.length).to be < text.length + end + end + end + + describe DiscourseAi::Tokenizer::OpenAiTokenizer do + include_examples "safe encoding", DiscourseAi::Tokenizer::OpenAiTokenizer + end + + describe DiscourseAi::Tokenizer::OpenAiCl100kTokenizer do + include_examples "safe encoding", + DiscourseAi::Tokenizer::OpenAiCl100kTokenizer + end +end