Merge branch 'jp-cs-prompt2' into 'master'

Extend code generation prompt

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/131161



Merged-by: Mikołaj Wawrzyniak <mwawrzyniak@gitlab.com>
Approved-by: Mikołaj Wawrzyniak <mwawrzyniak@gitlab.com>
Reviewed-by: Mikołaj Wawrzyniak <mwawrzyniak@gitlab.com>
Reviewed-by: Rajendra Kadam <rkadam@gitlab.com>
Co-authored-by: Jan Provaznik <jprovaznik@gitlab.com>
Co-authored-by: Rajendra Kadam <rkadam@gitlab.com>

ee/lib/code_suggestions/task_selector.rb

@@ -8,14 +8,34 @@ class TaskSelector
     # It is case-insensitive.
     # It searches for the last instance of a match by looking for the end
     # of a text block and an optional line break.
-    GENERATE_COMMENT_PREFIX = %r{(--|#|//)\s?GitLab Duo Generate:(.{10,})\s*\z}i
+    GENERATE_COMMENT_PREFIX =
+      %r{(?<prefix>.*)(?<comment>--|#|//)[ \t]?GitLab Duo Generate:[ \t]*(?<instruction>[^\r\n]{10,})\s*\Z}im
+    PREFIX_MAX_SIZE = 100_000
 
     def self.task(params:)
       prefix = params.dig('current_file', 'content_above_cursor')
+      match = comment_match(prefix)
 
-      return CodeSuggestions::Tasks::CodeGeneration::FromComment.new(params) if prefix&.match?(GENERATE_COMMENT_PREFIX)
+      return CodeSuggestions::Tasks::CodeCompletion.new(params.merge(prefix: prefix)) unless match
 
-      CodeSuggestions::Tasks::CodeCompletion.new(params)
+      CodeSuggestions::Tasks::CodeGeneration::FromComment.new(
+        params.merge(
+          prefix: match[:prefix].chomp,
+          instruction: match[:instruction]
+        )
+      )
     end
+
+    def self.comment_match(prefix)
+      return unless prefix
+      # This is a short-term fix for avoiding processing too long strings,
+      # but we should rather define proper set of input parameters for REST API
+      # endpoint with its limits:
+      # https://gitlab.com/gitlab-org/gitlab/-/issues/424724
+      return if prefix.size > PREFIX_MAX_SIZE
+
+      prefix&.match(GENERATE_COMMENT_PREFIX)
+    end
+    private_class_method :comment_match
   end
 end

ee/lib/code_suggestions/tasks/base.rb

@@ -6,6 +6,7 @@ class Base
       DEFAULT_CODE_SUGGESTIONS_URL = 'https://codesuggestions.gitlab.com'
 
       def initialize(params)
+        @prefix = params.delete(:prefix)
         @params = params
       end
 
@@ -21,11 +22,7 @@ def body
 
       private
 
-      attr_reader :params
-
-      def prefix
-        params.dig('current_file', 'content_above_cursor')
-      end
+      attr_reader :params, :prefix
 
       def endpoint_name
         raise NotImplementedError

ee/lib/code_suggestions/tasks/code_generation/from_comment.rb

@@ -8,6 +8,46 @@ class FromComment < CodeSuggestions::Tasks::Base
 
         GATEWAY_PROMPT_VERSION = 2
 
+        # https://cloud.google.com/vertex-ai/docs/generative-ai/code/code-models-overview
+        SUPPORTED_LANGUAGES = {
+          "C" => %w[c],
+          "C++" => %w[cc cpp],
+          "C#" => %w[cs],
+          "Clojure" => %w[clj cljs cljc],
+          "Dart" => %w[dart],
+          "Elixir" => %w[ex],
+          "Erlang" => %w[erl],
+          "Fortran" => %w[f],
+          "Go" => %w[go],
+          "GoogleSQL" => %w[sql],
+          "Groovy" => %w[groovy],
+          "Haskell" => %w[hs],
+          "HTML" => %w[html],
+          "Java" => %w[java],
+          "JavaScript" => %w[js],
+          "Kotlin" => %w[kt kts],
+          "Lean (proof assistant)" => %w[lean],
+          "Objective-C" => %w[m],
+          "OCaml" => %w[ml],
+          "Perl" => %w[pl],
+          "PHP" => %w[php],
+          "Python" => %w[py],
+          "Ruby" => %w[rb],
+          "Rust" => %w[rs],
+          "Scala" => %w[scala],
+          "Shell script" => %w[sh],
+          "Solidity" => %w[sol],
+          "Swift" => %w[swift],
+          "TypeScript" => %w[ts],
+          "Verilog" => %w[v]
+        }.freeze
+
+        def initialize(params)
+          @instruction = params.delete(:instruction)
+
+          super
+        end
+
         override :endpoint_name
         def endpoint_name
           'generations'
@@ -23,17 +63,31 @@ def body
 
         private
 
+        attr_reader :instruction
+
         def file_name
           params.dig('current_file', 'file_name').to_s
         end
 
         def prompt
           extension = File.extname(file_name).delete_prefix('.')
+          language = SUPPORTED_LANGUAGES.find { |_, extensions| extensions.include?(extension) }&.first
+          file_path_info = File.path(file_name)
 
           <<~PROMPT
+            This is a task to write new #{language} code in a file '#{file_path_info}' based on a given description.
+            You get first the already existing code file and then the description of the code that needs to be created.
+            It is your task to write valid and working #{language} code.
+            Only return in your response new code.
+
+            Already existing code:
+
             ```#{extension}
             #{prefix}
             ```
+
+            Create new code for the following description:
+            `#{instruction}`
           PROMPT
         end
       end

ee/spec/lib/code_suggestions/task_selector_spec.rb

@@ -67,6 +67,16 @@ def fibonacci(x)
       it 'only takes the last example in to account' do
         expect(subject).to be_an_instance_of(CodeSuggestions::Tasks::CodeGeneration::FromComment)
       end
+
+      context 'when prefix is too long' do
+        before do
+          stub_const('CodeSuggestions::TaskSelector::PREFIX_MAX_SIZE', 10)
+        end
+
+        it 'does not parse prefix and uses completion' do
+          expect(subject).to be_an_instance_of(CodeSuggestions::Tasks::CodeCompletion)
+        end
+      end
     end
 
     context 'when the last comment is a code suggestion' do

ee/spec/lib/code_suggestions/tasks/code_generation/from_comment_spec.rb

@@ -4,11 +4,71 @@
 
 RSpec.describe CodeSuggestions::Tasks::CodeGeneration::FromComment, feature_category: :code_suggestions do
   let(:prefix) { 'some text' }
-  let(:params) { { 'current_file' => { 'file_name' => 'test.py', 'content_above_cursor' => prefix } } }
+  let(:instruction) { 'Add code for validating function' }
+  let(:file_name) { 'test.py' }
 
-  it_behaves_like 'code suggestion task' do
+  let(:params) do
+    {
+      :prefix => prefix,
+      :instruction => instruction,
+      'current_file' => {
+        'file_name' => file_name,
+        'content_above_cursor' => 'some text'
+      }
+    }
+  end
+
+  let(:prompt) do
+    <<~PROMPT
+      This is a task to write new Python code in a file 'test.py' based on a given description.
+      You get first the already existing code file and then the description of the code that needs to be created.
+      It is your task to write valid and working Python code.
+      Only return in your response new code.
+
+      Already existing code:
+
+      ```py
+      some text
+      ```
+
+      Create new code for the following description:
+      `#{instruction}`
+    PROMPT
+  end
+
+  describe 'prompt build' do
     let(:task) { described_class.new(params) }
     let(:endpoint) { 'https://codesuggestions.gitlab.com/v2/code/generations' }
-    let(:body) { params.merge('prompt' => "```py\nsome text\n```\n", 'prompt_version' => 2) }
+    let(:body) do
+      params.merge(
+        'prompt' => prompt,
+        'prompt_version' => 2
+      )
+    end
+
+    it_behaves_like 'code suggestion task'
+
+    context 'when there is no filename extension' do
+      let(:file_name) { 'README' }
+      let(:prompt) do
+        <<~PROMPT
+          This is a task to write new  code in a file 'README' based on a given description.
+          You get first the already existing code file and then the description of the code that needs to be created.
+          It is your task to write valid and working  code.
+          Only return in your response new code.
+
+          Already existing code:
+
+          ```
+          some text
+          ```
+
+          Create new code for the following description:
+          `#{instruction}`
+        PROMPT
+      end
+
+      it_behaves_like 'code suggestion task'
+    end
   end
 end

ee/spec/requests/api/code_suggestions_spec.rb

@@ -318,12 +318,36 @@
 
         context 'when the task is code generation' do
           let(:current_user) { create(:user) }
-          let(:prefix) { '# GitLab Duo Generate: A function that outputs the first 20 fibonacci numbers' }
+          let(:instruction) { 'A function that outputs the first 20 fibonacci numbers' }
+          let(:prefix) do
+            <<~PREFIX
+              def is_even(n: int) ->
+              # Gitlab Duo Generate: #{instruction}
+            PREFIX
+          end
+
+          let(:prompt) do
+            <<~PROMPT
+              This is a task to write new Python code in a file 'test.py' based on a given description.
+              You get first the already existing code file and then the description of the code that needs to be created.
+              It is your task to write valid and working Python code.
+              Only return in your response new code.
+
+              Already existing code:
+
+              ```py
+              def is_even(n: int) ->
+              ```
+
+              Create new code for the following description:
+              `#{instruction}`
+            PROMPT
+          end
 
           it 'sends requests to the code generation endpoint' do
             expected_body = body.merge(
               prompt_version: 2,
-              prompt: "```py\n#{prefix}\n```\n"
+              prompt: prompt
             )
 
             expect(Gitlab::Workhorse)