lib/boxcars/train/xml_train.rb in boxcars-0.3.2 vs lib/boxcars/train/xml_train.rb in boxcars-0.3.3

@@ -11,10 +11,11 @@
     # @param prompt [Boxcars::Prompt] The prompt to use.
     # @param engine [Boxcars::Engine] The engine to use for this train.
     # @param kwargs [Hash] Additional arguments including: name, description, top_k, return_direct, and stop
     # @abstract
     def initialize(boxcars:, prompt:, engine: nil, **kwargs)
+      @using_xml = true
       super
     end
 
     def init_prefixes
       @thought_prefix ||= "<thought>"
@@ -39,10 +40,12 @@
     def prediction_additional(_inputs)
       { boxcars_xml: boxcars_xml, next_actions: next_actions }.merge super
     end
 
     def build_output(text)
-      if text =~ /#{close_tag(thought_prefix)}/
+      if text.end_with?("</usetool>")
+        "<data>#{engine_prefix}#{text}</output></data>"
+      elsif text =~ /#{close_tag(thought_prefix)}/
         "<data>#{engine_prefix}#{text}</data>"
       else
         "<data>#{text}</data>"
       end
     end