From 442681a3d3e93bdf9f58ff4dcd08e2f3373ab9f7 Mon Sep 17 00:00:00 2001
From: Roman Rizzi <roman@discourse.org>
Date: Thu, 4 Jul 2024 13:23:03 -0300
Subject: [PATCH] FIX: Mixtral models have system role support. (#703)

Using assistant role for system produces an error because
they expect alternating roles like user/assistant/user and so on.
Prompts cannot start with the assistant role.
---
 lib/completions/dialects/dialect.rb           |  1 -
 lib/completions/dialects/mistral.rb           | 59 -------------------
 .../endpoints/hugging_face_spec.rb            |  7 ++-
 spec/lib/completions/endpoints/vllm_spec.rb   | 11 +++-
 spec/lib/completions/llm_spec.rb              |  2 +-
 5 files changed, 16 insertions(+), 64 deletions(-)
 delete mode 100644 lib/completions/dialects/mistral.rb

diff --git a/lib/completions/dialects/dialect.rb b/lib/completions/dialects/dialect.rb
index 0f64c2a4..d36c8f1d 100644
--- a/lib/completions/dialects/dialect.rb
+++ b/lib/completions/dialects/dialect.rb
@@ -13,7 +13,6 @@ module DiscourseAi
             [
               DiscourseAi::Completions::Dialects::ChatGpt,
               DiscourseAi::Completions::Dialects::Gemini,
-              DiscourseAi::Completions::Dialects::Mistral,
               DiscourseAi::Completions::Dialects::Claude,
               DiscourseAi::Completions::Dialects::Command,
               DiscourseAi::Completions::Dialects::OpenAiCompatible,
diff --git a/lib/completions/dialects/mistral.rb b/lib/completions/dialects/mistral.rb
deleted file mode 100644
index b5ba1aeb..00000000
--- a/lib/completions/dialects/mistral.rb
+++ /dev/null
@@ -1,59 +0,0 @@
-# frozen_string_literal: true
-
-module DiscourseAi
-  module Completions
-    module Dialects
-      class Mistral < Dialect
-        class << self
-          def can_translate?(model_name)
-            %w[
-              mistralai/Mixtral-8x7B-Instruct-v0.1
-              mistralai/Mistral-7B-Instruct-v0.2
-              mistral
-            ].include?(model_name)
-          end
-        end
-
-        def tokenizer
-          llm_model&.tokenizer_class || DiscourseAi::Tokenizer::MixtralTokenizer
-        end
-
-        def tools
-          @tools ||= tools_dialect.translated_tools
-        end
-
-        def max_prompt_tokens
-          return llm_model.max_prompt_tokens if llm_model&.max_prompt_tokens
-
-          32_000
-        end
-
-        private
-
-        def system_msg(msg)
-          { role: "assistant", content: "<s>#{msg[:content]}</s>" }
-        end
-
-        def model_msg(msg)
-          { role: "assistant", content: msg[:content] }
-        end
-
-        def tool_call_msg(msg)
-          tools_dialect.from_raw_tool_call(msg)
-        end
-
-        def tool_msg(msg)
-          tools_dialect.from_raw_tool(msg)
-        end
-
-        def user_msg(msg)
-          content = +""
-          content << "#{msg[:id]}: " if msg[:id]
-          content << msg[:content]
-
-          { role: "user", content: content }
-        end
-      end
-    end
-  end
-end
diff --git a/spec/lib/completions/endpoints/hugging_face_spec.rb b/spec/lib/completions/endpoints/hugging_face_spec.rb
index f84fc43a..d5ed7446 100644
--- a/spec/lib/completions/endpoints/hugging_face_spec.rb
+++ b/spec/lib/completions/endpoints/hugging_face_spec.rb
@@ -94,7 +94,12 @@ RSpec.describe DiscourseAi::Completions::Endpoints::HuggingFace do
   let(:hf_mock) { HuggingFaceMock.new(endpoint) }
 
   let(:compliance) do
-    EndpointsCompliance.new(self, endpoint, DiscourseAi::Completions::Dialects::Mistral, user)
+    EndpointsCompliance.new(
+      self,
+      endpoint,
+      DiscourseAi::Completions::Dialects::OpenAiCompatible,
+      user,
+    )
   end
 
   describe "#perform_completion!" do
diff --git a/spec/lib/completions/endpoints/vllm_spec.rb b/spec/lib/completions/endpoints/vllm_spec.rb
index d879ad09..a361f62a 100644
--- a/spec/lib/completions/endpoints/vllm_spec.rb
+++ b/spec/lib/completions/endpoints/vllm_spec.rb
@@ -69,10 +69,17 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Vllm do
   let(:anthropic_mock) { VllmMock.new(endpoint) }
 
   let(:compliance) do
-    EndpointsCompliance.new(self, endpoint, DiscourseAi::Completions::Dialects::Mistral, user)
+    EndpointsCompliance.new(
+      self,
+      endpoint,
+      DiscourseAi::Completions::Dialects::OpenAiCompatible,
+      user,
+    )
   end
 
-  let(:dialect) { DiscourseAi::Completions::Dialects::Mistral.new(generic_prompt, model_name) }
+  let(:dialect) do
+    DiscourseAi::Completions::Dialects::OpenAiCompatible.new(generic_prompt, model_name)
+  end
   let(:prompt) { dialect.translate }
 
   let(:request_body) { model.default_options.merge(messages: prompt).to_json }
diff --git a/spec/lib/completions/llm_spec.rb b/spec/lib/completions/llm_spec.rb
index e91a4a60..9e5b935f 100644
--- a/spec/lib/completions/llm_spec.rb
+++ b/spec/lib/completions/llm_spec.rb
@@ -3,7 +3,7 @@
 RSpec.describe DiscourseAi::Completions::Llm do
   subject(:llm) do
     described_class.new(
-      DiscourseAi::Completions::Dialects::Mistral,
+      DiscourseAi::Completions::Dialects::OpenAiCompatible,
       canned_response,
       "hugging_face:Upstage-Llama-2-*-instruct-v2",
       gateway: canned_response,