WIP: automation triage using ai tool

Author: SamSaffron

config/locales/client.en.yml

@@ -81,6 +81,15 @@ en:
               label: "Top P"
               description: "Top P to use for the LLM, increase to increase randomness (0 to use model default)"
 
+        llm_tool_triage:
+          fields:
+            model:
+              label: "Model"
+              description: "The default language model used for triage"
+            tool:
+              label: "Tool"
+              description: "Tool to use for triage (tool must have no parameters defined)"
+
         llm_triage:
           fields:
             system_prompt:
@@ -122,7 +131,6 @@ en:
             model:
               label: "Model"
               description: "Language model used for triage"
-
     discourse_ai:
       title: "AI"
 

config/locales/server.en.yml

@@ -6,6 +6,9 @@ en:
         spam: "Flag as spam and hide post"
         spam_silence: "Flag as spam, hide post and silence user"
     scriptables:
+      llm_tool_triage:
+        title: Triage posts using AI Tool
+        description: "Triage posts using custom logic in an AI tool"
       llm_triage:
         title: Triage posts using AI
         description: "Triage posts using a large language model"

discourse_automation/llm_tool_triage.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+if defined?(DiscourseAutomation)
+  DiscourseAutomation::Scriptable.add("llm_tool_triage") do
+    version 1
+    run_in_background
+
+    triggerables %i[post_created_edited]
+
+    field :model,
+          component: :choices,
+          required: true,
+          extra: {
+            content: DiscourseAi::Automation.available_models,
+          }
+
+    field :tool,
+          component: :choices,
+          required: true,
+          extra: {
+            content: DiscourseAi::Automation.available_custom_tools,
+          }
+
+    script do |context, fields|
+      model = fields["model"]["value"]
+      tool_id = fields["tool"]["value"]
+
+      category_id = fields.dig("category", "value")
+      tags = fields.dig("tags", "value")
+
+      if post.topic.private_message?
+        include_personal_messages = fields.dig("include_personal_messages", "value")
+        next if !include_personal_messages
+      end
+
+      begin
+        RateLimiter.new(
+          Discourse.system_user,
+          "llm_tool_triage_#{post.id}",
+          SiteSetting.ai_automation_max_triage_per_post_per_minute,
+          1.minute,
+        ).performed!
+
+        RateLimiter.new(
+          Discourse.system_user,
+          "llm_tool_triage",
+          SiteSetting.ai_automation_max_triage_per_minute,
+          1.minute,
+        ).performed!
+
+        DiscourseAi::Automation::LlmToolTriage.handle(
+          post: post,
+          model: model,
+          tool_id: tool_id,
+          category_id: category_id,
+          tags: tags,
+          automation: self.automation,
+        )
+      rescue => e
+        Discourse.warn_exception(e, message: "llm_tool_triage: skipped triage on post #{post.id}")
+      end
+    end
+  end
+end

lib/ai_bot/tool_runner.rb

@@ -56,6 +56,7 @@ def framework_script
 
         const llm = {
           truncate: _llm_truncate,
+          generate: _llm_generate,
         };
 
         const index = {
@@ -175,20 +176,60 @@ def attach_truncate(mini_racer_context)
           "_llm_truncate",
           ->(text, length) { @llm.tokenizer.truncate(text, length) },
         )
+
+        mini_racer_context.attach(
+          "_llm_generate",
+          ->(prompt) do
+            in_attached_function do
+              @llm.generate(
+                convert_js_prompt_to_ruby(prompt),
+                user: llm_user,
+                feature_name: "custom_tool_#{tool.name}",
+              )
+            end
+          end,
+        )
+      end
+
+      def convert_js_prompt_to_ruby(prompt)
+        if prompt.is_a?(String)
+          prompt
+        elsif prompt.is_a?(Hash)
+          messages = prompt["messages"]
+          if messages.blank? || !messages.is_a?(Array)
+            raise Discourse::InvalidParameters.new("Prompt must have messages")
+          end
+          messages.each(&:symbolize_keys!)
+          messages.each { |message| message[:type] = message[:type].to_sym }
+          DiscourseAi::Completions::Prompt.new(messages: prompt["messages"])
+        else
+          raise Discourse::InvalidParameters.new("Prompt must be a string or a hash")
+        end
+      end
+
+      def llm_user
+        @llm_user ||=
+          begin
+            @context[:llm_user] || post&.user || @bot_user
+          end
+      end
+
+      def post
+        return @post if defined?(@post)
+        post_id = @context[:post_id]
+        @post = post_id && Post.find_by(id: post_id)
       end
 
       def attach_index(mini_racer_context)
         mini_racer_context.attach(
           "_index_search",
           ->(*params) do
-            begin
+            in_attached_function do
               query, options = params
               self.running_attached_function = true
               options ||= {}
               options = options.symbolize_keys
               self.rag_search(query, **options)
-            ensure
-              self.running_attached_function = false
             end
           end,
         )
@@ -203,26 +244,25 @@ def attach_upload(mini_racer_context)
           "_upload_create",
           ->(filename, base_64_content) do
             begin
