From 0cbbf130b95d3cfee0b60c51efb1ddf3e919ccdf Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Thu, 11 Apr 2024 11:24:22 +1000
Subject: [PATCH] FIX: never mention the word JSON in tool preamble (#572)

Just having the word JSON can confuse models when we expect them
to deal solely in XML

Instead provide an example of how string arrays should be returned

Technically the tool framework supports int arrays and more, but
our current implementation only does string arrays.

Also tune the prompt construction not to give any tips about arrays
if none exist
---
 lib/completions/dialects/dialect.rb           | 41 +++++++++++--------
 spec/lib/completions/dialects/dialect_spec.rb | 32 +++++++++++++++
 2 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/lib/completions/dialects/dialect.rb b/lib/completions/dialects/dialect.rb
index d84d9e38..9361d8ba 100644
--- a/lib/completions/dialects/dialect.rb
+++ b/lib/completions/dialects/dialect.rb
@@ -33,7 +33,17 @@ module DiscourseAi
             raise NotImplemented
           end
 
-          def tool_preamble
+          def tool_preamble(include_array_tip: true)
+            array_tip =
+              if include_array_tip
+                <<~TEXT
+                If a parameter type is an array, return an array of values. For example:
+                <$PARAMETER_NAME>["one","two","three"]</$PARAMETER_NAME>
+              TEXT
+              else
+                ""
+              end
+
             <<~TEXT
               In this environment you have access to a set of tools you can use to answer the user's question.
               You may call them like this.
@@ -47,16 +57,12 @@ module DiscourseAi
               </parameters>
               </invoke>
               </function_calls>
-
-              If a parameter type is an array, return a JSON array of values. For example:
-              [1,"two",3.0]
-
+              #{array_tip}
               If you wish to call multiple function in one reply, wrap multiple <invoke>
               block in a single <function_calls> block.
 
               Always prefer to lead with tool calls, if you need to execute any.
               Avoid all niceties prior to tool calls, Eg: "Let me look this up for you.." etc.
-
               Here are the complete list of tools available:
             TEXT
           end
@@ -148,6 +154,19 @@ module DiscourseAi
 
         attr_reader :prompt
 
+        def build_tools_prompt
+          return "" if prompt.tools.blank?
+
+          has_arrays =
+            prompt.tools.any? { |tool| tool[:parameters].any? { |p| p[:type] == "array" } }
+
+          (<<~TEXT).strip
+            #{self.class.tool_preamble(include_array_tip: has_arrays)}
+            <tools>
+            #{tools}</tools>
+          TEXT
+        end
+
         private
 
         attr_reader :model_name, :opts
@@ -211,16 +230,6 @@ module DiscourseAi
         def calculate_message_token(msg)
           self.class.tokenizer.size(msg[:content].to_s)
         end
-
-        def build_tools_prompt
-          return "" if prompt.tools.blank?
-
-          (<<~TEXT).strip
-            #{self.class.tool_preamble}
-            <tools>
-            #{tools}</tools>
-          TEXT
-        end
       end
     end
   end
diff --git a/spec/lib/completions/dialects/dialect_spec.rb b/spec/lib/completions/dialects/dialect_spec.rb
index c54d1838..db481e94 100644
--- a/spec/lib/completions/dialects/dialect_spec.rb
+++ b/spec/lib/completions/dialects/dialect_spec.rb
@@ -17,6 +17,38 @@ class TestDialect < DiscourseAi::Completions::Dialects::Dialect
 end
 
 RSpec.describe DiscourseAi::Completions::Dialects::Dialect do
+  describe "#build_tools_prompt" do
+    it "can exclude array instructions" do
+      prompt = DiscourseAi::Completions::Prompt.new("12345")
+      prompt.tools = [
+        {
+          name: "weather",
+          description: "lookup weather in a city",
+          parameters: [{ name: "city", type: "string", description: "city name", required: true }],
+        },
+      ]
+
+      dialect = TestDialect.new(prompt, "test")
+
+      expect(dialect.build_tools_prompt).not_to include("array")
+    end
+
+    it "can include array instructions" do
+      prompt = DiscourseAi::Completions::Prompt.new("12345")
+      prompt.tools = [
+        {
+          name: "weather",
+          description: "lookup weather in a city",
+          parameters: [{ name: "city", type: "array", description: "city names", required: true }],
+        },
+      ]
+
+      dialect = TestDialect.new(prompt, "test")
+
+      expect(dialect.build_tools_prompt).to include("array")
+    end
+  end
+
   describe "#trim_messages" do
     it "should trim tool messages if tool_calls are trimmed" do
       prompt = DiscourseAi::Completions::Prompt.new("12345")