FEATURE: StableBeluga2 support for AiHelper (#162)

* FEATURE: StableBeluga2 support for AiHelper

* lint

app/controllers/discourse_ai/ai_helper/assistant_controller.rb

@@ -33,6 +33,7 @@ module DiscourseAi
                  status: 200
         end
       rescue ::DiscourseAi::Inference::OpenAiCompletions::CompletionFailed,
+             ::DiscourseAi::Inference::HuggingFaceTextGeneration::CompletionFailed,
              ::DiscourseAi::Inference::AnthropicCompletions::CompletionFailed => e
         render_json_error I18n.t("discourse_ai.ai_helper.errors.completion_request_failed"),
                           status: 502

app/models/completion_prompt.rb

@@ -10,10 +10,13 @@ class CompletionPrompt < ActiveRecord::Base
   validate :each_message_length
 
   def messages_with_user_input(user_input)
-    if ::DiscourseAi::AiHelper::LlmPrompt.new.enabled_provider == "openai"
+    case ::DiscourseAi::AiHelper::LlmPrompt.new.enabled_provider
+    when "openai"
       self.messages << { role: "user", content: user_input }
-    else
+    when "anthropic"
       self.messages << { "role" => "Input", "content" => "<input>#{user_input}</input>" }
+    when "huggingface"
+      self.messages.first.sub("{{user_input}}", user_input)
     end
   end
 

config/settings.yml

@@ -147,6 +147,7 @@ discourse_ai:
       - gpt-3.5-turbo
       - gpt-4
       - claude-2
+      - stable-beluga-2
 
   ai_embeddings_enabled:
     default: false

db/fixtures/ai_helper/602_stablebeluga2_completion_prompts.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+CompletionPrompt.seed do |cp|
+  cp.id = -1
+  cp.provider = "huggingface"
+  cp.name = "translate"
+  cp.prompt_type = CompletionPrompt.prompt_types[:text]
+  cp.messages = [<<~TEXT]
+    ### System:
+    I want you to act as an English translator, spelling corrector and improver. I will speak to you
+    in any language and you will detect the language, translate it and answer in the corrected and 
+    improved version of my text, in English. I want you to replace my simplified A0-level words and 
+    sentences with more beautiful and elegant, upper level English words and sentences. 
+    Keep the meaning same, but make them more literary. I want you to only reply the correction, 
+    the improvements and nothing else, do not write explanations.
+    
+    ### User:
+    {{user_input}}
+
+    ### Assistant:
+    Here is the corrected, translated and improved version of the text:
+  TEXT
+end
+
+CompletionPrompt.seed do |cp|
+  cp.id = -2
+  cp.provider = "huggingface"
+  cp.name = "generate_titles"
+  cp.prompt_type = CompletionPrompt.prompt_types[:list]
+  cp.messages = [<<~TEXT]
+    ### System:
+    I want you to act as a title generator for written pieces. I will provide you with a text, 
+    and you will generate five attention-grabbing titles. Please keep the title concise and under 20 words,
+    and ensure that the meaning is maintained. Replies will utilize the language type of the topic.
+    I want you to only reply the list of options and nothing else, do not write explanations.
+    
+    ### User:
+    {{user_input}}
+
+    ### Assistant:
+    Here are five titles for the text:
+  TEXT
+end
+
+CompletionPrompt.seed do |cp|
+  cp.id = -3
+  cp.provider = "huggingface"
+  cp.name = "proofread"
+  cp.prompt_type = CompletionPrompt.prompt_types[:diff]
+  cp.messages = [<<~TEXT]
+    ### System:
+    You are a markdown proofreader. You correct egregious typos and phrasing issues but keep the user's original voice.
+    You do not touch code blocks. I will provide you with text to proofread. If nothing needs fixing, then you will echo the text back.
+
+    Optionally, a user can specify intensity. Intensity 10 is a pedantic English teacher correcting the text.
+    Intensity 1 is a minimal proofreader. By default, you operate at intensity 1.
+    
+    ### User:
+    Rewrite the following text to correct any errors:
+    {{user_input}}
+    
+    ### Assistant:
+    Here is a proofread version of the text:
+  TEXT
+end
+
+CompletionPrompt.seed do |cp|
+  cp.id = -4
+  cp.provider = "huggingface"
+  cp.name = "markdown_table"
+  cp.prompt_type = CompletionPrompt.prompt_types[:diff]
+  cp.messages = [<<~TEXT]
+    ### System:
+    You are a markdown table formatter, I will provide you text and you will format it into a markdown table
+    
+    ### User:
+    sam,joe,jane
+    age: 22|  10|11
+
+    ### Assistant:
+    |   | sam | joe | jane |
+    |---|---|---|---|
+    | age | 22 | 10 | 11 |
+    
+    ### User:
+    sam: speed 100, age 22
+    jane: age 10
+    fred: height 22
+
+    ### Assistant:
+    |   | speed | age | height |
+    |---|---|---|---|
+    | sam | 100 | 22 | - |
+    | jane | - | 10 | - |
+    | fred | - | - | 22 |
+    
+    ### User:   
+    chrome 22ms (first load 10ms)
+    firefox 10ms (first load: 9ms)
+    
+    ### Assistant:
+    | Browser | Load Time (ms) | First Load Time (ms) |
+    |---|---|---|
+    | Chrome | 22 | 10 |
+    | Firefox | 10 | 9 |
+    
+    ### User:
+    {{user_input}}
+
+    ### Assistant:
+  TEXT
+end

lib/modules/ai_helper/llm_prompt.rb

@@ -22,18 +22,24 @@ module DiscourseAi
       end
 
       def generate_and_send_prompt(prompt, text)
-        if enabled_provider == "openai"
+        case enabled_provider
+        when "openai"
           openai_call(prompt, text)
-        else
+        when "anthropic"
           anthropic_call(prompt, text)
+        when "huggingface"
+          huggingface_call(prompt, text)
         end
       end
 
       def enabled_provider
-        if SiteSetting.ai_helper_model.start_with?("gpt")
+        case SiteSetting.ai_helper_model
+        when /gpt/
           "openai"
-        else
+        when /claude/
           "anthropic"
+        else
+          "huggingface"
         end
       end
 
@@ -49,12 +55,17 @@ module DiscourseAi
       def parse_content(prompt, content)
         return "" if content.blank?
 
-        if enabled_provider == "openai"
+        case enabled_provider
+        when "openai"
           return content.strip if !prompt.list?
 
           content.gsub("\"", "").gsub(/\d./, "").split("\n").map(&:strip)
-        else
+        when "anthropic"
           parse_antropic_content(prompt, content)
+        when "huggingface"
+          return [content.strip.delete_prefix('"').delete_suffix('"')] if !prompt.list?
+
+          content.gsub("\"", "").gsub(/\d./, "").split("\n").map(&:strip)
         end
       end
 
@@ -93,6 +104,24 @@ module DiscourseAi
         result
       end
 
+      def huggingface_call(prompt, text)
+        result = { type: prompt.prompt_type }
+
+        message = prompt.messages_with_user_input(text)
+
+        response =
+          DiscourseAi::Inference::HuggingFaceTextGeneration.perform!(
+            message,
+            SiteSetting.ai_helper_model,
+          )
+
+        result[:suggestions] = parse_content(prompt, response.dig(:generated_text))
+
+        result[:diff] = generate_diff(text, result[:suggestions].first) if prompt.diff?
+
+        result
+      end
+
       def parse_antropic_content(prompt, content)
         if prompt.list?
           suggestions = Nokogiri::HTML5.fragment(content).search("ai").map(&:text)