Merge branch '490585-migrate_libraries_to_aigw' into 'master'

migrate context_block, libraries_block, user_instruction for anthropic

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



Merged-by: Leaminn Ma <lma@gitlab.com>
Approved-by: Missy Davies <mdavies@gitlab.com>
Approved-by: Leaminn Ma <lma@gitlab.com>
Reviewed-by: Leaminn Ma <lma@gitlab.com>
Reviewed-by: Duo Code Reviewer <duo-code-review-bot@gitlab.com>
Co-authored-by: tgao3701908 <tgao@gitlab.com>

ee/lib/code_suggestions/prompts/code_generation/ai_gateway_messages.rb

@@ -12,7 +12,6 @@ class AiGatewayMessages < CodeSuggestions::Prompts::Base
         # response time grows with prompt size, so we don't use upper limit size of prompt window
         MAX_INPUT_CHARS = 50000
         GATEWAY_PROMPT_VERSION = 3
-        CONTENT_TYPES = { file: 'file', snippet: 'snippet' }.freeze
 
         def request_params
           {
@@ -37,7 +36,10 @@ def request_params
         def code_generation_enhancer
           {
             **examples_section_params,
-            **existing_code_block_params
+            **existing_code_block_params,
+            **context_block_params,
+            **libraries_block_params,
+            **user_instruction_params
           }
         end
 
@@ -50,8 +52,6 @@ def examples_section_params
         end
 
         def existing_code_block_params
-          return {} unless params[:prefix].present?
-
           trimmed_prefix = prefix.to_s.last(MAX_INPUT_CHARS)
           trimmed_suffix = suffix.to_s.first(MAX_INPUT_CHARS - trimmed_prefix.size)
 
@@ -60,6 +60,64 @@ def existing_code_block_params
             trimmed_suffix: trimmed_suffix
           }
         end
+
+        def context_block_params
+          related_files = []
+          related_snippets = []
+
+          params[:context]&.each do |context|
+            if context[:type] == ::Ai::AdditionalContext::CODE_SUGGESTIONS_CONTEXT_TYPES[:file]
+              related_files << <<~FILE_CONTENT
+              <file_content file_name="#{context[:name]}">
+              #{context[:content]}
+              </file_content>
+              FILE_CONTENT
+            elsif context[:type] == ::Ai::AdditionalContext::CODE_SUGGESTIONS_CONTEXT_TYPES[:snippet]
+              related_snippets << <<~SNIPPET_CONTENT
+              <snippet_content name="#{context[:name]}">
+              #{context[:content]}
+              </snippet_content>
+              SNIPPET_CONTENT
+            end
+          end
+
+          {
+            related_files: related_files,
+            related_snippets: related_snippets
+          }
+        end
+
+        def libraries_block_params
+          if libraries.any?
+            Gitlab::InternalEvents.track_event(
+              'include_repository_xray_data_into_code_generation_prompt',
+              project: params[:project],
+              namespace: params[:project]&.namespace,
+              user: params[:current_user]
+            )
+          end
+
+          { libraries: libraries }
+        end
+
+        def libraries
+          return [] unless xray_report
+
+          xray_report.libs.map { |l| l['name'] } # rubocop:disable Rails/Pluck -- libs is an array
+        end
+        strong_memoize_attr :libraries
+
+        def xray_report
+          ::Projects::XrayReport.for_project(params[:project]).for_lang(language.x_ray_lang).first
+        end
+        strong_memoize_attr :xray_report
+
+        def user_instruction_params
+          instruction = params[:instruction]&.instruction.presence ||
+            'Generate the best possible code based on instructions.'
+
+          { user_instruction: instruction }
+        end
       end
     end
   end

ee/spec/lib/code_suggestions/prompts/code_generation/ai_gateway_messages_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_relative 'ai_gateway_shared_examples'
+
+RSpec.describe CodeSuggestions::Prompts::CodeGeneration::AiGatewayMessages, feature_category: :code_suggestions do
+  let(:prompt_version) { 3 }
+
+  it_behaves_like 'code generation AI Gateway request params' do
+    def expected_request_params
+      {
+        prompt_components: [
+          {
+            type: 'code_editor_generation',
+            payload: {
+              file_name: expected_file_name,
+              content_above_cursor: expected_content_above_cursor,
+              content_below_cursor: expected_content_below_cursor,
+              language_identifier: expected_language_identifier,
+              prompt_id: 'code_suggestions/generations',
+              prompt_enhancer: {
+                examples_array: expected_examples_array,
+                trimmed_prefix: expected_trimmed_prefix,
+                trimmed_suffix: expected_trimmed_suffix,
+                related_files: expected_related_files,
+                related_snippets: expected_related_snippets,
+                libraries: expected_libraries,
+                user_instruction: expected_user_instruction
+              }
+            }
+          }
+        ]
+      }
+    end
+  end
+end

