FEATURE: support chain halting and upload creation support (#821)

This adds chain halting (ability to terminate llm chain in a tool)
and the ability to create uploads in a tool

Together this lets us integrate custom image generators into a
custom tool.

Author: SamSaffron

app/controllers/discourse_ai/admin/ai_tools_controller.rb

@@ -49,7 +49,7 @@ def test
         if params[:id].present?
           ai_tool = AiTool.find(params[:id])
         else
-          ai_tool = AiTool.new(ai_tool_params)
+          ai_tool = AiTool.new(ai_tool_params.except(:rag_uploads))
         end
 
         parameters = params[:parameters].to_unsafe_h

app/models/ai_tool.rb

@@ -93,6 +93,19 @@ def self.preamble
       *    Returns:
       *      Array of { fragment: string, metadata: string }
       *
+      * 4. upload
+      *    upload.create(filename, base_64_content): Uploads a file.
+      *    Parameters:
+      *      filename (string): Name of the file.
+      *      base_64_content (string): Base64 encoded file content.
+      *    Returns:
+      *      { id: number, short_url: string }
+      *
+      * 5. chain
+      *    chain.setCustomRaw(raw): Sets the body of the post and exist chain.
+      *    Parameters:
+      *      raw (string): raw content to add to post.
+      *
       * Constraints
       *
       * Execution Time: ≤ 2000ms
@@ -236,6 +249,70 @@ def self.presets
       SCRIPT
         summary: "Get real-time stock quotes using AlphaVantage API",
       },
+      {
+        preset_id: "image_generation",
+        name: "image_generation",
+        description:
+          "Generate images using the FLUX model from Black Forest Labs using together.ai",
+        parameters: [
+          {
+            name: "prompt",
+            type: "string",
+            required: true,
+            description: "The text prompt for image generation",
+          },
+          {
+            name: "seed",
+            type: "number",
+            required: false,
+            description: "Optional seed for random number generation",
+          },
+        ],
+        script: <<~SCRIPT,
+          #{preamble}
+          const apiKey = "YOUR_KEY";
+          const model = "black-forest-labs/FLUX.1.1-pro";
+
+          function invoke(params) {
+            let seed = parseInt(params.seed);
+            if (!(seed > 0)) {
+              seed = Math.floor(Math.random() * 1000000) + 1;
+            }
+
+            const prompt = params.prompt;
+            const body = {
+              model: model,
+              prompt: prompt,
+              width: 1024,
+              height: 768,
+              steps: 10,
+              n: 1,
+              seed: seed,
+              response_format: "b64_json",
+            };
+
+            const result = http.post("https://api.together.xyz/v1/images/generations", {
+              headers: {
+                "Authorization": `Bearer ${apiKey}`,
+                "Content-Type": "application/json",
+              },
+              body: JSON.stringify(body),
+            });
+
+            const base64Image = JSON.parse(result.body).data[0].b64_json;
+            const image = upload.create("generated_image.png", base64Image);
+            const raw = `\n![${prompt}](${image.short_url})\n`;
+            chain.setCustomRaw(raw);
+
+            return { result: "Image generated successfully", seed: seed };
+          }
+
+          function details() {
+            return "Generates images based on a text prompt using the FLUX model.";
+          }
+  SCRIPT
+        summary: "Generate image",
+      },
       { preset_id: "empty_tool", script: <<~SCRIPT },
           #{preamble}
           function invoke(params) {

config/locales/server.en.yml

@@ -128,11 +128,13 @@ en:
       custom_name: "%{name} (custom)"
       presets:
         browse_web_jina:
-          name: "Browse web using jina.ai"
+          name: "Browse web (jina.ai)"
         exchange_rate:
           name: "Exchange rate"
         stock_quote:
           name: "Stock quote (AlphaVantage)"
+        image_generation:
+          name: "Flux image generator (Together.ai)"
         empty_tool:
           name: "Start from blank..."
 

lib/ai_bot/tool_runner.rb

@@ -4,7 +4,7 @@ module DiscourseAi
   module AiBot
     class ToolRunner
       attr_reader :tool, :parameters, :llm
-      attr_accessor :running_attached_function, :timeout
+      attr_accessor :running_attached_function, :timeout, :custom_raw
 
       TooManyRequestsError = Class.new(StandardError)
 
@@ -36,6 +36,8 @@ def mini_racer_context
             attach_truncate(ctx)
             attach_http(ctx)
             attach_index(ctx)
+            attach_upload(ctx)
+            attach_chain(ctx)
             ctx.eval(framework_script)
             ctx
           end
@@ -55,6 +57,15 @@ def framework_script
         const index = {
           search: _index_search,
         }
+
+        const upload = {
+          create: _upload_create,
+        }
+
+        const chain = {
+          setCustomRaw: _chain_set_custom_raw,
+        };
+
         function details() { return ""; };
       JS
       end
@@ -176,6 +187,40 @@ def attach_index(mini_racer_context)
         )
       end
 
