FIX: display search correctly, bug when stripping XML (#668)

- Display filtered search correctly, so it is not confusing
- When XML stripping, if a chunk was `<` it would crash
- SQL Helper improved to be better aware of Data Explorer

Author: SamSaffron

lib/ai_bot/personas/sql_helper.rb

@@ -58,7 +58,18 @@ def system_prompt
             - When generating SQL always use ```sql Markdown code blocks.
             - When generating SQL NEVER end SQL samples with a semicolon (;).
 
-            Eg:
+            - You also understand the special formatting rules for Data Explorer in Discourse.
+               - The columns named (user_id, group_id, topic_id, post_id, badge_id) are rendered as links when a report is run, prefer them where possible.
+               - You can define custom params to create flexible queries, example:
+                  -- [params]
+                  -- int :num = 1
+                  -- text :name
+
+                  SELECT :num, :name
+               - You support the types (integer, text, boolean, date)
+
+
+            - When generating SQL use markdown formatting for code blocks, example:
 
             ```sql
             select 1 from table

lib/ai_bot/tools/search.rb

@@ -4,6 +4,8 @@ module DiscourseAi
   module AiBot
     module Tools
       class Search < Tool
+        attr_reader :last_query
+
         MIN_SEMANTIC_RESULTS = 5
 
         class << self
@@ -95,40 +97,38 @@ def search_args
           parameters.slice(:category, :user, :order, :max_posts, :tags, :before, :after, :status)
         end
 
-        def invoke
-          search_string =
-            search_args.reduce(+parameters[:search_query].to_s) do |memo, (key, value)|
-              return memo if value.blank?
-              memo << " " << "#{key}:#{value}"
-            end
-
-          @last_query = search_string
+        def search_query
+          parameters[:search_query]
+        end
 
-          yield(I18n.t("discourse_ai.ai_bot.searching", query: search_string))
+        def invoke
+          search_terms = []
 
-          if options[:base_query].present?
-            search_string = "#{search_string} #{options[:base_query]}"
-          end
+          search_terms << options[:base_query] if options[:base_query].present?
+          search_terms << search_query.strip if search_query.present?
+          search_args.each { |key, value| search_terms << "#{key}:#{value}" if value.present? }
 
-          safe_search_string = search_string.to_s
           guardian = nil
-
           if options[:search_private] && context[:user]
             guardian = Guardian.new(context[:user])
           else
             guardian = Guardian.new
-            safe_search_string += " status:public"
+            search_terms << "status:public"
           end
 
-          results =
-            ::Search.execute(safe_search_string, search_type: :full_page, guardian: guardian)
+          search_string = search_terms.join(" ").to_s
+          @last_query = search_string
+
+          yield(I18n.t("discourse_ai.ai_bot.searching", query: search_string))
+
+          results = ::Search.execute(search_string, search_type: :full_page, guardian: guardian)
 
           max_results = calculate_max_results(llm)
           results_limit = parameters[:limit] || max_results
           results_limit = max_results if parameters[:limit].to_i > max_results
 
           should_try_semantic_search =
-            SiteSetting.ai_embeddings_semantic_search_enabled && parameters[:search_query].present?
+            SiteSetting.ai_embeddings_semantic_search_enabled && search_query.present?
 
           max_semantic_results = max_results / 4
           results_limit = results_limit - max_semantic_results if should_try_semantic_search

lib/completions/xml_tag_stripper.rb

@@ -17,6 +17,7 @@ def <<(text)
           end
         end
         @parsed.concat(parse_tags(text))
+
         @parsed, result = process_parsed(@parsed)
         result
       end
@@ -69,10 +70,15 @@ def parse_tags(text)
         while true
           before, after = text.split("<", 2)
 
-          parsed << { type: :text, content: before }
+          parsed << { type: :text, content: before } if before && !before.empty?
 
           break if after.nil?
 
+          if before.empty? && after.empty?
+            parsed << { type: :maybe_tag, content: "<" }
+            break
+          end
+
           tag, after = after.split(">", 2)
 
           is_end_tag = tag[0] == "/"

spec/lib/completions/xml_tag_stripper_spec.rb

@@ -29,6 +29,32 @@
     expect(result).to eq("\nhello\n")
   end
 
+  it "does not crash when we send a <" do
+    result = +""
+    result << (tag_stripper << "based:\n")
+    result << (tag_stripper << "<").to_s
+    result << (tag_stripper << " href")
+    result << (tag_stripper << ">")
+    result << (tag_stripper << "test ")
+
+    expect(result).to eq("based:\n< href>test ")
+  end
+
+  it "strips thinking correctly in a stream" do
+    result = +""
+    result << (tag_stripper << "hello")
+    result << (tag_stripper << "<").to_s
+    result << (tag_stripper << "thinking").to_s
+    result << (tag_stripper << ">").to_s
+    result << (tag_stripper << "test").to_s
+    result << (tag_stripper << "<").to_s
+    result << (tag_stripper << "/").to_s
+    result << (tag_stripper << "thinking").to_s
+    result << (tag_stripper << "> world")
+
+    expect(result).to eq("hello world")
+  end
+
   it "works when nesting unrelated tags it strips correctly" do
     text = <<~TEXT
       <thinking>

spec/lib/modules/ai_bot/tools/search_spec.rb

@@ -73,6 +73,8 @@
       results = search.invoke(&progress_blk)
       expect(results[:rows].length).to eq(1)
 
+      expect(search.last_query).to eq("#funny order:latest")
+
       GroupUser.create!(group: group, user: user)
 
       results = search.invoke(&progress_blk)