From 541f24ab0bde4ee3026d2bb523535c36f2df726f Mon Sep 17 00:00:00 2001
From: Jorge Manrubia <jorge@37signals.com>
Date: Fri, 16 May 2025 16:39:18 +0200
Subject: [PATCH] Limit the tokens we send for embeds

See insight by Donal here https://github.com/basecamp/fizzy/pull/483#discussion_r2084766023
---
 Gemfile                           |  2 ++
 Gemfile.lock                      |  7 +++++++
 app/models/ai/tokenizer.rb        | 31 +++++++++++++++++++++++++++++++
 app/models/concerns/searchable.rb |  4 ++--
 4 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 app/models/ai/tokenizer.rb

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