Define prompts within translator instead of relying on prompts from discourse-ai

Author: nattsw

app/services/discourse_ai/language_detector.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  class LanguageDetector
+    PROMPT_TEXT = <<~TEXT
+      I want you to act as a language expert, determining the locale for a set of text.
+      The locale is a language identifier, such as "en" for English, "de" for German, etc,
+      and can also include a region identifier, such as "en-GB" for British English, or "zh-Hans" for Simplified Chinese.
+      I will provide you with text, and you will determine the locale of the text.
+      Include your locale between <language></language> XML tags.
+    TEXT
+
+    def initialize(text)
+      @text = text
+    end
+
+    def detect
+      prompt =
+        DiscourseAi::Completions::Prompt.new(
+          PROMPT_TEXT,
+          messages: [{ type: :user, content: @text, id: "user" }],
+        )
+
+      response =
+        DiscourseAi::Completions::Llm.proxy(SiteSetting.ai_helper_model).generate(
+          prompt,
+          user: Discourse.system_user,
+          feature_name: "translator-language-detect",
+        )
+
+      (Nokogiri::HTML5.fragment(response).at("language")&.text || response)
+    end
+  end
+end

app/services/discourse_ai/translator.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  class Translator
+    def initialize(text, target_language)
+      @text = text
+      @target_language = target_language
+    end
+
+    def translate
+      prompt =
+        DiscourseAi::Completions::Prompt.new(
+          build_prompt(@target_language),
+          messages: [{ type: :user, content: @text, id: "user" }],
+        )
+
+      llm_translation =
+        DiscourseAi::Completions::Llm.proxy(SiteSetting.ai_helper_model).generate(
+          prompt,
+          user: Discourse.system_user,
+          feature_name: "translator-translate",
+        )
+
+      (Nokogiri::HTML5.fragment(llm_translation).at("translation")&.text || llm_translation)
+    end
+
+    private
+
+    def build_prompt(target_language)
+      <<~TEXT
+          You are a highly skilled translator with expertise in many languages.
+          Your task is to identify the language of the text I provide and accurately translate it into this language locale "#{target_language}" while preserving the meaning, tone, and nuance of the original text.
+          The text may also contain html tags, which should be preserved in the translation.
+          Please maintain proper grammar, spelling, and punctuation in the translated version.
+          Wrap the translated text in a <translation> tag.
+        TEXT
+    end
+  end
+end

app/services/discourse_translator/discourse_ai.rb

@@ -14,14 +14,11 @@ def self.detect(topic_or_post)
       return unless required_settings_enabled
 
       topic_or_post.custom_fields[DiscourseTranslator::DETECTED_LANG_CUSTOM_FIELD] ||= begin
-        # we don't need that much text to determine the locale
-        text = get_text(topic_or_post).truncate(MAX_DETECT_LOCALE_TEXT_LENGTH)
-
-        get_ai_helper_output(
-          text,
-          CompletionPrompt.find_by(id: CompletionPrompt::DETECT_TEXT_LOCALE),
-        )
+        ::DiscourseAi::LanguageDetector.new(text_for_detection(topic_or_post)).detect
       end
+    rescue => e
+      e.message
+      Rails.logger.warn("Failed to detect language: #{e}")
     end
 
     def self.translate(topic_or_post)
@@ -30,27 +27,14 @@ def self.translate(topic_or_post)
       detected_lang = detect(topic_or_post)
       translated_text =
         from_custom_fields(topic_or_post) do
-          get_ai_helper_output(
-            get_text(topic_or_post),
-            CompletionPrompt.find_by(id: CompletionPrompt::TRANSLATE),
-          )
+          ::DiscourseAi::Translator.new(text_for_translation(topic_or_post), I18n.locale).translate
         end
 
       [detected_lang, translated_text]
     end
 
     private
 
-    def self.get_ai_helper_output(text, prompt)
-      ::DiscourseAi::AiHelper::Assistant.new.generate_and_send_prompt(
-        prompt,
-        text,
-        Discourse.system_user,
-      )[
-        :suggestions
-      ].first
-    end
-
     def self.required_settings_enabled
       SiteSetting.translator_enabled && SiteSetting.translator == "DiscourseAi" &&
         SiteSetting.discourse_ai_enabled && SiteSetting.ai_helper_enabled

spec/services/discourse_ai/language_detector_spec.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "rails_helper"
+
+describe DiscourseAi::LanguageDetector do
+  before do
+    Fabricate(:fake_model).tap do |fake_llm|
+      SiteSetting.public_send("ai_helper_model=", "custom:#{fake_llm.id}")
+    end
+    SiteSetting.ai_helper_enabled = true
+  end
+
+  describe ".detect" do
+    it "creates the correct prompt" do
+      allow(DiscourseAi::Completions::Prompt).to receive(:new).with(
+        DiscourseAi::LanguageDetector::PROMPT_TEXT,
+        messages: [{ type: :user, content: "meow", id: "user" }],
+      ).and_call_original
+
+      described_class.new("meow").detect
+    end
+
+    it "sends the language detection prompt to the ai helper model" do
+      mock_prompt = instance_double(DiscourseAi::Completions::Prompt)
+      mock_llm = instance_double(DiscourseAi::Completions::Llm)
+
+      allow(DiscourseAi::Completions::Prompt).to receive(:new).and_return(mock_prompt)
+      allow(DiscourseAi::Completions::Llm).to receive(:proxy).with(
+        SiteSetting.ai_helper_model,
+      ).and_return(mock_llm)
+      allow(mock_llm).to receive(:generate).with(
+        mock_prompt,
+        user: Discourse.system_user,
+        feature_name: "translator-language-detect",
+      )
+
+      described_class.new("meow").detect
+    end
+
+    it "returns the language from the llm's response in the language tag" do
+      DiscourseAi::Completions::Llm.with_prepared_responses(["<language>de</language>"]) do
+        expect(described_class.new("meow").detect).to eq "de"
+      end
+    end
+  end
+end

