295 lines
8.1 KiB
Ruby
295 lines
8.1 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
RSpec.describe AiTool do
|
|
fab!(:llm_model) { Fabricate(:llm_model, name: "claude-2") }
|
|
let(:llm) { DiscourseAi::Completions::Llm.proxy("custom:#{llm_model.id}") }
|
|
|
|
def create_tool(
|
|
parameters: nil,
|
|
script: nil,
|
|
rag_chunk_tokens: nil,
|
|
rag_chunk_overlap_tokens: nil
|
|
)
|
|
AiTool.create!(
|
|
name: "test",
|
|
description: "test",
|
|
parameters: parameters || [{ name: "query", type: "string", desciption: "perform a search" }],
|
|
script: script || "function invoke(params) { return params; }",
|
|
created_by_id: 1,
|
|
summary: "Test tool summary",
|
|
rag_chunk_tokens: rag_chunk_tokens || 374,
|
|
rag_chunk_overlap_tokens: rag_chunk_overlap_tokens || 10,
|
|
)
|
|
end
|
|
|
|
it "it can run a basic tool" do
|
|
tool = create_tool
|
|
|
|
expect(tool.signature).to eq(
|
|
{
|
|
name: "test",
|
|
description: "test",
|
|
parameters: [{ name: "query", type: "string", desciption: "perform a search" }],
|
|
},
|
|
)
|
|
|
|
runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil, context: {})
|
|
|
|
expect(runner.invoke).to eq("query" => "test")
|
|
end
|
|
|
|
it "can perform POST HTTP requests" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
result = http.post("https://example.com/api",
|
|
{
|
|
headers: { TestHeader: "TestValue" },
|
|
body: JSON.stringify({ data: params.data })
|
|
}
|
|
);
|
|
|
|
return result.body;
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
runner = tool.runner({ "data" => "test data" }, llm: nil, bot_user: nil, context: {})
|
|
|
|
stub_request(:post, "https://example.com/api").with(
|
|
body: "{\"data\":\"test data\"}",
|
|
headers: {
|
|
"Accept" => "*/*",
|
|
"Testheader" => "TestValue",
|
|
"User-Agent" => "Discourse AI Bot 1.0 (https://www.discourse.org)",
|
|
},
|
|
).to_return(status: 200, body: "Success", headers: {})
|
|
|
|
result = runner.invoke
|
|
|
|
expect(result).to eq("Success")
|
|
end
|
|
|
|
it "can perform GET HTTP requests, with 1 param" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
result = http.get("https://example.com/" + params.query);
|
|
return result.body;
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil, context: {})
|
|
|
|
stub_request(:get, "https://example.com/test").with(
|
|
headers: {
|
|
"Accept" => "*/*",
|
|
"User-Agent" => "Discourse AI Bot 1.0 (https://www.discourse.org)",
|
|
},
|
|
).to_return(status: 200, body: "Hello World", headers: {})
|
|
|
|
result = runner.invoke
|
|
|
|
expect(result).to eq("Hello World")
|
|
end
|
|
|
|
it "is limited to MAX http requests" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
let i = 0;
|
|
while (i < 21) {
|
|
http.get("https://example.com/");
|
|
i += 1;
|
|
}
|
|
return "will not happen";
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
runner = tool.runner({}, llm: nil, bot_user: nil, context: {})
|
|
|
|
stub_request(:get, "https://example.com/").to_return(
|
|
status: 200,
|
|
body: "Hello World",
|
|
headers: {
|
|
},
|
|
)
|
|
|
|
expect { runner.invoke }.to raise_error(DiscourseAi::AiBot::ToolRunner::TooManyRequestsError)
|
|
end
|
|
|
|
it "can perform GET HTTP requests" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
result = http.get("https://example.com/" + params.query,
|
|
{ headers: { TestHeader: "TestValue" } }
|
|
);
|
|
|
|
return result.body;
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil, context: {})
|
|
|
|
stub_request(:get, "https://example.com/test").with(
|
|
headers: {
|
|
"Accept" => "*/*",
|
|
"Testheader" => "TestValue",
|
|
"User-Agent" => "Discourse AI Bot 1.0 (https://www.discourse.org)",
|
|
},
|
|
).to_return(status: 200, body: "Hello World", headers: {})
|
|
|
|
result = runner.invoke
|
|
|
|
expect(result).to eq("Hello World")
|
|
end
|
|
|
|
it "will not timeout on slow HTTP reqs" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
result = http.get("https://example.com/" + params.query,
|
|
{ headers: { TestHeader: "TestValue" } }
|
|
);
|
|
|
|
return result.body;
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil, context: {})
|
|
|
|
stub_request(:get, "https://example.com/test").to_return do
|
|
sleep 0.01
|
|
{ status: 200, body: "Hello World", headers: {} }
|
|
end
|
|
|
|
runner.timeout = 5
|
|
|
|
result = runner.invoke
|
|
|
|
expect(result).to eq("Hello World")
|
|
end
|
|
|
|
it "has access to llm truncation tools" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
return llm.truncate("Hello World", 1);
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
|
|
runner = tool.runner({}, llm: llm, bot_user: nil, context: {})
|
|
result = runner.invoke
|
|
|
|
expect(result).to eq("Hello")
|
|
end
|
|
|
|
it "can timeout slow JS" do
|
|
script = <<~JS
|
|
function invoke(params) {
|
|
while (true) {}
|
|
}
|
|
JS
|
|
|
|
tool = create_tool(script: script)
|
|
runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil, context: {})
|
|
|
|
runner.timeout = 5
|
|
|
|
result = runner.invoke
|
|
expect(result[:error]).to eq("Script terminated due to timeout")
|
|
end
|
|
|
|
context "when defining RAG fragments" do
|
|
before do
|
|
SiteSetting.authorized_extensions = "txt"
|
|
SiteSetting.ai_embeddings_enabled = true
|
|
SiteSetting.ai_embeddings_discourse_service_api_endpoint = "http://test.com"
|
|
SiteSetting.ai_embeddings_model = "bge-large-en"
|
|
|
|
Jobs.run_immediately!
|
|
end
|
|
|
|
def create_upload(content, filename)
|
|
upload = nil
|
|
Tempfile.create(filename) do |file|
|
|
file.write(content)
|
|
file.rewind
|
|
|
|
upload = UploadCreator.new(file, filename).create_for(Discourse.system_user.id)
|
|
end
|
|
upload
|
|
end
|
|
|
|
def stub_embeddings
|
|
# this is a trick, we get ever increasing embeddings, this gives us in turn
|
|
# 100% consistent search results
|
|
@counter = 0
|
|
stub_request(:post, "http://test.com/api/v1/classify").to_return(
|
|
status: 200,
|
|
body: lambda { |req| ([@counter += 1] * 1024).to_json },
|
|
headers: {
|
|
},
|
|
)
|
|
end
|
|
|
|
it "allows search within uploads" do
|
|
stub_embeddings
|
|
|
|
upload1 = create_upload(<<~TXT, "test.txt")
|
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
|
|
TXT
|
|
|
|
upload2 = create_upload(<<~TXT, "test.txt")
|
|
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
|
|
TXT
|
|
|
|
tool = create_tool(rag_chunk_tokens: 10, rag_chunk_overlap_tokens: 4, script: <<~JS)
|
|
function invoke(params) {
|
|
let result1 = index.search("testing a search", { limit: 1 });
|
|
let result2 = index.search("testing another search", { limit: 3, filenames: ["test.txt"] });
|
|
|
|
return [result1, result2];
|
|
}
|
|
JS
|
|
|
|
RagDocumentFragment.link_target_and_uploads(tool, [upload1.id, upload2.id])
|
|
|
|
result = tool.runner({}, llm: nil, bot_user: nil, context: {}).invoke
|
|
|
|
expected = [
|
|
[{ "fragment" => "44 45 46 47 48 49 50", "metadata" => nil }],
|
|
[
|
|
{ "fragment" => "44 45 46 47 48 49 50", "metadata" => nil },
|
|
{ "fragment" => "36 37 38 39 40 41 42 43 44 45", "metadata" => nil },
|
|
{ "fragment" => "30 31 32 33 34 35 36 37", "metadata" => nil },
|
|
],
|
|
]
|
|
|
|
expect(result).to eq(expected)
|
|
|
|
# will force a reindex
|
|
tool.rag_chunk_tokens = 5
|
|
tool.rag_chunk_overlap_tokens = 2
|
|
tool.save!
|
|
|
|
# this part of the API is a bit awkward, maybe we should do it
|
|
# automatically
|
|
RagDocumentFragment.update_target_uploads(tool, [upload1.id, upload2.id])
|
|
result = tool.runner({}, llm: nil, bot_user: nil, context: {}).invoke
|
|
|
|
expected = [
|
|
[{ "fragment" => "48 49 50", "metadata" => nil }],
|
|
[
|
|
{ "fragment" => "48 49 50", "metadata" => nil },
|
|
{ "fragment" => "45 46 47", "metadata" => nil },
|
|
{ "fragment" => "42 43 44", "metadata" => nil },
|
|
],
|
|
]
|
|
|
|
expect(result).to eq(expected)
|
|
end
|
|
end
|
|
end
|