improve test

Author: SamSaffron

lib/completions/dialects/xml_tools.rb

@@ -54,6 +54,9 @@ def instructions
             end
         end
 
+        DONE_MESSAGE =
+          "Regardless of what you think, REPLY IMMEDIATELY, WITHOUT MAKING ANY FURTHER TOOL CALLS, YOU ARE OUT OF TOOL CALL QUOTA!"
+
         def from_raw_tool(raw_message)
           result = (<<~TEXT).strip
             <function_results>
@@ -67,8 +70,7 @@ def from_raw_tool(raw_message)
           TEXT
 
           if @injecting_done
-            result +
-              "\n\nRegardless of what you think, REPLY IMMEDIATELY, WITHOUT MAKING ANY FURTHER TOOL CALLS, YOU ARE OUT OF TOOL CALL QUOTA!"
+            "#{result}\n\n#{DONE_MESSAGE}"
           else
             result
           end

spec/lib/completions/dialects/dialect_spec.rb

@@ -7,6 +7,18 @@ def trim(messages)
     trim_messages(messages)
   end
 
+  def system_msg(msg)
+    msg
+  end
+
+  def user_msg(msg)
+    msg
+  end
+
+  def model_msg(msg)
+    msg
+  end
+
   def tokenizer
     DiscourseAi::Tokenizer::OpenAiTokenizer
   end
@@ -15,6 +27,57 @@ def tokenizer
 RSpec.describe DiscourseAi::Completions::Dialects::Dialect do
   fab!(:llm_model)
 
+  describe "#translate" do
+    let(:five_token_msg) { "This represents five tokens." }
+    let(:tools) do
+      [
+        {
+          name: "echo",
+          description: "echo a string",
+          parameters: [
+            { name: "text", type: "string", description: "string to echo", required: true },
+          ],
+        },
+      ]
+    end
+
+    it "injects done message when tool_choice is :none and last message follows tool pattern" do
+      tool_call_prompt = { name: "echo", arguments: { text: "test message" } }
+
+      prompt = DiscourseAi::Completions::Prompt.new("System instructions", tools: tools)
+      prompt.push(type: :user, content: "echo test message")
+      prompt.push(type: :tool_call, content: tool_call_prompt.to_json, id: "123", name: "echo")
+      prompt.push(type: :tool, content: "test message".to_json, name: "echo", id: "123")
+      prompt.tool_choice = :none
+
+      dialect = TestDialect.new(prompt, llm_model)
+      dialect.max_prompt_tokens = 100 # Set high enough to avoid trimming
+
+      translated = dialect.translate
+
+      expect(translated).to eq(
+        [
+          { type: :system, content: "System instructions" },
+          { type: :user, content: "echo test message" },
+          {
+            type: :tool_call,
+            content:
+              "<function_calls>\n<invoke>\n<tool_name>echo</tool_name>\n<parameters>\n<text>test message</text>\n</parameters>\n</invoke>\n</function_calls>",
+            id: "123",
+            name: "echo",
+          },
+          {
+            type: :tool,
+            id: "123",
+            name: "echo",
+            content:
+              "<function_results>\n<result>\n<tool_name>echo</tool_name>\n<json>\n\"test message\"\n</json>\n</result>\n</function_results>\n\n#{::DiscourseAi::Completions::Dialects::XmlTools::DONE_MESSAGE}",
+          },
+        ],
+      )
+    end
+  end
+
   describe "#trim_messages" do
     let(:five_token_msg) { "This represents five tokens." }