From 0069256efd9a3b95d81345e5d3a2fb441be90e5a Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Mon, 13 May 2024 19:40:11 +1000
Subject: [PATCH] FIX: improve function call parsing (#613)

- support " / ' wrapped values
- coerce integer to integer
- enforce enum at boundary
---
 lib/ai_bot/personas/persona.rb                |  23 ++++
 lib/completions/endpoints/aws_bedrock.rb      |   3 +-
 .../modules/ai_bot/personas/persona_spec.rb   | 104 ++++++++++++++++++
 3 files changed, 129 insertions(+), 1 deletion(-)

diff --git a/lib/ai_bot/personas/persona.rb b/lib/ai_bot/personas/persona.rb
index 24010190..b7f73f29 100644
--- a/lib/ai_bot/personas/persona.rb
+++ b/lib/ai_bot/personas/persona.rb
@@ -208,6 +208,15 @@ module DiscourseAi
                 rescue JSON::ParserError
                   [value.to_s]
                 end
+            elsif param[:type] == "string" && value
+              value = strip_quotes(value).to_s
+            elsif param[:type] == "integer" && value
+              value = strip_quotes(value).to_i
+            end
+
+            if param[:enum] && value && !param[:enum].include?(value)
+              # invalid enum value
+              value = nil
             end
 
             arguments[name.to_sym] = value if value
@@ -223,6 +232,20 @@ module DiscourseAi
           )
         end
 
+        def strip_quotes(value)
+          if value.is_a?(String)
+            if value.start_with?('"') && value.end_with?('"')
+              value = value[1..-2]
+            elsif value.start_with?("'") && value.end_with?("'")
+              value = value[1..-2]
+            else
+              value
+            end
+          else
+            value
+          end
+        end
+
         def rag_fragments_prompt(conversation_context, llm:, user:)
           upload_refs =
             UploadReference.where(target_id: id, target_type: "AiPersona").pluck(:upload_id)
diff --git a/lib/completions/endpoints/aws_bedrock.rb b/lib/completions/endpoints/aws_bedrock.rb
index 7163ff47..9d15081d 100644
--- a/lib/completions/endpoints/aws_bedrock.rb
+++ b/lib/completions/endpoints/aws_bedrock.rb
@@ -121,7 +121,8 @@ module DiscourseAi
           i = 0
           while decoded
             parsed = JSON.parse(decoded.payload.string)
-            messages << Base64.decode64(parsed["bytes"])
+            # perhaps some control message we can just ignore
+            messages << Base64.decode64(parsed["bytes"]) if parsed && parsed["bytes"]
 
             decoded, _done = @decoder.decode_chunk
 
diff --git a/spec/lib/modules/ai_bot/personas/persona_spec.rb b/spec/lib/modules/ai_bot/personas/persona_spec.rb
index b803297e..145e0dbc 100644
--- a/spec/lib/modules/ai_bot/personas/persona_spec.rb
+++ b/spec/lib/modules/ai_bot/personas/persona_spec.rb
@@ -70,6 +70,110 @@ RSpec.describe DiscourseAi::AiBot::Personas::Persona do
     expect(tools.find { |t| t[:name] == "image" }).to be_nil
   end
 
+  it "can parse string that are wrapped in quotes" do
+    SiteSetting.ai_stability_api_key = "123"
+    xml = <<~XML
+      <function_calls>
+        <invoke>
+        <tool_name>image</tool_name>
+        <tool_id>call_JtYQMful5QKqw97XFsHzPweB</tool_id>
+        <parameters>
+        <prompts>["cat oil painting", "big car"]</prompts>
+        <aspect_ratio>"16:9"</aspect_ratio>
+        </parameters>
+        </invoke>
+        <invoke>
+        <tool_name>image</tool_name>
+        <tool_id>call_JtYQMful5QKqw97XFsHzPweB</tool_id>
+        <parameters>
+        <prompts>["cat oil painting", "big car"]</prompts>
+        <aspect_ratio>'16:9'</aspect_ratio>
+        </parameters>
+        </invoke>
+      </function_calls>
+    XML
+
+    image1, image2 =
+      tools =
+        DiscourseAi::AiBot::Personas::Artist.new.find_tools(
+          xml,
+          bot_user: nil,
+          llm: nil,
+          context: nil,
+        )
+    expect(image1.parameters[:prompts]).to eq(["cat oil painting", "big car"])
+    expect(image1.parameters[:aspect_ratio]).to eq("16:9")
+    expect(image2.parameters[:aspect_ratio]).to eq("16:9")
+
+    expect(tools.length).to eq(2)
+  end
+
+  it "enforces enums" do
+    xml = <<~XML
+      <function_calls>
+        <invoke>
+        <tool_name>search</tool_name>
+        <tool_id>call_JtYQMful5QKqw97XFsHzPweB</tool_id>
+        <parameters>
+        <max_posts>"3.2"</max_posts>
+        <status>cow</status>
+        <foo>bar</foo>
+        </parameters>
+        </invoke>
+        <invoke>
+        <tool_name>search</tool_name>
+        <tool_id>call_JtYQMful5QKqw97XFsHzPweB</tool_id>
+        <parameters>
+        <max_posts>"3.2"</max_posts>
+        <status>open</status>
+        <foo>bar</foo>
+        </parameters>
+        </invoke>
+      </function_calls>
+    XML
+
+    search1, search2 =
+      tools =
+        DiscourseAi::AiBot::Personas::General.new.find_tools(
+          xml,
+          bot_user: nil,
+          llm: nil,
+          context: nil,
+        )
+
+    expect(search1.parameters.key?(:status)).to eq(false)
+    expect(search2.parameters[:status]).to eq("open")
+  end
+
+  it "can coerce integers" do
+    xml = <<~XML
+      <function_calls>
+        <invoke>
+        <tool_name>search</tool_name>
+        <tool_id>call_JtYQMful5QKqw97XFsHzPweB</tool_id>
+        <parameters>
+        <max_posts>"3.2"</max_posts>
+        <search_query>hello world</search_query>
+        <foo>bar</foo>
+        </parameters>
+        </invoke>
+      </function_calls>
+    XML
+
+    search, =
+      tools =
+        DiscourseAi::AiBot::Personas::General.new.find_tools(
+          xml,
+          bot_user: nil,
+          llm: nil,
+          context: nil,
+        )
+
+    expect(search.parameters[:max_posts]).to eq(3)
+    expect(search.parameters[:search_query]).to eq("hello world")
+    expect(search.parameters.key?(:foo)).to eq(false)
+  end
+
   it "can correctly parse arrays in tools" do
     SiteSetting.ai_openai_api_key = "123"