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.
This commit is contained in:
Sam 2024-10-09 08:17:45 +11:00 committed by GitHub
parent 3170e14acb
commit e1a0eb6131
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 207 additions and 33 deletions

View File

@ -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

View File

@ -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) {

View File

@ -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..."

View File

@ -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",

View File

@ -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"

View File

@ -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

View File

@ -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)