diff --git a/Gemfile b/Gemfile index 86ffeb0e0..64a32d01a 100644 --- a/Gemfile +++ b/Gemfile @@ -33,8 +33,10 @@ gem "sentry-ruby" gem "sentry-rails" gem "rails_structured_logging", bc: "rails-structured-logging" +# AI gem "ruby_llm", git: "https://github.com/crmne/ruby_llm.git" gem "sqlite-vec", "0.1.7.alpha.2" +gem "tiktoken_ruby" group :development, :test do gem "debug" diff --git a/Gemfile.lock b/Gemfile.lock index 95140273e..b8dd7eebb 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -304,6 +304,7 @@ GEM nokogiri (>= 1.15.7, != 1.16.7, != 1.16.6, != 1.16.5, != 1.16.4, != 1.16.3, != 1.16.2, != 1.16.1, != 1.16.0.rc1, != 1.16.0) rainbow (3.1.1) rake (13.2.1) + rb_sys (0.9.106) rdoc (6.13.1) psych (>= 4.0.0) redcarpet (3.6.1) @@ -400,6 +401,11 @@ GEM thruster (0.1.13-arm64-darwin) thruster (0.1.13-x86_64-darwin) thruster (0.1.13-x86_64-linux) + tiktoken_ruby (0.0.11.1) + rb_sys (= 0.9.106) + tiktoken_ruby (0.0.11.1-arm64-darwin) + tiktoken_ruby (0.0.11.1-x86_64-darwin) + tiktoken_ruby (0.0.11.1-x86_64-linux) timeout (0.4.3) turbo-rails (2.0.13) actionpack (>= 7.1.0) @@ -464,6 +470,7 @@ DEPENDENCIES sqlite3 (>= 2.0) stimulus-rails thruster + tiktoken_ruby turbo-rails vcr webmock diff --git a/app/models/ai/tokenizer.rb b/app/models/ai/tokenizer.rb new file mode 100644 index 000000000..dfc8db657 --- /dev/null +++ b/app/models/ai/tokenizer.rb @@ -0,0 +1,31 @@ +class Ai::Tokenizer + attr_reader :text, :max_input_tokens + + class << self + def truncate(text, max_input_tokens: 8196) + new(text, max_input_tokens:).truncated + end + end + + def initialize(text, max_input_tokens: 8196) + @text = text + @max_input_tokens = max_input_tokens + end + + def truncated + # Truncating the tokens might split a unicode character so if we get an error + # we'll try removing an extra token + # The encode/decode round trip seems to add a token, so we start with max_input_tokens - 1 + (1..4).each do |i| + tokens = tokenizer.encode(text)[0..(max_input_tokens - 20 - i)] + return tokenizer.decode(tokens) + rescue Tiktoken::UnicodeError + raise if i == 4 + end + end + + private + def tokenizer + @tokenizer ||= Tiktoken.encoding_for_model("text-embedding-3-small") + end +end diff --git a/app/models/concerns/searchable.rb b/app/models/concerns/searchable.rb index ef9f32306..5946212dd 100644 --- a/app/models/concerns/searchable.rb +++ b/app/models/concerns/searchable.rb @@ -21,7 +21,7 @@ module Searchable scope :search, ->(query) { joins("join #{using} idx on #{table_name}.id = idx.rowid").where("idx.#{as} match ?", query) } scope :search_similar, ->(query) do - query_embedding = Rails.cache.fetch("embed-search:#{query}") { RubyLLM.embed(query) } + query_embedding = Rails.cache.fetch("embed-search:#{query}") { RubyLLM.embed(Ai::Tokenizer.truncate(query)) } joins(:search_embedding) .where("embedding MATCH ? AND k = ?", query_embedding.vectors.to_json, 20) .order(:distance) @@ -34,7 +34,7 @@ module Searchable end def refresh_search_embedding - embedding = RubyLLM.embed(search_embedding_content) + embedding = RubyLLM.embed(Ai::Tokenizer.truncate(search_embedding_content)) search_embedding = self.search_embedding || build_search_embedding search_embedding.update! embedding: embedding.vectors.to_json end