-              self.running_attached_function = true
-              # protect against misuse
-              filename = File.basename(filename)
-
-              Tempfile.create(filename) do |file|
-                file.binmode
-                file.write(Base64.decode64(base_64_content))
-                file.rewind
-
-                upload =
-                  UploadCreator.new(
-                    file,
-                    filename,
-                    for_private_message: @context[:private_message],
-                  ).create_for(@bot_user.id)
-
-                { id: upload.id, short_url: upload.short_url, url: upload.url }
+              in_attached_function do
+                # protect against misuse
+                filename = File.basename(filename)
+
+                Tempfile.create(filename) do |file|
+                  file.binmode
+                  file.write(Base64.decode64(base_64_content))
+                  file.rewind
+
+                  upload =
+                    UploadCreator.new(
+                      file,
+                      filename,
+                      for_private_message: @context[:private_message],
+                    ).create_for(@bot_user.id)
+
+                  { id: upload.id, short_url: upload.short_url, url: upload.url }
+                end
               end
-            ensure
-              self.running_attached_function = false
             end
           end,
         )
@@ -238,18 +278,20 @@ def attach_http(mini_racer_context)
                 raise TooManyRequestsError.new("Tool made too many HTTP requests")
               end
 
-              self.running_attached_function = true
-              headers = (options && options["headers"]) || {}
+              in_attached_function do
+                headers = (options && options["headers"]) || {}
 
-              result = {}
-              DiscourseAi::AiBot::Tools::Tool.send_http_request(url, headers: headers) do |response|
-                result[:body] = response.body
-                result[:status] = response.code.to_i
-              end
+                result = {}
+                DiscourseAi::AiBot::Tools::Tool.send_http_request(
+                  url,
+                  headers: headers,
+                ) do |response|
+                  result[:body] = response.body
+                  result[:status] = response.code.to_i
+                end
 
-              result
-            ensure
-              self.running_attached_function = false
+                result
+              end
             end
           end,
         )
@@ -264,35 +306,43 @@ def attach_http(mini_racer_context)
                   raise TooManyRequestsError.new("Tool made too many HTTP requests")
                 end
 
-                self.running_attached_function = true
-                headers = (options && options["headers"]) || {}
-                body = options && options["body"]
-
-                result = {}
-                DiscourseAi::AiBot::Tools::Tool.send_http_request(
-                  url,
-                  method: method,
-                  headers: headers,
-                  body: body,
-                ) do |response|
-                  result[:body] = response.body
-                  result[:status] = response.code.to_i
+                in_attached_function do
+                  headers = (options && options["headers"]) || {}
+                  body = options && options["body"]
+
+                  result = {}
+                  DiscourseAi::AiBot::Tools::Tool.send_http_request(
+                    url,
+                    method: method,
+                    headers: headers,
+                    body: body,
+                  ) do |response|
+                    result[:body] = response.body
+                    result[:status] = response.code.to_i
+                  end
+
+                  result
+                rescue => e
+                  if Rails.env.development?
+                    p url
+                    p options
+                    p e
+                    puts e.backtrace
+                  end
+                  raise e
                 end
-
-                result
-              rescue => e
-                p url
-                p options
-                p e
-                puts e.backtrace
-                raise e
-              ensure
-                self.running_attached_function = false
               end
             end,
           )
         end
       end
+
+      def in_attached_function
+        self.running_attached_function = true
+        yield
+      ensure
+        self.running_attached_function = false
+      end
     end
   end
 end

lib/automation.rb

@@ -12,6 +12,15 @@ def self.flag_types
         },
       ]
     end
+
+    def self.available_custom_tools
+      AiTool
+        .where(enabled: true)
+        .where("parameters = '[]'::jsonb")
+        .pluck(:id, :name, :description)
+        .map { |id, name, description| { id: id, translated_name: name, description: description } }
+    end
+
     def self.available_models
       values = DB.query_hash(<<~SQL)
         SELECT display_name AS translated_name, id AS id

lib/completions/prompt.rb

@@ -65,6 +65,21 @@ def encoded_uploads(message)
         UploadEncoder.encode(upload_ids: message[:upload_ids], max_pixels: max_pixels)
       end
 
+      def ==(other)
+        return false unless other.is_a?(Prompt)
+        messages == other.messages && tools == other.tools && topic_id == other.topic_id &&
+          post_id == other.post_id && max_pixels == other.max_pixels &&
+          tool_choice == other.tool_choice
+      end
+
+      def eql?(other)
+        self == other
+      end
+
+      def hash
+        [messages, tools, topic_id, post_id, max_pixels, tool_choice].hash
+      end
+
       private
 
       def validate_message(message)

plugin.rb

@@ -65,6 +65,7 @@ def self.public_asset_path(name)
   # do not autoload this cause we may have no namespace
   require_relative "discourse_automation/llm_triage"
   require_relative "discourse_automation/llm_report"
+  require_relative "discourse_automation/llm_tool_triage"
 
   add_admin_route("discourse_ai.title", "discourse-ai", { use_new_show_route: true })