From e1a0eb61312aa9a720a1896d6362e30b3d01c73e Mon Sep 17 00:00:00 2001 From: Sam Date: Wed, 9 Oct 2024 08:17:45 +1100 Subject: [PATCH] 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. --- .../discourse_ai/admin/ai_tools_controller.rb | 2 +- app/models/ai_tool.rb | 77 +++++++++++++++++++ config/locales/server.en.yml | 4 +- lib/ai_bot/tool_runner.rb | 47 ++++++++++- lib/ai_bot/tools/custom.rb | 16 +++- spec/lib/modules/ai_bot/playground_spec.rb | 53 ++++++++++--- spec/system/ai_bot/tool_spec.rb | 41 +++++----- 7 files changed, 207 insertions(+), 33 deletions(-) diff --git a/app/controllers/discourse_ai/admin/ai_tools_controller.rb b/app/controllers/discourse_ai/admin/ai_tools_controller.rb index 4f231870..f5aa081b 100644 --- a/app/controllers/discourse_ai/admin/ai_tools_controller.rb +++ b/app/controllers/discourse_ai/admin/ai_tools_controller.rb @@ -49,7 +49,7 @@ module DiscourseAi 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 diff --git a/app/models/ai_tool.rb b/app/models/ai_tool.rb index b3638223..9267f827 100644 --- a/app/models/ai_tool.rb +++ b/app/models/ai_tool.rb @@ -93,6 +93,19 @@ class AiTool < ActiveRecord::Base * 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 @@ class AiTool < ActiveRecord::Base 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) { diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml index c42deb03..3f564bcb 100644 --- a/config/locales/server.en.yml +++ b/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..." diff --git a/lib/ai_bot/tool_runner.rb b/lib/ai_bot/tool_runner.rb index 8d725901..21c5d68f 100644 --- a/lib/ai_bot/tool_runner.rb +++ b/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 @@ module DiscourseAi 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 @@ module DiscourseAi 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 @@ module DiscourseAi ) 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", diff --git a/lib/ai_bot/tools/custom.rb b/lib/ai_bot/tools/custom.rb index db3ed379..5c77860f 100644 --- a/lib/ai_bot/tools/custom.rb +++ b/lib/ai_bot/tools/custom.rb @@ -30,8 +30,18 @@ module DiscourseAi 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 @@ module DiscourseAi 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" diff --git a/spec/lib/modules/ai_bot/playground_spec.rb b/spec/lib/modules/ai_bot/playground_spec.rb index c3f054df..5aae0b7e 100644 --- a/spec/lib/modules/ai_bot/playground_spec.rb +++ b/spec/lib/modules/ai_bot/playground_spec.rb @@ -94,21 +94,52 @@ RSpec.describe DiscourseAi::AiBot::Playground do let(:playground) { DiscourseAi::AiBot::Playground.new(bot) } - it "can force usage of a tool" do - tool_name = "custom-#{custom_tool.id}" - ai_persona.update!(tools: [[tool_name, nil, "force"]]) - responses = [function_call, "custom tool did stuff (maybe)"] + 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 - prompt = nil - DiscourseAi::Completions::Llm.with_prepared_responses(responses) do |_, _, _prompt| + 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) - prompt = _prompt + reply_post = playground.reply_to(new_post) + 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(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, true]]) + responses = [function_call, "custom tool did stuff (maybe)"] + + 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) + prompts = _prompts + end + + 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 diff --git a/spec/system/ai_bot/tool_spec.rb b/spec/system/ai_bot/tool_spec.rb index 4a97ed5c..2d250345 100644 --- a/spec/system/ai_bot/tool_spec.rb +++ b/spec/system/ai_bot/tool_spec.rb @@ -10,26 +10,9 @@ describe "AI Tool Management", type: :system do 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,7 +31,27 @@ describe "AI Tool Management", type: :system do 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") @@ -56,6 +59,8 @@ describe "AI Tool Management", type: :system do 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)