- halt after tools

- post streamer ensures we don't have half completed stuff on screen when a tool is slow
- reimplemnt xml tools to have a more relaxed parse

Author: SamSaffron

lib/ai_bot/bot.rb

@@ -106,6 +106,8 @@ def reply(context, &update_blk)
           tool_found = false
           force_tool_if_needed(prompt, context)
 
+          tool_halted = false
+
           result =
             llm.generate(prompt, feature_name: "bot", **llm_kwargs) do |partial, cancel|
               tool = persona.find_tool(partial, bot_user: user, llm: llm, context: context)
@@ -122,7 +124,12 @@ def reply(context, &update_blk)
                 process_tool(tool, raw_context, llm, cancel, update_blk, prompt, context)
                 tools_ran += 1
                 ongoing_chain &&= tool.chain_next_response?
+
+                if !tool.chain_next_response?
+                  tool_halted = true
+                end
               else
+                next if tool_halted
                 needs_newlines = true
                 if partial.is_a?(DiscourseAi::Completions::ToolCall)
                   Rails.logger.warn("DiscourseAi: Tool not found: #{partial.name}")

lib/ai_bot/playground.rb

@@ -399,7 +399,7 @@ def reply_to(post, custom_instructions: nil, &blk)
         PostCustomPrompt.none
 
         reply = +""
-        start = Time.now
+        post_streamer = nil
 
         post_type =
           post.post_type == Post.types[:whisper] ? Post.types[:whisper] : Post.types[:regular]
@@ -448,6 +448,8 @@ def reply_to(post, custom_instructions: nil, &blk)
 
         context[:skip_tool_details] ||= !bot.persona.class.tool_details
 
+        post_streamer = PostStreamer.new(delay: Rails.env.test? ? 0 : 0.5) if stream_reply
+
         new_custom_prompts =
           bot.reply(context) do |partial, cancel, placeholder, type|
             reply << partial
@@ -461,22 +463,20 @@ def reply_to(post, custom_instructions: nil, &blk)
               reply_post.update!(raw: reply, cooked: PrettyText.cook(reply))
             end
 
-            if stream_reply
-              # Minor hack to skip the delay during tests.
-              if placeholder.blank?
-                next if (Time.now - start < 0.5) && !Rails.env.test?
-                start = Time.now
-              end
-
-              Discourse.redis.expire(redis_stream_key, 60)
-
-              publish_update(reply_post, { raw: raw })
+            if post_streamer
+              post_streamer.run_later {
+                Discourse.redis.expire(redis_stream_key, 60)
+                publish_update(reply_post, { raw: raw })
+              }
             end
           end
 
         return if reply.blank?
 
         if stream_reply
+          post_streamer.finish
+          post_streamer = nil
+
           # land the final message prior to saving so we don't clash
           reply_post.cooked = PrettyText.cook(reply)
           publish_final_update(reply_post)
@@ -514,6 +514,7 @@ def reply_to(post, custom_instructions: nil, &blk)
 
         reply_post
       ensure
+        post_streamer&.finish(skip_callback: true)
         publish_final_update(reply_post) if stream_reply
         if reply_post && post.post_number == 1 && post.topic.private_message?
           title_playground(reply_post)

lib/ai_bot/post_streamer.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module AiBot
+    class PostStreamer
+      def initialize(delay: 0.5)
+        @mutex = Mutex.new
+        @callback = nil
+        @delay = delay
+        @done = false
+      end
+
+      def run_later(&callback)
+        @mutex.synchronize { @callback = callback }
+        ensure_worker!
+      end
+
+      def finish(skip_callback: false)
+        @mutex.synchronize do
+          @callback&.call if skip_callback
+          @callback = nil
+          @done = true
+        end
+
+        begin
+          @worker_thread&.wakeup
+        rescue StandardError
+          ThreadError
+        end
+        @worker_thread&.join
+        @worker_thread = nil
+      end
+
+      private
+
+      def run
+        while !@done
+          @mutex.synchronize do
+            callback = @callback
+            @callback = nil
+            callback&.call
+          end
+          sleep @delay
+        end
+      end
+
+      def ensure_worker!
+        return if @worker_thread
+        @mutex.synchronize do
+          return if @worker_thread
+          db = RailsMultisite::ConnectionManagement.current_db
+          @worker_thread =
+            Thread.new { RailsMultisite::ConnectionManagement.with_connection(db) { run } }
+        end
+      end
+    end
+  end
+end

