Limit the tokens we send for embeds

See insight by Donal here https://github.com/basecamp/fizzy/pull/483#discussion_r2084766023

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"

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

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

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