FEATURE: move summary to use llm_model (#699)

This allows summary to use the new LLM models and migrates of API key based model selection Claude 3.5 etc... all work now. --------- Co-authored-by: Roman Rizzi <rizziromanalejandro@gmail.com>
discourse · Jul 4, 2024 · 1320eed · 1320eed
1 parent fc081d9
commit 1320eed
Show file tree

Hide file tree

Showing 33 changed files with 279 additions and 519 deletions.
diff --git a/Gemfile.lock b/Gemfile.lock
@@ -1,7 +1,7 @@
 GEM
   remote: https://rubygems.org/
   specs:
-    activesupport (7.1.3.3)
+    activesupport (7.1.3.4)
       base64
       bigdecimal
       concurrent-ruby (~> 1.0, >= 1.0.2)
@@ -14,27 +14,27 @@ GEM
     ast (2.4.2)
     base64 (0.2.0)
     bigdecimal (3.1.8)
-    concurrent-ruby (1.2.3)
+    concurrent-ruby (1.3.3)
     connection_pool (2.4.1)
     drb (2.2.1)
     i18n (1.14.5)
       concurrent-ruby (~> 1.0)
     json (2.7.2)
     language_server-protocol (3.17.0.3)
-    minitest (5.23.1)
+    minitest (5.24.1)
     mutex_m (0.2.0)
-    parallel (1.24.0)
-    parser (3.3.1.0)
+    parallel (1.25.1)
+    parser (3.3.3.0)
       ast (~> 2.4.1)
       racc
     prettier_print (1.2.1)
     racc (1.8.0)
-    rack (3.0.11)
+    rack (3.1.6)
     rainbow (3.1.1)
     regexp_parser (2.9.2)
-    rexml (3.2.8)
-      strscan (>= 3.0.9)
-    rubocop (1.64.0)
+    rexml (3.3.1)
+      strscan
+    rubocop (1.64.1)
       json (~> 2.3)
       language_server-protocol (>= 3.17.0)
       parallel (~> 1.10)
@@ -47,29 +47,28 @@ GEM
       unicode-display_width (>= 2.4.0, < 3.0)
     rubocop-ast (1.31.3)
       parser (>= 3.3.1.0)
-    rubocop-capybara (2.20.0)
+    rubocop-capybara (2.21.0)
       rubocop (~> 1.41)
-    rubocop-discourse (3.8.0)
+    rubocop-discourse (3.8.1)
       activesupport (>= 6.1)
       rubocop (>= 1.59.0)
       rubocop-capybara (>= 2.0.0)
       rubocop-factory_bot (>= 2.0.0)
       rubocop-rails (>= 2.25.0)
-      rubocop-rspec (>= 2.25.0)
-    rubocop-factory_bot (2.25.1)
-      rubocop (~> 1.41)
-    rubocop-rails (2.25.0)
+      rubocop-rspec (>= 3.0.1)
+      rubocop-rspec_rails (>= 2.30.0)
+    rubocop-factory_bot (2.26.1)
+      rubocop (~> 1.61)
+    rubocop-rails (2.25.1)
       activesupport (>= 4.2.0)
       rack (>= 1.1)
       rubocop (>= 1.33.0, < 2.0)
       rubocop-ast (>= 1.31.1, < 2.0)
-    rubocop-rspec (2.29.2)
-      rubocop (~> 1.40)
-      rubocop-capybara (~> 2.17)
-      rubocop-factory_bot (~> 2.22)
-      rubocop-rspec_rails (~> 2.28)
-    rubocop-rspec_rails (2.28.3)
-      rubocop (~> 1.40)
+    rubocop-rspec (3.0.2)
+      rubocop (~> 1.61)
+    rubocop-rspec_rails (2.30.0)
+      rubocop (~> 1.61)
+      rubocop-rspec (~> 3, >= 3.0.1)
     ruby-progressbar (1.13.0)
     strscan (3.1.0)
     syntax_tree (6.2.0)

diff --git a/app/controllers/discourse_ai/summarization/chat_summary_controller.rb b/app/controllers/discourse_ai/summarization/chat_summary_controller.rb
@@ -15,11 +15,10 @@ def show
         channel = ::Chat::Channel.find(params[:channel_id])
         guardian.ensure_can_join_chat_channel!(channel)
 
-        strategy = DiscourseAi::Summarization::Models::Base.selected_strategy
+        strategy = DiscourseAi::Summarization.default_strategy
         raise Discourse::NotFound.new unless strategy
-        unless DiscourseAi::Summarization::Models::Base.can_request_summary_for?(current_user)
-          raise Discourse::InvalidAccess
-        end
+
+        guardian.ensure_can_request_summary!
 
         RateLimiter.new(current_user, "channel_summary", 6, 5.minutes).performed!
 

diff --git a/app/controllers/discourse_ai/summarization/summary_controller.rb b/app/controllers/discourse_ai/summarization/summary_controller.rb
@@ -8,12 +8,8 @@ class SummaryController < ::ApplicationController
       def show
         topic = Topic.find(params[:topic_id])
         guardian.ensure_can_see!(topic)
-        strategy = DiscourseAi::Summarization::Models::Base.selected_strategy
 
-        if strategy.nil? ||
-             !DiscourseAi::Summarization::Models::Base.can_see_summary?(topic, current_user)
-          raise Discourse::NotFound
-        end
+        raise Discourse::NotFound if !guardian.can_see_summary?(topic)
 
         RateLimiter.new(current_user, "summary", 6, 5.minutes).performed! if current_user
 
@@ -30,9 +26,7 @@ def show
           render json: success_json
         else
           hijack do
-            summary =
-              DiscourseAi::TopicSummarization.new(strategy).summarize(topic, current_user, opts)
-
+            summary = DiscourseAi::TopicSummarization.summarize(topic, current_user, opts)
             render_serialized(summary, AiTopicSummarySerializer)
           end
         end

diff --git a/app/jobs/regular/stream_topic_ai_summary.rb b/app/jobs/regular/stream_topic_ai_summary.rb
@@ -8,10 +8,8 @@ def execute(args)
       return unless topic = Topic.find_by(id: args[:topic_id])
       return unless user = User.find_by(id: args[:user_id])
 
-      strategy = DiscourseAi::Summarization::Models::Base.selected_strategy
-      if strategy.nil? || !DiscourseAi::Summarization::Models::Base.can_see_summary?(topic, user)
-        return
-      end
+      strategy = DiscourseAi::Summarization.default_strategy
+      return if strategy.nil? || !Guardian.new(user).can_see_summary?(topic)
 
       guardian = Guardian.new(user)
       return unless guardian.can_see?(topic)

diff --git a/app/serializers/ai_topic_summary_serializer.rb b/app/serializers/ai_topic_summary_serializer.rb
@@ -4,7 +4,7 @@ class AiTopicSummarySerializer < ApplicationSerializer
   attributes :summarized_text, :algorithm, :outdated, :can_regenerate, :new_posts_since_summary
 
   def can_regenerate
-    DiscourseAi::Summarization::Models::Base.can_request_summary_for?(scope.current_user)
+    scope.can_request_summary?
   end
 
   def new_posts_since_summary

diff --git a/app/services/discourse_ai/topic_summarization.rb b/app/services/discourse_ai/topic_summarization.rb
@@ -2,6 +2,10 @@
 
 module DiscourseAi
   class TopicSummarization
+    def self.summarize(topic, user, opts = {}, &on_partial_blk)
+      new(DiscourseAi::Summarization.default_strategy).summarize(topic, user, opts, &on_partial_blk)
+    end
+
     def initialize(strategy)
       @strategy = strategy
     end
@@ -15,7 +19,7 @@ def summarize(topic, user, opts = {}, &on_partial_blk)
       targets_data = summary_targets(topic).pluck(:post_number, :raw, :username)
 
       current_topic_sha = build_sha(targets_data.map(&:first))
-      can_summarize = DiscourseAi::Summarization::Models::Base.can_request_summary_for?(user)
+      can_summarize = Guardian.new(user).can_request_summary?
 
       if use_cached?(existing_summary, can_summarize, current_topic_sha, !!opts[:skip_age_check])
         # It's important that we signal a cached summary is outdated

diff --git a/assets/javascripts/initializers/ai-chat-summarization.js b/assets/javascripts/initializers/ai-chat-summarization.js
@@ -6,10 +6,7 @@ export default apiInitializer("1.34.0", (api) => {
   const currentUser = api.getCurrentUser();
   const chatService = api.container.lookup("service:chat");
   const modal = api.container.lookup("service:modal");
-  const canSummarize =
-    siteSettings.ai_summarization_strategy &&
-    currentUser &&
-    currentUser.can_summarize;
+  const canSummarize = currentUser && currentUser.can_summarize;
 
   if (
     !siteSettings.chat_enabled ||

diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml
@@ -92,11 +92,9 @@ en:
     ai_embeddings_semantic_related_include_closed_topics: "Include closed topics in semantic search results"
     ai_embeddings_semantic_search_hyde_model: "Model used to expand keywords to get better results during a semantic search"
     ai_embeddings_per_post_enabled: Generate embeddings for each post
-
-    ai_summarization_discourse_service_api_endpoint: "URL where the Discourse summarization API is running."
-    ai_summarization_discourse_service_api_key: "API key for the Discourse summarization API."
-    ai_summarization_strategy: "Additional ways to summarize content registered by plugins"
-    ai_custom_summarization_allowed_groups: "Groups allowed to summarize contents using the `summarization_strategy`."
+    ai_summarization_enabled: "Enable the topic summarization module."
+    ai_summarization_model: "Model to use for summarization."
+    ai_custom_summarization_allowed_groups: "Groups allowed to use create new summaries."
 
     ai_bot_enabled: "Enable the AI Bot module."
     ai_bot_enable_chat_warning: "Display a warning when PM chat is initiated. Can be overriden by editing the translation string: discourse_ai.ai_bot.pm_warning"

diff --git a/config/settings.yml b/config/settings.yml
@@ -319,23 +319,24 @@ discourse_ai:
     default: false
     client: true
     hidden: true
-
-  ai_summarization_discourse_service_api_endpoint: ""
-  ai_summarization_discourse_service_api_endpoint_srv:
-    default: ""
-    hidden: true
-  ai_summarization_discourse_service_api_key:
-    default: ""
-    secret: true
-  ai_summarization_strategy:
-    client: true
+  ai_summarization_enabled:
+    default: false
+    validator: "DiscourseAi::Configuration::LlmDependencyValidator"
+  ai_summarization_model:
     default: ""
-    enum: "DiscourseAi::Configuration::SummarizationEnumerator"
-    validator: "DiscourseAi::Configuration::SummarizationValidator"
+    allow_any: false
+    type: enum
+    enum: "DiscourseAi::Configuration::LlmEnumerator"
+    validator: "DiscourseAi::Configuration::LlmValidator"
   ai_custom_summarization_allowed_groups:
     type: group_list
     list_type: compact
     default: "3|13" # 3: @staff, 13: @trust_level_3
+  ai_summarization_strategy: # TODO(roman): Deprecated. Remove by Sept 2024
+    type: enum
+    default: ""
+    hidden: true
+    choices: "DiscourseAi::Configuration::LlmEnumerator.old_summarization_options + ['']"
 
   ai_bot_enabled:
     default: false
@@ -359,7 +360,7 @@ discourse_ai:
     default: "1|2" # 1: admins, 2: moderators
     allow_any: false
     refresh: true
-  ai_bot_enabled_chat_bots: # TODO(roman): Remove setting. Deprecated
+  ai_bot_enabled_chat_bots: # TODO(roman): Deprecated. Remove by Sept 2024
     type: list
     default: "gpt-3.5-turbo"
     hidden: true

diff --git a/db/post_migrate/20240703135444_llm_models_for_summarization.rb b/db/post_migrate/20240703135444_llm_models_for_summarization.rb
@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+class LlmModelsForSummarization < ActiveRecord::Migration[7.0]
+  def up
+    setting_value =
+      DB
+        .query_single(
+          "SELECT value FROM site_settings WHERE name = :llm_setting",
+          llm_setting: "ai_summarization_strategy",
+        )
+        .first
+        .to_s
+
+    return if setting_value.empty?
+
+    gpt_models = %w[gpt-4 gpt-4-32k gpt-4-turbo gpt-4o gpt-3.5-turbo gpt-3.5-turbo-16k]
+    gemini_models = %w[gemini-pro gemini-1.5-pro gemini-1.5-flash]
+    claude_models = %w[claude-2 claude-instant-1 claude-3-haiku claude-3-sonnet claude-3-opus]
+    oss_models = %w[mistralai/Mixtral-8x7B-Instruct-v0.1 mistralai/Mixtral-8x7B-Instruct-v0.1]
+
+    providers = []
+    prov_priority = ""
+
+    if gpt_models.include?(setting_value)
+      providers = %w[azure open_ai]
+      prov_priority = "azure"
+    elsif gemini_models.include?(setting_value)
+      providers = %w[google]
+      prov_priority = "google"
+    elsif claude_models.include?(setting_value)
+      providers = %w[aws_bedrock anthropic]
+      prov_priority = "aws_bedrock"
+    elsif oss_models.include?(setting_value)
+      providers = %w[hugging_face vllm]
+      prov_priority = "vllm"
+    end
+
+    insert_llm_model(setting_value, providers, prov_priority) if providers.present?
+  end
+
+  def insert_llm_model(old_value, providers, priority)
+    matching_models = DB.query(<<~SQL, model_name: old_value, providers: providers)
+      SELECT * FROM llm_models WHERE name = :model_name AND provider IN (:providers)
+    SQL
+
+    return if matching_models.empty?
+
+    priority_model = matching_models.find { |m| m.provider == priority } || matching_models.first
+    new_value = "custom:#{priority_model.id}"
+
+    DB.exec(<<~SQL, new_value: new_value)
+      INSERT INTO site_settings(name, data_type, value, created_at, updated_at)
+      VALUES ('ai_summarization_model', 1, :new_value, NOW(), NOW())
+    SQL
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration
+  end
+end
diff --git a/lib/completions/llm.rb b/lib/completions/llm.rb
@@ -200,7 +200,9 @@ def proxy_from_obj(llm_model)
               raise "Invalid call LLM call, expected #{@canned_llm} but got #{model_name}"
             end
 
-            return new(dialect_klass, nil, model_name, gateway: @canned_response)
+            return(
+              new(dialect_klass, nil, model_name, gateway: @canned_response, llm_model: llm_model)
+            )
           end
 
           gateway_klass = DiscourseAi::Completions::Endpoints::Base.endpoint_for(provider_name)
@@ -293,11 +295,11 @@ def tokenizer
           dialect_klass.new(DiscourseAi::Completions::Prompt.new(""), model_name).tokenizer
       end
 
-      attr_reader :model_name
+      attr_reader :model_name, :llm_model
 
       private
 
-      attr_reader :dialect_klass, :gateway_klass, :llm_model
+      attr_reader :dialect_klass, :gateway_klass
     end
   end
 end
diff --git a/lib/configuration/llm_enumerator.rb b/lib/configuration/llm_enumerator.rb
@@ -20,6 +20,29 @@ def self.values
         values
       end
 
+      # TODO(roman): Deprecated. Remove by Sept 2024
+      def self.old_summarization_options
+        %w[
+          gpt-4
+          gpt-4-32k
+          gpt-4-turbo
+          gpt-4o
+          gpt-3.5-turbo
+          gpt-3.5-turbo-16k
+          gemini-pro
+          gemini-1.5-pro
+          gemini-1.5-flash
+          claude-2
+          claude-instant-1
+          claude-3-haiku
+          claude-3-sonnet
+          claude-3-opus
+          mistralai/Mixtral-8x7B-Instruct-v0.1
+          mistralai/Mixtral-8x7B-Instruct-v0.1
+        ]
+      end
+
+      # TODO(roman): Deprecated. Remove by Sept 2024
       def self.available_ai_bots
         %w[
           gpt-3.5-turbo

diff --git a/lib/configuration/llm_validator.rb b/lib/configuration/llm_validator.rb
@@ -61,6 +61,7 @@ def modules_and_choose_llm_settings
         {
           ai_embeddings_semantic_search_enabled: :ai_embeddings_semantic_search_hyde_model,
           composer_ai_helper_enabled: :ai_helper_model,
+          ai_summarization_enabled: :ai_summarization_model,
         }
       end
     end

diff --git a/lib/configuration/summarization_enumerator.rb b/lib/configuration/summarization_enumerator.rb