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:
parent
3170e14acb
commit
e1a0eb6131
|
@ -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
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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..."
|
||||
|
||||
|
|
|
@ -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",
|
||||
|
|
|
@ -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"
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
Loading…
Reference in New Issue