From 0dba6623a0a420b95b3b0afb70255310a11d67ee Mon Sep 17 00:00:00 2001
From: Rafael dos Santos Silva <xfalcox@gmail.com>
Date: Fri, 9 Feb 2024 14:49:28 -0300
Subject: [PATCH] FIX: Better AI chat thread titles (#467)

* FIX: Better AI chat thread titles

- Fix quote removal when multi-line

- Use XML tags for better LLM output parsing

- Use stop_sequences for faster and less wasteful LLM calls

- Adds truncation as the last line of defense
---
 lib/ai_helper/chat_thread_titler.rb             |  9 +++++++--
 .../ai_helper/chat_thread_titler_spec.rb        | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/lib/ai_helper/chat_thread_titler.rb b/lib/ai_helper/chat_thread_titler.rb
index 755a8426..36b4517f 100644
--- a/lib/ai_helper/chat_thread_titler.rb
+++ b/lib/ai_helper/chat_thread_titler.rb
@@ -23,7 +23,7 @@ module DiscourseAi
             and you will generate a single attention-grabbing title. Please keep the title concise and under 15 words
             and ensure that the meaning is maintained. The title will utilize the same language type of the chat.
             I want you to only reply the suggested title and nothing else, do not write explanations.
-            You will find the chat between <input></input> XML tags.
+            You will find the chat between <input></input> XML tags. Return the suggested title between <title> tags.
           TEXT
             messages: [{ type: :user, content: chat, id: "User" }],
           )
@@ -31,11 +31,16 @@ module DiscourseAi
         DiscourseAi::Completions::Llm.proxy(SiteSetting.ai_helper_model).generate(
           prompt,
           user: Discourse.system_user,
+          stop_sequences: ["</input>"],
         )
       end
 
       def cleanup(title)
-        title.split("\n").first.then { _1.match?(/^("|')(.*)("|')$/) ? title[1..-2] : _1 }
+        (Nokogiri::HTML5.fragment(title).at("title")&.text || title)
+          .split("\n")
+          .first
+          .then { _1.match?(/^("|')(.*)("|')$/) ? _1[1..-2] : _1 }
+          .truncate(100, separator: " ")
       end
 
       def thread_content(thread)
diff --git a/spec/lib/modules/ai_helper/chat_thread_titler_spec.rb b/spec/lib/modules/ai_helper/chat_thread_titler_spec.rb
index 89ad6674..3ff9a4a3 100644
--- a/spec/lib/modules/ai_helper/chat_thread_titler_spec.rb
+++ b/spec/lib/modules/ai_helper/chat_thread_titler_spec.rb
@@ -45,6 +45,23 @@ RSpec.describe DiscourseAi::AiHelper::ChatThreadTitler do
 
       expect(result).to eq(expected_title)
     end
+
+    it "parses the XML" do
+      titles = "Here is your title <title>The solitary horse</title> my friend"
+      expected_title = "The solitary horse"
+
+      result = titler.cleanup(titles)
+
+      expect(result).to eq(expected_title)
+    end
+
+    it "truncates long titles" do
+      titles = "O cavalo trota pelo campo" + " Pocotó" * 100
+
+      result = titler.cleanup(titles)
+
+      expect(result.size).to be <= 100
+    end
   end
 
   describe "#thread_content" do