spec/services/discourse_ai/translator_spec.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require "rails_helper"
+
+describe DiscourseAi::Translator do
+  before do
+    Fabricate(:fake_model).tap do |fake_llm|
+      SiteSetting.public_send("ai_helper_model=", "custom:#{fake_llm.id}")
+    end
+    SiteSetting.ai_helper_enabled = true
+  end
+
+  describe ".translate" do
+    let(:text_to_translate) { "cats are great" }
+    let(:target_language) { "de" }
+
+    it "creates the correct prompt" do
+      allow(DiscourseAi::Completions::Prompt).to receive(:new).with(
+        <<~TEXT,
+          You are a highly skilled translator with expertise in many languages.
+          Your task is to identify the language of the text I provide and accurately translate it into this language locale "de" while preserving the meaning, tone, and nuance of the original text.
+          The text may also contain html tags, which should be preserved in the translation.
+          Please maintain proper grammar, spelling, and punctuation in the translated version.
+          Wrap the translated text in a <translation> tag.
+        TEXT
+        messages: [{ type: :user, content: text_to_translate, id: "user" }],
+      ).and_call_original
+
+      described_class.new(text_to_translate, target_language).translate
+    end
+
+    it "sends the translation prompt to the selected ai helper model" do
+      mock_prompt = instance_double(DiscourseAi::Completions::Prompt)
+      mock_llm = instance_double(DiscourseAi::Completions::Llm)
+
+      allow(DiscourseAi::Completions::Prompt).to receive(:new).and_return(mock_prompt)
+      allow(DiscourseAi::Completions::Llm).to receive(:proxy).with(
+        SiteSetting.ai_helper_model,
+      ).and_return(mock_llm)
+      allow(mock_llm).to receive(:generate).with(
+        mock_prompt,
+        user: Discourse.system_user,
+        feature_name: "translator-translate",
+      )
+
+      described_class.new(text_to_translate, target_language).translate
+    end
+
+    it "returns the translation from the llm's response in the translation tag" do
+      DiscourseAi::Completions::Llm.with_prepared_responses(
+        ["<translation>hur dur hur dur!</translation>"],
+      ) do
+        expect(
+          described_class.new(text_to_translate, target_language).translate,
+        ).to eq "hur dur hur dur!"
+      end
+    end
+
+    it "returns the raw response if the translation tag is not present" do
+      DiscourseAi::Completions::Llm.with_prepared_responses(["raw response."]) do
+        expect(
+          described_class.new(text_to_translate, target_language).translate,
+        ).to eq "raw response."
+      end
+    end
+  end
+end

spec/services/discourse_ai_spec.rb

@@ -14,32 +14,22 @@
     SiteSetting.translator = "DiscourseAi"
   end
 
+  describe ".language_supported?" do
+    it "returns true for any language" do
+      expect(described_class.language_supported?("any-language")).to eq(true)
+    end
+  end
+
   describe ".detect" do
     it "stores the detected language in a custom field" do
       locale = "de"
-      DiscourseAi::Completions::Llm.with_prepared_responses(["<output>de</output>"]) do
+      DiscourseAi::Completions::Llm.with_prepared_responses(["<language>de</language>"]) do
         DiscourseTranslator::DiscourseAi.detect(post)
         post.save_custom_fields
       end
 
       expect(post.custom_fields[DiscourseTranslator::DETECTED_LANG_CUSTOM_FIELD]).to eq locale
     end
-
-    it "truncates to MAX LENGTH" do
-      truncated_text =
-        post.cooked.truncate(DiscourseTranslator::DiscourseAi::MAX_DETECT_LOCALE_TEXT_LENGTH)
-      expect_any_instance_of(::DiscourseAi::AiHelper::Assistant).to receive(
-        :generate_and_send_prompt,
-      ).with(
-        CompletionPrompt.find_by(id: CompletionPrompt::DETECT_TEXT_LOCALE),
-        truncated_text,
-        Discourse.system_user,
-      ).and_call_original
-
-      DiscourseAi::Completions::Llm.with_prepared_responses(["<output>de</output>"]) do
-        DiscourseTranslator::DiscourseAi.detect(post)
-      end
-    end
   end
 
   describe ".translate" do
@@ -50,7 +40,7 @@
 
     it "translates the post and returns [locale, translated_text]" do
       DiscourseAi::Completions::Llm.with_prepared_responses(
-        ["<output>some translated text</output>", "<output>translated</output>"],
+        ["<translation>some translated text</translation>"],
       ) do
         locale, translated_text = DiscourseTranslator::DiscourseAi.translate(post)
         expect(locale).to eq "de"