ee/spec/lib/code_suggestions/prompts/code_generation/anthropic_shared_examples.rb → ee/spec/lib/code_suggestions/prompts/code_generation/ai_gateway_shared_examples.rb

 # frozen_string_literal: true
 
-RSpec.shared_examples 'anthropic prompt' do
+# AIGW v2 api for code generation receives prompt
+RSpec.shared_examples 'code generation AI Gateway request prompt' do
   let(:language) { instance_double(CodeSuggestions::ProgrammingLanguage, x_ray_lang: x_ray_lang) }
   let(:language_name) { 'Go' }
   let(:x_ray_lang) { nil }
@@ -316,19 +317,19 @@
       context 'when context is available' do
         let(:main_go_content) do
           <<~CONTENT
-          package main
+            package main
 
-          func main()
-            fullName("John", "Doe")
-          }
+            func main()
+              fullName("John", "Doe")
+            }
           CONTENT
         end
 
         let(:full_name_func_content) do
           <<~CONTENT
-          func fullName(first, last string) {
-            fmt.Println(first, last)
-          }
+            func fullName(first, last string) {
+              fmt.Println(first, last)
+            }
           CONTENT
         end
 
@@ -615,3 +616,163 @@
     end
   end
 end
+
+# AIGW v3 api for code generation receives params
+RSpec.shared_examples 'code generation AI Gateway request params' do
+  let_it_be(:current_user) { create(:user) }
+  let_it_be(:project) { create(:project) }
+
+  let(:file_name) { 'main.go' }
+  let(:content_above_cursor) { "package main\n\nimport \"fmt\"\n\nfunc main() {\n" }
+  let(:content_below_cursor) { "func test() {\n" }
+  let(:comment) { 'My comment instructions' }
+  let(:instruction) { instance_double(CodeSuggestions::Instruction, instruction: comment, trigger_type: 'comment') }
+  let(:examples) { [{ example: 'func hello() {', response: 'func hello() {<new_code>fmt.Println("hello")' }] }
+
+  let(:context) do
+    [
+      { type: 'file', name: 'main.go', content:
+        <<~CONTENT
+          package main
+
+          func main()
+            fullName("John", "Doe")
+          }
+        CONTENT
+      },
+      { type: 'snippet', name: 'fullName', content:
+        <<~CONTENT
+          func fullName(first, last string) {
+            fmt.Println(first, last)
+          }
+        CONTENT
+      }
+    ]
+  end
+
+  let(:current_file_params) do
+    {
+      file_name: file_name,
+      content_above_cursor: content_above_cursor,
+      content_below_cursor: content_below_cursor
+    }
+  end
+
+  let(:params) do
+    {
+      current_user: current_user,
+      project: project,
+      instruction: instruction,
+      current_file: current_file_params,
+      context: context
+    }
+  end
+
+  subject { described_class.new(params) }
+
+  describe '#request_params' do
+    context 'when all parameters are present' do
+      before_all do
+        create(:xray_report, lang: 'go', project: project,
+          payload: { libs: [{ name: 'zlib (1.2.3)' }, { name: 'boost (2.0.0)' }] })
+      end
+
+      let(:expected_file_name) { file_name }
+      let(:expected_content_above_cursor) { content_above_cursor }
+      let(:expected_content_below_cursor) { content_below_cursor }
+      let(:expected_language_identifier) { 'Go' }
+      let(:expected_examples_array) { examples }
+      let(:expected_trimmed_prefix) { content_above_cursor }
+      let(:expected_trimmed_suffix) { content_below_cursor }
+      let(:expected_libraries) { ['zlib (1.2.3)', 'boost (2.0.0)'] }
+      let(:expected_user_instruction) { comment }
+
+      let(:expected_related_files) do
+        [
+          "<file_content file_name=\"main.go\">\npackage main\n\nfunc main()\n  " \
+            "fullName(\"John\", \"Doe\")\n}\n\n</file_content>\n"
+        ]
+      end
+
+      let(:expected_related_snippets) do
+        [
+          "<snippet_content name=\"fullName\">\nfunc fullName(first, last string) {\n  " \
+            "fmt.Println(first, last)\n}\n\n</snippet_content>\n"
+        ]
+      end
+
+      before do
+        allow_next_instance_of(CodeSuggestions::ProgrammingLanguage) do |instance|
+          allow(instance).to receive(:generation_examples).with(type: instruction.trigger_type).and_return(examples)
+        end
+      end
+
+      it 'returns expected request params' do
+        expect(subject.request_params).to eq(expected_request_params)
+      end
+
+      it 'tracks an X-Ray event' do
+        expect(Gitlab::InternalEvents).to receive(:track_event).with(
+          'include_repository_xray_data_into_code_generation_prompt',
+          project: project,
+          namespace: project.namespace,
+          user: current_user
+        )
+
+        subject.request_params
+      end
+
+      context 'when the prefix length exceeds the prompt limit' do
+        let(:limit) { 10 }
+        let(:expected_trimmed_prefix) { content_above_cursor.last(limit) }
+        let(:expected_trimmed_suffix) { '' }
+
+        before do
+          stub_const('CodeSuggestions::Prompts::CodeGeneration::AiGatewayMessages::MAX_INPUT_CHARS', limit)
+        end
+
+        it 'returns expected request params' do
+          expect(subject.request_params).to eq(expected_request_params)
+        end
+
+        context 'when the combined prefix and suffix length exceeds the prompt limit' do
+          let(:limit) { content_above_cursor.size + 5 }
+          let(:expected_trimmed_prefix) { content_above_cursor }
+          let(:expected_trimmed_suffix) { content_below_cursor.first(5) }
+
+          it 'returns expected request params' do
+            expect(subject.request_params).to eq(expected_request_params)
+          end
+        end
+      end
+    end
+
+    context 'when all parameters are blank' do
+      let(:instruction) { nil }
+      let(:context) { nil }
+      let(:current_file_params) { nil }
+
+      let(:expected_file_name) { '' }
+      let(:expected_content_above_cursor) { nil }
+      let(:expected_content_below_cursor) { nil }
+      let(:expected_language_identifier) { '' }
+      let(:expected_examples_array) { [] }
+      let(:expected_trimmed_prefix) { '' }
+      let(:expected_trimmed_suffix) { '' }
+      let(:expected_libraries) { [] }
+      let(:expected_user_instruction) { 'Generate the best possible code based on instructions.' }
+      let(:expected_related_files) { [] }
+      let(:expected_related_snippets) { [] }
+
+      it 'returns expected request params' do
+        expect(subject.request_params).to eq(expected_request_params)
+      end
+
+      it 'does not track an X-Ray event' do
+        expect(Gitlab::InternalEvents).not_to receive(:track_event)
+
+        subject.request_params
+      end
+    end
+  end
+end

