Add ai_custom_models_prompts_migration FF

Add a feature flag to control prompts migration
Use the new FF to route to AiGatewayCodeCompletionMessage
Adjust the base prompt in AiGatewayCodeCompletionMessage to be nil

Changelog: changed
EE: true

config/feature_flags/development/ai_custom_models_prompts_migration.yml

+---
+name: ai_custom_models_prompts_migration
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/473156
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/160050
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/473358
+milestone: '17.3'
+group: group::custom models
+type: development
+default_enabled: false

ee/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message.rb

@@ -18,12 +18,12 @@ def request_params
           end
         end
 
-        private
-
         def prompt
-          raise NotImplementedError, "#{self.class} has not implemented method #{__method__}"
+          nil
         end
 
+        private
+
         def pick_prefix
           prefix.last(500)
         end

ee/lib/code_suggestions/tasks/self_hosted_code_completion.rb

@@ -37,16 +37,23 @@ def params
 
       def prompt
         model_name = feature_setting.self_hosted_model.model.to_sym
-        case model_name
-        when :codegemma
-          CodeSuggestions::Prompts::CodeCompletion::CodeGemmaMessages.new(params)
-        when :codestral
-          CodeSuggestions::Prompts::CodeCompletion::CodestralMessages.new(params)
-        when :'codellama:code'
-          CodeSuggestions::Prompts::CodeCompletion::CodellamaMessages.new(params)
-        else
-          raise "Unknown model: #{model_name}"
-        end
+        # rubocop:disable Gitlab/FeatureFlagWithoutActor -- Global development flag for migrating the prompts
+        prompt_migration_enabled = ::Feature.enabled?(:ai_custom_models_prompts_migration)
+        # rubocop:enable Gitlab/FeatureFlagWithoutActor
+        ai_gateway_class = CodeSuggestions::Prompts::CodeCompletion::AiGatewayCodeCompletionMessage
+        model_classes = {
+          codegemma: CodeSuggestions::Prompts::CodeCompletion::CodeGemmaMessages,
+          codestral: CodeSuggestions::Prompts::CodeCompletion::CodestralMessages,
+          'codellama:code': CodeSuggestions::Prompts::CodeCompletion::CodellamaMessages
+        }
+
+        message_class = if prompt_migration_enabled
+                          ai_gateway_class
+                        else
+                          model_classes.fetch(model_name) { raise "Unknown model: #{model_name}" }
+                        end
+
+        message_class.new(params)
       end
       strong_memoize_attr :prompt
     end

ee/spec/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message_spec.rb

@@ -35,10 +35,8 @@ def prompt
   end
 
   describe '#prompt' do
-    it 'raises NotImplementedError for the abstract class' do
-      expect do
-        described_class.new(params).send(:prompt)
-      end.to raise_error(NotImplementedError, "#{described_class} has not implemented method prompt")
+    it 'returns an empty prompt' do
+      expect(described_class.new(params).prompt).to be_nil
     end
   end
 end

ee/spec/lib/code_suggestions/tasks/self_hosted_code_completion_spec.rb

@@ -52,6 +52,10 @@
       params: params, unsafe_passthrough_params: unsafe_params)
   end
 
+  before do
+    stub_feature_flags(ai_custom_models_prompts_migration: false)
+  end
+
   describe '#body' do
     before do
       allow(CodeSuggestions::Prompts::CodeCompletion::CodeGemmaMessages)
@@ -85,6 +89,21 @@
 
       expect(CodeSuggestions::Prompts::CodeCompletion::CodeGemmaMessages).to have_received(:new).with(params)
     end
+
+    context 'when the ai_custom_models_prompts_migration FF is enabled' do
+      before do
+        stub_feature_flags(ai_custom_models_prompts_migration: true)
+        allow(CodeSuggestions::Prompts::CodeCompletion::AiGatewayCodeCompletionMessage)
+          .to receive(:new).and_return(codgemma_messages_prompt)
+      end
+
+      it 'calls the base AiGatewayCodeCompletionMessage class' do
+        task.body
+
+        expect(CodeSuggestions::Prompts::CodeCompletion::AiGatewayCodeCompletionMessage)
+          .to have_received(:new).with(params)
+      end
+    end
   end
 
   describe 'prompt selection per model name' do