diff --git a/config/feature_flags/development/ai_custom_models_prompts_migration.yml b/config/feature_flags/development/ai_custom_models_prompts_migration.yml
new file mode 100644
index 0000000000000000000000000000000000000000..01ba89c7199858008e22c405fed6be323f9ca5fc
--- /dev/null
+++ b/config/feature_flags/development/ai_custom_models_prompts_migration.yml
@@ -0,0 +1,9 @@
+---
+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
diff --git a/ee/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message.rb b/ee/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message.rb
index b8a1aad944f9a33c1eb5b486e11635ea1e78a528..e8ab9c04528828213b00dfde775507f50a9f2a25 100644
--- a/ee/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message.rb
+++ b/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
diff --git a/ee/lib/code_suggestions/tasks/self_hosted_code_completion.rb b/ee/lib/code_suggestions/tasks/self_hosted_code_completion.rb
index 756d6bd9e785b9bde197ed84ddb4f2d073fcf773..c26bd9efa2a18bf86f05cd9d9f2affe9899832eb 100644
--- a/ee/lib/code_suggestions/tasks/self_hosted_code_completion.rb
+++ b/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
diff --git a/ee/spec/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message_spec.rb b/ee/spec/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message_spec.rb
index ea1bde358710f56443160f6ca353266af389bd2b..8eb38cd10d7e775e4fe52d906749684951bb8073 100644
--- a/ee/spec/lib/code_suggestions/prompts/code_completion/ai_gateway_code_completion_message_spec.rb
+++ b/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
diff --git a/ee/spec/lib/code_suggestions/tasks/self_hosted_code_completion_spec.rb b/ee/spec/lib/code_suggestions/tasks/self_hosted_code_completion_spec.rb
index a4f4cf219211183719f40a37a52bc4f67661bcfb..bfb7251257eba283b3254fa21d10bfbb4ed4867f 100644
--- a/ee/spec/lib/code_suggestions/tasks/self_hosted_code_completion_spec.rb
+++ b/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