lib/ai_bot/tools/create_artifact.rb

@@ -39,6 +39,7 @@ def self.signature
         end
 
         def invoke
+          yield parameters[:name] || "Web Artifact"
           # Get the current post from context
           post = Post.find_by(id: context[:post_id])
           return error_response("No post context found") unless post

lib/completions/xml_tool_processor.rb

@@ -62,31 +62,14 @@ def <<(text)
       def finish
         return [] if @function_buffer.blank?
 
-        xml = Nokogiri::HTML5.fragment(@function_buffer)
-        normalize_function_ids!(xml)
-        last_invoke = xml.at("invoke:last")
-        if last_invoke
-          last_invoke.next_sibling.remove while last_invoke.next_sibling
-          xml.at("invoke:last").add_next_sibling("\n") if !last_invoke.next_sibling
+        idx = -1
+        parse_malformed_xml(@function_buffer).map do |tool|
+          ToolCall.new(
+            id: "tool_#{idx += 1}",
+            name: tool[:tool_name],
+            parameters: tool[:parameters]
+          )
         end
-
-        xml
-          .css("invoke")
-          .map do |invoke|
-            tool_name = invoke.at("tool_name").content.force_encoding("UTF-8")
-            tool_id = invoke.at("tool_id").content.force_encoding("UTF-8")
-            parameters = {}
-            invoke
-              .at("parameters")
-              &.children
-              &.each do |node|
-                next if node.text?
-                name = node.name
-                value = node.content.to_s
-                parameters[name.to_sym] = value.to_s.force_encoding("UTF-8")
-              end
-            ToolCall.new(id: tool_id, name: tool_name, parameters: parameters)
-          end
       end
 
       def should_cancel?
@@ -95,6 +78,40 @@ def should_cancel?
 
       private
 
+      def parse_malformed_xml(input)
+        input
+          .scan(
+            %r{
+      <invoke>
+        \s*
+        <tool_name>
+          ([^<]+)
+        </tool_name>
+        \s*
+        <parameters>
+          (.*?)
+        </parameters>
+        \s*
+      </invoke>
+    }mx,
+          )
+          .map do |tool_name, params|
+            {
+              tool_name: tool_name.strip,
+              parameters:
+                params
+                  .scan(%r{
+          <([^>]+)>
+            (.*?)
+          </\1>
+        }mx)
+                  .each_with_object({}) do |(name, value), hash|
+                    hash[name.to_sym] = value.gsub(/^<!\[CDATA\[|\]\]>$/, "")
+                  end,
+            }
+          end
+      end
+
       def normalize_function_ids!(function_buffer)
         function_buffer
           .css("invoke")

spec/lib/completions/xml_tool_processor_spec.rb

@@ -12,6 +12,26 @@
     expect(processor.should_cancel?).to eq(false)
   end
 
+  it "can handle mix and match xml cause tool llms may not encode" do
+    xml = (<<~XML).strip
+      <function_calls>
+      <invoke>
+      <tool_name>hello</tool_name>
+      <parameters>
+       <hello>world <sam>sam</sam></hello>
+       <test><![CDATA[</h1>\n</div>\n]]></test>
+      </parameters>
+      </invoke>
+    XML
+
+    result = []
+    result << (processor << xml)
+    result << (processor.finish)
+
+    tool_call = result.last.first
+    expect(tool_call.parameters).to eq(hello: "world <sam>sam</sam>", test: "</h1>\n</div>\n")
+  end
+
   it "is usable for simple single message mode" do
     xml = (<<~XML).strip
       hello
@@ -149,8 +169,7 @@
     result << (processor.finish)
 
     # Should just do its best to parse the XML
-    tool_call =
-      DiscourseAi::Completions::ToolCall.new(id: "tool_0", name: "test", parameters: { param: "" })
+    tool_call = DiscourseAi::Completions::ToolCall.new(id: "tool_0", name: "test", parameters: {})
     expect(result).to eq([["text"], [tool_call]])
   end