From c352054d4ef422caf334e0f17d94f8f162257f06 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 4 Nov 2024 10:07:17 +1100 Subject: [PATCH] FIX: encode parameters returned from LLMs correctly (#889) Fixes encoding of params on LLM function calls. Previously we would improperly return results if a function parameter returned an HTML tag. Additionally adds some missing HTTP verbs to tool calls. --- app/models/ai_tool.rb | 2 + lib/ai_bot/tool_runner.rb | 74 +++++++++++-------- lib/ai_bot/tools/tool.rb | 6 ++ .../anthropic_message_processor.rb | 2 +- lib/completions/endpoints/gemini.rb | 2 +- lib/completions/endpoints/open_ai.rb | 2 +- .../completions/endpoints/anthropic_spec.rb | 4 +- spec/lib/completions/endpoints/gemini_spec.rb | 28 +++++++ .../lib/completions/endpoints/open_ai_spec.rb | 4 +- spec/models/ai_tool_spec.rb | 50 +++++++------ 10 files changed, 112 insertions(+), 62 deletions(-) diff --git a/app/models/ai_tool.rb b/app/models/ai_tool.rb index 9267f827..ee2026cf 100644 --- a/app/models/ai_tool.rb +++ b/app/models/ai_tool.rb @@ -73,6 +73,8 @@ class AiTool < ActiveRecord::Base * Returns: * { status: number, body: string } * + * (also available: http.put, http.patch, http.delete) + * * Note: Max 20 HTTP requests per execution. * * 2. llm diff --git a/lib/ai_bot/tool_runner.rb b/lib/ai_bot/tool_runner.rb index 67adeed1..80ebbda4 100644 --- a/lib/ai_bot/tool_runner.rb +++ b/lib/ai_bot/tool_runner.rb @@ -44,10 +44,14 @@ module DiscourseAi end def framework_script + http_methods = %i[get post put patch delete].map { |method| <<~JS }.join("\n") + #{method}: function(url, options) { + return _http_#{method}(url, options); + }, + JS <<~JS const http = { - get: function(url, options) { return _http_get(url, options) }, - post: function(url, options) { return _http_post(url, options) }, + #{http_methods} }; const llm = { @@ -249,36 +253,44 @@ module DiscourseAi end, ) - mini_racer_context.attach( - "_http_post", - ->(url, options) do - begin - @http_requests_made += 1 - if @http_requests_made > MAX_HTTP_REQUESTS - raise TooManyRequestsError.new("Tool made too many HTTP requests") + %i[post put patch delete].each do |method| + mini_racer_context.attach( + "_http_#{method}", + ->(url, options) do + begin + @http_requests_made += 1 + if @http_requests_made > MAX_HTTP_REQUESTS + raise TooManyRequestsError.new("Tool made too many HTTP requests") + end + + self.running_attached_function = true + headers = (options && options["headers"]) || {} + body = options && options["body"] + + result = {} + DiscourseAi::AiBot::Tools::Tool.send_http_request( + url, + method: method, + headers: headers, + body: body, + ) do |response| + result[:body] = response.body + result[:status] = response.code.to_i + end + + result + rescue => e + p url + p options + p e + puts e.backtrace + raise e + ensure + self.running_attached_function = false end - - self.running_attached_function = true - headers = (options && options["headers"]) || {} - body = options && options["body"] - - result = {} - DiscourseAi::AiBot::Tools::Tool.send_http_request( - url, - method: :post, - headers: headers, - body: body, - ) do |response| - result[:body] = response.body - result[:status] = response.code.to_i - end - - result - ensure - self.running_attached_function = false - end - end, - ) + end, + ) + end end end end diff --git a/lib/ai_bot/tools/tool.rb b/lib/ai_bot/tools/tool.rb index 32688a18..89ab9d6c 100644 --- a/lib/ai_bot/tools/tool.rb +++ b/lib/ai_bot/tools/tool.rb @@ -188,6 +188,12 @@ module DiscourseAi request = FinalDestination::HTTP::Get.new(uri) elsif method == :post request = FinalDestination::HTTP::Post.new(uri) + elsif method == :put + request = FinalDestination::HTTP::Put.new(uri) + elsif method == :patch + request = FinalDestination::HTTP::Patch.new(uri) + elsif method == :delete + request = FinalDestination::HTTP::Delete.new(uri) end raise ArgumentError, "Invalid method: #{method}" if !request diff --git a/lib/completions/anthropic_message_processor.rb b/lib/completions/anthropic_message_processor.rb index 96a9a169..14be04ec 100644 --- a/lib/completions/anthropic_message_processor.rb +++ b/lib/completions/anthropic_message_processor.rb @@ -39,7 +39,7 @@ class DiscourseAi::Completions::AnthropicMessageProcessor ) params = JSON.parse(tool_call.raw_json, symbolize_names: true) - xml = params.map { |name, value| "<#{name}>#{value}" }.join("\n") + xml = params.map { |name, value| "<#{name}>#{CGI.escapeHTML(value)}" }.join("\n") node.at("tool_name").content = tool_call.name node.at("tool_id").content = tool_call.id diff --git a/lib/completions/endpoints/gemini.rb b/lib/completions/endpoints/gemini.rb index f2f8fa9b..f626e41a 100644 --- a/lib/completions/endpoints/gemini.rb +++ b/lib/completions/endpoints/gemini.rb @@ -179,7 +179,7 @@ module DiscourseAi if partial[:args] argument_fragments = partial[:args].reduce(+"") do |memo, (arg_name, value)| - memo << "\n<#{arg_name}>#{value}" + memo << "\n<#{arg_name}>#{CGI.escapeHTML(value)}" end argument_fragments << "\n" diff --git a/lib/completions/endpoints/open_ai.rb b/lib/completions/endpoints/open_ai.rb index 0677af94..b3960568 100644 --- a/lib/completions/endpoints/open_ai.rb +++ b/lib/completions/endpoints/open_ai.rb @@ -173,7 +173,7 @@ module DiscourseAi argument_fragments = json_args.reduce(+"") do |memo, (arg_name, value)| - memo << "\n<#{arg_name}>#{value}" + memo << "\n<#{arg_name}>#{CGI.escapeHTML(value)}" end argument_fragments << "\n" diff --git a/spec/lib/completions/endpoints/anthropic_spec.rb b/spec/lib/completions/endpoints/anthropic_spec.rb index 3ab38644..94d5d655 100644 --- a/spec/lib/completions/endpoints/anthropic_spec.rb +++ b/spec/lib/completions/endpoints/anthropic_spec.rb @@ -74,7 +74,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"y\\": \\"s"} } event: content_block_delta - data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"am"} } + data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"m"} } event: content_block_delta data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":" "} } @@ -118,7 +118,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do search - sam sam + s<a>m sam general toolu_01DjrShFRRHp9SnHYRFRc53F diff --git a/spec/lib/completions/endpoints/gemini_spec.rb b/spec/lib/completions/endpoints/gemini_spec.rb index b02d5bb5..2f602d3a 100644 --- a/spec/lib/completions/endpoints/gemini_spec.rb +++ b/spec/lib/completions/endpoints/gemini_spec.rb @@ -182,6 +182,34 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Gemini do expect(parsed[:tool_config]).to eq({ function_calling_config: { mode: "AUTO" } }) end + it "properly encodes tool calls" do + prompt = DiscourseAi::Completions::Prompt.new("Hello", tools: [echo_tool]) + + llm = DiscourseAi::Completions::Llm.proxy("custom:#{model.id}") + url = "#{model.url}:generateContent?key=123" + + response_json = { "functionCall" => { name: "echo", args: { text: "ydney" } } } + response = gemini_mock.response(response_json, tool_call: true).to_json + + stub_request(:post, url).to_return(status: 200, body: response) + + response = llm.generate(prompt, user: user) + + expected = (<<~XML).strip + + + echo + + <S>ydney + + tool_0 + + + XML + + expect(response.strip).to eq(expected) + end + it "Supports Vision API" do prompt = DiscourseAi::Completions::Prompt.new( diff --git a/spec/lib/completions/endpoints/open_ai_spec.rb b/spec/lib/completions/endpoints/open_ai_spec.rb index 44ff136c..60df1d67 100644 --- a/spec/lib/completions/endpoints/open_ai_spec.rb +++ b/spec/lib/completions/endpoints/open_ai_spec.rb @@ -294,7 +294,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::OpenAi do type: "function", function: { name: "echo", - arguments: "{\"text\":\"hello\"}", + arguments: "{\"text\":\"hllo\"}", }, }, ], @@ -325,7 +325,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::OpenAi do echo - hello + h<e>llo call_I8LKnoijVuhKOM85nnEQgWwd diff --git a/spec/models/ai_tool_spec.rb b/spec/models/ai_tool_spec.rb index 7e0412f6..a0d9fc01 100644 --- a/spec/models/ai_tool_spec.rb +++ b/spec/models/ai_tool_spec.rb @@ -38,35 +38,37 @@ RSpec.describe AiTool do 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 }) - } - ); + it "can perform HTTP requests with various verbs" do + %i[post put delete patch].each do |verb| + script = <<~JS + function invoke(params) { + result = http.#{verb}("https://example.com/api", + { + headers: { TestHeader: "TestValue" }, + body: JSON.stringify({ data: params.data }) + } + ); - return result.body; - } - JS + return result.body; + } + JS - tool = create_tool(script: script) - runner = tool.runner({ "data" => "test data" }, llm: nil, bot_user: nil, context: {}) + 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: {}) + stub_request(verb, "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 + result = runner.invoke - expect(result).to eq("Success") + expect(result).to eq("Success") + end end it "can perform GET HTTP requests, with 1 param" do