+      def attach_chain(mini_racer_context)
+        mini_racer_context.attach("_chain_set_custom_raw", ->(raw) { self.custom_raw = raw })
+      end
+
+      def attach_upload(mini_racer_context)
+        mini_racer_context.attach(
+          "_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 }
+              end
+            ensure
+              self.running_attached_function = false
+            end
+          end,
+        )
+      end
+
       def attach_http(mini_racer_context)
         mini_racer_context.attach(
           "_http_get",

lib/ai_bot/tools/custom.rb

@@ -30,8 +30,18 @@ def self.name
           AiTool.where(id: tool_id).pluck(:name).first
         end
 
+        def initialize(*args, **kwargs)
+          @chain_next_response = true
+          super(*args, **kwargs)
+        end
+
         def invoke
-          runner.invoke
+          result = runner.invoke
+          if runner.custom_raw
+            self.custom_raw = runner.custom_raw
+            @chain_next_response = false
+          end
+          result
         end
 
         def runner
@@ -50,6 +60,10 @@ def details
           runner.details
         end
 
+        def chain_next_response?
+          !!@chain_next_response
+        end
+
         def help
           # I do not think this is called, but lets make sure
           raise "Not implemented"

spec/lib/modules/ai_bot/playground_spec.rb

@@ -94,21 +94,52 @@
 
     let(:playground) { DiscourseAi::AiBot::Playground.new(bot) }
 
+    it "can create uploads from a tool" do
+      custom_tool.update!(script: <<~JS)
+        let imageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/wcAAgEB/awxUE0AAAAASUVORK5CYII="
+        function invoke(params) {
+          let image = upload.create("image.png", imageBase64);
+          chain.setCustomRaw(`![image](${image.short_url})`);
+          return image.id;
+        };
+      JS
+
+      tool_name = "custom-#{custom_tool.id}"
+      ai_persona.update!(tools: [[tool_name, nil, true]], tool_details: false)
+
+      reply_post = nil
+      prompts = nil
+
+      responses = [function_call]
+      DiscourseAi::Completions::Llm.with_prepared_responses(responses) do |_, _, _prompts|
+        new_post = Fabricate(:post, raw: "Can you use the custom tool?")
+        reply_post = playground.reply_to(new_post)
+        prompts = _prompts
+      end
+
+      expect(prompts.length).to eq(1)
+      upload_id = prompts[0].messages[3][:content].to_i
+
+      upload = Upload.find(upload_id)
+
+      expect(reply_post.raw).to eq("![image](#{upload.short_url})")
+    end
+
     it "can force usage of a tool" do
       tool_name = "custom-#{custom_tool.id}"
-      ai_persona.update!(tools: [[tool_name, nil, "force"]])
+      ai_persona.update!(tools: [[tool_name, nil, true]])
       responses = [function_call, "custom tool did stuff (maybe)"]
 
-      prompt = nil
-      DiscourseAi::Completions::Llm.with_prepared_responses(responses) do |_, _, _prompt|
+      prompts = nil
+      DiscourseAi::Completions::Llm.with_prepared_responses(responses) do |_, _, _prompts|
         new_post = Fabricate(:post, raw: "Can you use the custom tool?")
         _reply_post = playground.reply_to(new_post)
-        prompt = _prompt
+        prompts = _prompts
       end
 
-      expect(prompt.length).to eq(2)
-      expect(prompt[0].tool_choice).to eq("search")
-      expect(prompt[1].tool_choice).to eq(nil)
+      expect(prompts.length).to eq(2)
+      expect(prompts[0].tool_choice).to eq("search")
+      expect(prompts[1].tool_choice).to eq(nil)
     end
 
     it "uses custom tool in conversation" do

spec/system/ai_bot/tool_spec.rb

@@ -10,26 +10,9 @@
     sign_in(admin)
   end
 
-  it "allows admin to create a new AI tool from preset" do
-    visit "/admin/plugins/discourse-ai/ai-tools"
-
-    expect(page).to have_content("Tools")
-
-    find(".ai-tool-list-editor__new-button").click
-
-    select_kit = PageObjects::Components::SelectKit.new(".ai-tool-editor__presets")
-    select_kit.expand
-    select_kit.select_row_by_value("exchange_rate")
-
-    find(".ai-tool-editor__next").click
-
-    expect(page.first(".parameter-row__required-toggle").checked?).to eq(true)
-    expect(page.first(".parameter-row__enum-toggle").checked?).to eq(false)
-
+  def ensure_can_run_test
     find(".ai-tool-editor__test-button").click
 
-    expect(page).not_to have_button(".ai-tool-editor__delete")
-
     modal = PageObjects::Modals::AiToolTest.new
     modal.base_currency = "USD"
     modal.target_currency = "EUR"
@@ -48,14 +31,36 @@
     expect(modal).to have_content("0.85")
 
     modal.close
+  end
 
+  it "allows admin to create a new AI tool from preset" do
+    visit "/admin/plugins/discourse-ai/ai-tools"
+
+    expect(page).to have_content("Tools")
+
+    find(".ai-tool-list-editor__new-button").click
+
+    select_kit = PageObjects::Components::SelectKit.new(".ai-tool-editor__presets")
+    select_kit.expand
+    select_kit.select_row_by_value("exchange_rate")
+
+    find(".ai-tool-editor__next").click
+
+    expect(page.first(".parameter-row__required-toggle").checked?).to eq(true)
+    expect(page.first(".parameter-row__enum-toggle").checked?).to eq(false)
+
+    ensure_can_run_test
+
+    expect(page).not_to have_button(".ai-tool-editor__delete")
     find(".ai-tool-editor__save").click
 
     expect(page).to have_content("Tool saved")
 
     last_tool = AiTool.order("id desc").limit(1).first
     visit "/admin/plugins/discourse-ai/ai-tools/#{last_tool.id}"
 
+    ensure_can_run_test
+
     expect(page.first(".parameter-row__required-toggle").checked?).to eq(true)
     expect(page.first(".parameter-row__enum-toggle").checked?).to eq(false)