ee/spec/lib/code_suggestions/prompts/code_generation/anthropic_messages_spec.rb

 # frozen_string_literal: true
 
 require 'spec_helper'
-require_relative 'anthropic_shared_examples'
+require_relative 'ai_gateway_shared_examples'
 
 RSpec.describe CodeSuggestions::Prompts::CodeGeneration::AnthropicMessages, feature_category: :code_suggestions do
   let(:prompt_version) { 3 }
 
-  it_behaves_like 'anthropic prompt' do
+  it_behaves_like 'code generation AI Gateway request prompt' do
     def expected_prompt
       [
         { role: :system, content: system_prompt },

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

@@ -44,30 +44,34 @@
 
     let(:anthropic_request_params) do
       {
-        "prompt_components" => [
+        'prompt_components' => [
           {
-            "type" => "code_editor_generation",
-            "payload" => {
-              "file_name" => "test.py",
-              "content_above_cursor" => "some prefix",
-              "content_below_cursor" => "some suffix",
-              "language_identifier" => "Python",
-              "prompt_id" => "code_suggestions/generations",
-              "prompt_enhancer" => {
-                "examples_array" => [
+            'type' => 'code_editor_generation',
+            'payload' => {
+              'file_name' => 'test.py',
+              'content_above_cursor' => 'some prefix',
+              'content_below_cursor' => 'some suffix',
+              'language_identifier' => 'Python',
+              'prompt_id' => 'code_suggestions/generations',
+              'prompt_enhancer' => {
+                'examples_array' => [
                   {
-                    "example" => "class Project:\\n  def __init__(self, name, public):{{cursor}}\\n\\n ",
-                    "response" => "return self.visibility == 'PUBLIC'",
-                    "trigger_type" => "comment"
+                    'example' => 'class Project:\\n  def __init__(self, name, public):{{cursor}}\\n\\n ',
+                    'response' => "return self.visibility == 'PUBLIC'",
+                    'trigger_type' => 'comment'
                   },
                   {
-                    "example" => "# get the current user's name from the session data\\n{{cursor}}",
-                    "response" => "username = session['username']\\nreturn username",
-                    "trigger_type" => "comment"
+                    'example' => "# get the current user's name from the session data\\n{{cursor}}",
+                    'response' => "username = session['username']\\nreturn username",
+                    'trigger_type' => 'comment'
                   }
                 ],
-                "trimmed_prefix" => "some prefix",
-                "trimmed_suffix" => "some suffix"
+                'trimmed_prefix' => 'some prefix',
+                'trimmed_suffix' => 'some suffix',
+                'related_files' => '',
+                'related_snippets' => '',
+                'libraries' => '',
+                'user_instruction' => 'Generate the best possible code based on instructions.'
               }
             }
           }
@@ -94,40 +98,44 @@
       let(:endpoint_path) { 'v3/code/completions' }
       let(:expected_body) do
         {
-          "current_file" => {
-            "content_above_cursor" => "fix",
-            "content_below_cursor" => "som",
-            "file_name" => "test.py"
+          'current_file' => {
+            'content_above_cursor' => 'fix',
+            'content_below_cursor' => 'som',
+            'file_name' => 'test.py'
           },
-          "prompt_components" => [
+          'prompt_components' => [
             {
-              "payload" => {
-                "content_above_cursor" => "some prefix",
-                "content_below_cursor" => "some suffix",
-                "file_name" => "test.py",
-                "language_identifier" => "Python",
-                "prompt_enhancer" => {
-                  "examples_array" => [
+              'payload' => {
+                'content_above_cursor' => 'some prefix',
+                'content_below_cursor' => 'some suffix',
+                'file_name' => 'test.py',
+                'language_identifier' => 'Python',
+                'prompt_enhancer' => {
+                  'examples_array' => [
                     {
-                      "example" => "class Project:\\n  def __init__(self, name, public):{{cursor}}\\n\\n ",
-                      "response" => "return self.visibility == 'PUBLIC'",
-                      "trigger_type" => "comment"
+                      'example' => 'class Project:\\n  def __init__(self, name, public):{{cursor}}\\n\\n ',
+                      'response' => "return self.visibility == 'PUBLIC'",
+                      'trigger_type' => 'comment'
                     },
                     {
-                      "example" => "# get the current user's name from the session data\\n{{cursor}}",
-                      "response" => "username = session['username']\\nreturn username",
-                      "trigger_type" => "comment"
+                      'example' => "# get the current user's name from the session data\\n{{cursor}}",
+                      'response' => "username = session['username']\\nreturn username",
+                      'trigger_type' => 'comment'
                     }
                   ],
-                  "trimmed_prefix" => "some prefix",
-                  "trimmed_suffix" => "some suffix"
+                  'trimmed_prefix' => 'some prefix',
+                  'trimmed_suffix' => 'some suffix',
+                  'related_files' => '',
+                  'related_snippets' => '',
+                  'libraries' => '',
+                  'user_instruction' => 'Generate the best possible code based on instructions.'
                 },
-                "prompt_id" => "code_suggestions/generations"
+                'prompt_id' => 'code_suggestions/generations'
               },
-              "type" => "code_editor_generation"
+              'type' => 'code_editor_generation'
             }
           ],
-          "telemetry" => [{ "model_engine" => "anthropic" }]
+          'telemetry' => [{ 'model_engine' => 'anthropic' }]
         }
       end
 

ee/spec/requests/api/code_suggestions_spec.rb

@@ -191,9 +191,13 @@ def is_even(n: int) ->
                     "trigger_type" => "comment"
                   }
                 ],
-                "trimmed_prefix" => "def is_even(n: int) ->\n# A " \
+                'trimmed_prefix' => "def is_even(n: int) ->\n# A " \
                   "function that outputs the first 20 fibonacci numbers\n",
-                "trimmed_suffix" => ""
+                'trimmed_suffix' => '',
+                'related_files' => [],
+                'related_snippets' => [],
+                'libraries' => [],
+                'user_instruction' => 'Generate the best possible code based on instructions.'
               }
             }
           }