DEV: Move tokenizers to a gem (#1481)

Also renames the Mixtral tokenizer to Mistral.

See gem at github.com/discourse/discourse_ai-tokenizers


Co-authored-by: Roman Rizzi <roman@discourse.org>

Author: xfalcox

.gitignore

@@ -6,5 +6,4 @@ evals/log
 evals/cases
 config/eval-llms.local.yml
 # this gets rid of search results from ag, ripgrep, etc
-tokenizers/
 public/ai-share/highlight.min.js

app/models/embedding_definition.rb

@@ -23,7 +23,7 @@ def tokenizer_names
         DiscourseAi::Tokenizer::GeminiTokenizer,
         DiscourseAi::Tokenizer::MultilingualE5LargeTokenizer,
         DiscourseAi::Tokenizer::OpenAiTokenizer,
-        DiscourseAi::Tokenizer::MixtralTokenizer,
+        DiscourseAi::Tokenizer::MistralTokenizer,
         DiscourseAi::Tokenizer::QwenTokenizer,
       ].map(&:name)
     end

db/migrate/20250702073222_rename_mixtral_tokenizer_to_mistral_tokenizer.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+class RenameMixtralTokenizerToMistralTokenizer < ActiveRecord::Migration[7.2]
+  def up
+    execute <<~SQL
+      UPDATE
+        llm_models
+      SET
+        tokenizer = 'DiscourseAi::Tokenizer::Mistral'
+      WHERE
+        tokenizer = 'DiscourseAi::Tokenizer::Mixtral'
+    SQL
+
+    execute <<~SQL
+      UPDATE
+        embedding_definitions
+      SET
+        tokenizer_class = 'DiscourseAi::Tokenizer::Mistral'
+      WHERE
+        tokenizer_class = 'DiscourseAi::Tokenizer::Mixtral'
+    SQL
+  end
+
+  def down
+    execute <<~SQL
+      UPDATE
+        llm_models
+      SET
+        tokenizer = 'DiscourseAi::Tokenizer::Mixtral'
+      WHERE
+        tokenizer = 'DiscourseAi::Tokenizer::Mistral'
+    SQL
+
+    execute <<~SQL
+      UPDATE
+        embedding_definitions
+      SET
+        tokenizer_class = 'DiscourseAi::Tokenizer::Mixtral'
+      WHERE
+        tokenizer_class = 'DiscourseAi::Tokenizer::Mistral'
+    SQL
+  end
+end

lib/automation/llm_triage.rb

@@ -40,7 +40,12 @@ def self.handle(
 
         content = "title: #{post.topic.title}\n#{post.raw}"
 
-        content = llm.tokenizer.truncate(content, max_post_tokens) if max_post_tokens.present?
+        content =
+          llm.tokenizer.truncate(
+            content,
+            max_post_tokens,
+            strict: SiteSetting.ai_strict_token_counting,
+          ) if max_post_tokens.present?
 
         if post.upload_ids.present?
           content = [content]

lib/automation/report_context_generator.rb

@@ -99,7 +99,12 @@ def format_post(post)
         buffer << post.created_at.strftime("%Y-%m-%d %H:%M")
         buffer << "user: #{post.user&.username}"
         buffer << "likes: #{post.like_count}"
-        excerpt = @tokenizer.truncate(post.raw, @tokens_per_post)
+        excerpt =
+          @tokenizer.truncate(
+            post.raw,
+            @tokens_per_post,
+            strict: SiteSetting.ai_strict_token_counting,
+          )
         excerpt = "excerpt: #{excerpt}..." if excerpt.length < post.raw.length
         buffer << "#{excerpt}"
         { likes: post.like_count, info: buffer.join("\n") }

lib/completions/dialects/dialect.rb

@@ -147,6 +147,7 @@ def trim_messages(messages)
               system_message[:content] = tokenizer.truncate(
                 system_message[:content],
                 max_system_tokens,
+                strict: SiteSetting.ai_strict_token_counting,
               )
             end
 

lib/completions/llm.rb

@@ -190,7 +190,7 @@ def presets
                       display_name: "Pixtral Large",
                     },
                   ],
-                  tokenizer: DiscourseAi::Tokenizer::MixtralTokenizer,
+                  tokenizer: DiscourseAi::Tokenizer::MistralTokenizer,
                   endpoint: "https://api.mistral.ai/v1/chat/completions",
                   provider: "mistral",
                 },

lib/embeddings/strategies/truncation.rb

@@ -22,7 +22,11 @@ def prepare_target_text(target, vdef)
             when Post
               post_truncation(target, vdef.tokenizer, max_length)
             when RagDocumentFragment
-              vdef.tokenizer.truncate(target.fragment, max_length)
+              vdef.tokenizer.truncate(
+                target.fragment,
+                max_length,
+                strict: SiteSetting.ai_strict_token_counting,
+              )
             else
               raise ArgumentError, "Invalid target type"
             end
@@ -36,7 +40,7 @@ def prepare_query_text(text, vdef, asymetric: false)
           qtext = asymetric ? "#{vdef.search_prompt} #{text}" : text
           max_length = vdef.max_sequence_length - 2
 
-          vdef.tokenizer.truncate(qtext, max_length)
+          vdef.tokenizer.truncate(qtext, max_length, strict: SiteSetting.ai_strict_token_counting)
         end
 
         private
@@ -74,7 +78,7 @@ def topic_truncation(topic, tokenizer, max_length)
             text << "\n\n"
           end
 
-          tokenizer.truncate(text, max_length)
+          tokenizer.truncate(text, max_length, strict: SiteSetting.ai_strict_token_counting)
         end
 
         def post_truncation(post, tokenizer, max_length)
@@ -86,7 +90,7 @@ def post_truncation(post, tokenizer, max_length)
             text << Nokogiri::HTML5.fragment(post.cooked).text
           end
 
-          tokenizer.truncate(text, max_length)
+          tokenizer.truncate(text, max_length, strict: SiteSetting.ai_strict_token_counting)
         end
       end
     end

lib/personas/question_consolidator.rb

@@ -42,7 +42,12 @@ def revised_prompt
           truncated_content = content
 
           if current_tokens > allowed_tokens
-            truncated_content = @llm.tokenizer.truncate(content, allowed_tokens)
+            truncated_content =
+              @llm.tokenizer.truncate(
+                content,
+                allowed_tokens,
+                strict: SiteSetting.ai_strict_token_counting,
+              )
             current_tokens = allowed_tokens
           end
 

lib/personas/tool_runner.rb

@@ -278,7 +278,9 @@ def rag_search(query, filenames: nil, limit: 10)
       def attach_truncate(mini_racer_context)
         mini_racer_context.attach(
           "_llm_truncate",
-          ->(text, length) { @llm.tokenizer.truncate(text, length) },
+          ->(text, length) do
+            @llm.tokenizer.truncate(text, length, strict: SiteSetting.ai_strict_token_counting)
+          end,
         )
 
         mini_racer_context.attach(