De-obfuscate tofa/VertexAI code

This renames the codename Tofa to Vertex AI

ee/lib/api/ai/experimentation/tofa.rb → ee/lib/api/ai/experimentation/vertex_ai.rb

@@ -3,7 +3,7 @@
 module API
   module Ai
     module Experimentation
-      class Tofa < ::API::Base
+      class VertexAi < ::API::Base
         feature_category :not_owned # rubocop:todo Gitlab/AvoidFeatureCategoryNotOwned
         urgency :low
 
@@ -18,7 +18,7 @@ def check_feature_enabled
               Feature.enabled?(:ai_experimentation_api, current_user)
           end
 
-          def tofa_post(_endpoint, json_body: nil)
+          def vertex_ai_post(_endpoint, json_body: nil)
             headers = {
               "Accept" => ["application/json"],
               "Authorization" => ["Bearer #{tofa_api_token}"],
@@ -37,18 +37,29 @@ def tofa_post(_endpoint, json_body: nil)
 
           def default_payload_for(params)
             tofa_params = params.transform_keys { |name| name.camelize(:lower) }
-            json = JSON.parse(tofa_request_payload) # rubocop: disable Gitlab/Json
-            json_keys = tofa_request_json_keys.split(' ')
             content = tofa_params.delete(:content)
-            json[json_keys[0]][0][json_keys[1]][0][json_keys[2]] = content
-
-            json['parameters'].merge!(tofa_params)
+            json = {
+              instances: [
+                {
+                  messages: [
+                    {
+                      author: "content",
+                      content: content
+                    }
+                  ]
+                }
+              ],
+              parameters: {
+                temperature: ::Gitlab::Llm::VertexAi::Client::DEFAULT_TEMPERATURE
+              }
+            }
 
+            json[:parameters].merge!(tofa_params)
             json
           end
 
           def configuration
-            @configuration ||= Gitlab::Llm::Tofa::Configuration.new
+            @configuration ||= Gitlab::Llm::VertexAi::Configuration.new
           end
 
           def tofa_api_token
@@ -61,15 +72,13 @@ def tofa_api_token
 
           delegate(
             :host,
-            :tofa_request_json_keys,
-            :tofa_request_payload,
             :url,
             to: :configuration
           )
         end
 
         namespace 'ai/experimentation/tofa' do
-          desc 'Proxies request to Tofa chat endpoint'
+          desc 'Proxies request to Vertex AI chat endpoint'
           params do
             requires :content, type: String
             optional :temperature, type: Float, values: 0.0..1.0, default: 0.5
@@ -78,7 +87,7 @@ def tofa_api_token
             optional :top_p, type: Float, values: 0.0..1.0, default: 0.95
           end
           post 'chat' do
-            body tofa_post(
+            body vertex_ai_post(
               'chat', json_body: default_payload_for(declared(params, include_missing: false))
             )
           end

ee/lib/ee/api/api.rb

@@ -12,7 +12,7 @@ module API
 
         mount ::API::Admin::Search::Zoekt
         mount ::API::Ai::Experimentation::OpenAi
-        mount ::API::Ai::Experimentation::Tofa
+        mount ::API::Ai::Experimentation::VertexAi
         mount ::API::AuditEvents
         mount ::API::ProjectApprovalRules
         mount ::API::StatusChecks

ee/lib/gitlab/llm/open_ai/completions/explain_vulnerability.rb

@@ -6,7 +6,6 @@ module OpenAi
       module Completions
         class ExplainVulnerability
           DEFAULT_ERROR = 'An unexpected error has occurred.'
-          DIGIT_REGEX = /^\d$/
 
           def initialize(template_class)
             @template_class = template_class
@@ -14,27 +13,11 @@ def initialize(template_class)
 
           def execute(user, vulnerability, _options)
             template = template_class.new(vulnerability)
-            response = response_for(user, vulnerability, template)
+            response = response_for(user, template)
 
-            if tofa?(vulnerability)
-              json = Gitlab::Json.parse(response, symbolize_names: true)
-
-              GraphqlTriggers.ai_completion_response(
-                user.to_global_id,
-                vulnerability.to_global_id,
-                {
-                  id: SecureRandom.uuid,
-                  model_name: vulnerability.class.name,
-                  response_body: json.dig(*tofa_json_args),
-                  errors: tofa_errors(json)
-                }
-              )
-            else
-              ::Gitlab::Llm::OpenAi::ResponseService
-                .new(user, vulnerability, response, options: {})
-                .execute(Gitlab::Llm::OpenAi::ResponseModifiers::Chat.new)
-            end
-          # or refresh token grant
+            ::Gitlab::Llm::OpenAi::ResponseService
+              .new(user, vulnerability, response, options: {})
+              .execute(Gitlab::Llm::OpenAi::ResponseModifiers::Chat.new)
           rescue StandardError => error
             Gitlab::ErrorTracking.track_exception(error)
 
@@ -47,35 +30,12 @@ def execute(user, vulnerability, _options)
 
           attr_reader :template_class
 
-          def response_for(user, vulnerability, template)
-            client_class = client_class_for(vulnerability)
+          def response_for(user, template)
+            client_class = ::Gitlab::Llm::OpenAi::Client
             client_class
               .new(user)
               .chat(content: template.to_prompt, **template.options(client_class))
           end
-
-          def client_class_for(vulnerability)
-            if tofa?(vulnerability)
-              ::Gitlab::Llm::Tofa::Client
-            else
-              ::Gitlab::Llm::OpenAi::Client
-            end
-          end
-
-          def tofa_json_args
-            str = Gitlab::CurrentSettings.current_application_settings.tofa_response_json_keys
-            str.split.map do |s|
-              s.match?(DIGIT_REGEX) ? s.to_i : s.to_sym
-            end
-          end
-
-          def tofa_errors(json)
-            [json.dig(:error, :message)].compact
-          end
-
-          def tofa?(vulnerability)
-            Feature.enabled?(:tofa_experimentation, vulnerability.project)
-          end
         end
       end
     end

ee/lib/gitlab/llm/tofa/client.rb → ee/lib/gitlab/llm/vertex_ai/client.rb

@@ -2,9 +2,10 @@
 
 module Gitlab
   module Llm
-    module Tofa
+    module VertexAi
       class Client
         include ::Gitlab::Llm::Concerns::ExponentialBackoff
+        DEFAULT_TEMPERATURE = 0.5
 
         def initialize(_user, configuration = Configuration.new)
           @configuration = configuration
@@ -27,8 +28,6 @@ def chat(content:, **options)
         delegate(
           :access_token,
           :host,
-          :tofa_request_json_keys,
-          :tofa_request_payload,
           :url,
           to: :configuration
         )
@@ -43,11 +42,21 @@ def headers
         end
 
         def default_payload_for(content)
-          json = JSON.parse(tofa_request_payload) # rubocop: disable Gitlab/Json
-          json_keys = tofa_request_json_keys.split(' ')
-          json[json_keys[0]][0][json_keys[1]][0][json_keys[2]] = content
-
-          json
+          {
+            instances: [
+              {
+                messages: [
+                  {
+                    author: "content",
+                    content: content
+                  }
+                ]
+              }
+            ],
+            parameters: {
+              temperature: DEFAULT_TEMPERATURE
+            }
+          }
         end
       end
     end

ee/lib/gitlab/llm/vertex_ai/completions/explain_vulnerability.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Llm
+    module VertexAi
+      module Completions
+        class ExplainVulnerability
+          DEFAULT_ERROR = 'An unexpected error has occurred.'
+
+          def initialize(template_class)
+            @template_class = template_class
+          end
+
+          def execute(user, vulnerability, options)
+            unless vertex_ai?(vulnerability)
+              return ::Gitlab::Llm::OpenAi::Completions::ExplainVulnerability
+                .new(template_class)
+                .execute(user, vulnerability, options)
+            end
+
+            template = template_class.new(vulnerability)
+            response = response_for(user, template)
+
+            json = Gitlab::Json.parse(response, symbolize_names: true)
+
+            GraphqlTriggers.ai_completion_response(
+              user.to_global_id,
+              vulnerability.to_global_id,
+              {
+                id: SecureRandom.uuid,
+                model_name: vulnerability.class.name,
+                response_body: json.dig(:predictions, 0, :candidates, 0, :content),
+                errors: [json.dig(:error, :message)].compact
+              }
+            )
+          rescue StandardError => error
+            Gitlab::ErrorTracking.track_exception(error)
+
+            ::Gitlab::Llm::OpenAi::ResponseService
+              .new(user, vulnerability, { error: { message: DEFAULT_ERROR } }.to_json, options: {})
+              .execute
+          end
+
+          private
+
+          attr_reader :template_class
+
+          def response_for(user, template)
+            client_class = ::Gitlab::Llm::VertexAi::Client
+            client_class
+              .new(user)
+              .chat(content: template.to_prompt, **template.options(client_class))
+          end
+
+          def vertex_ai?(vulnerability)
+            Feature.enabled?(:tofa_experimentation, vulnerability.project)
+          end
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/llm/tofa/configuration.rb → ee/lib/gitlab/llm/vertex_ai/configuration.rb

@@ -2,12 +2,14 @@
 
 module Gitlab
   module Llm
-    module Tofa
+    module VertexAi
       class Configuration
+        DEFAULT_SCOPE = 'https://www.googleapis.com/auth/cloud-platform'
+
         def access_token
           Rails.cache.fetch(
             :tofa_access_token,
-            expires_in: tofa_access_token_expires_in.to_i.seconds,
+            expires_in: 3540.seconds,
             skip_nil: true
           ) do
             fresh_token
@@ -29,48 +31,20 @@ def settings
         end
 
         delegate(
-          :tofa_access_token_expires_in,
-          :tofa_client_library_args,
-          :tofa_client_library_class,
-          :tofa_client_library_create_credentials_method,
-          :tofa_client_library_fetch_access_token_method,
           :tofa_credentials,
           :tofa_host,
-          :tofa_request_json_keys,
-          :tofa_request_payload,
           :tofa_url,
           to: :settings
         )
 
         def fresh_token
-          client_library_class = tofa_client_library_class.constantize
-          client_library_args = string_to_hash(tofa_client_library_args)
-          first_key = client_library_args.first[0]
-          client_library_args[first_key] = StringIO.new(tofa_credentials)
-          create_credentials_method = tofa_client_library_create_credentials_method.to_sym
-          fetch_access_token_method = tofa_client_library_fetch_access_token_method.to_sym
-
-          # rubocop:disable GitlabSecurity/PublicSend
-          response = client_library_class.public_send(
-            create_credentials_method,
-            **client_library_args
-          ).public_send(fetch_access_token_method)
-          # rubocop:enable GitlabSecurity/PublicSend
+          response = ::Google::Auth::ServiceAccountCredentials.make_creds(
+            json_key_io: StringIO.new(tofa_credentials),
+            scope: DEFAULT_SCOPE
+          ).fetch_access_token!
 
           response["access_token"]
         end
-
-        def string_to_hash(str)
-          hash = {}
-
-          str.split(", ").each do |pair|
-            key, value = pair.split(" ")
-            value = value == "nil" ? nil : value
-            hash[key.to_sym] = value
-          end
-
-          hash
-        end
       end
     end
   end

ee/spec/lib/gitlab/llm/concerns/exponential_backoff_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe ::Gitlab::Llm::Concerns::ExponentialBackoff, feature_category: :no_category do # rubocop: disable RSpec/InvalidFeatureCategory
+RSpec.describe Gitlab::Llm::Concerns::ExponentialBackoff, feature_category: :no_category do # rubocop: disable RSpec/InvalidFeatureCategory
   let(:success) do
     instance_double(HTTParty::Response,
       code: 200, success?: true, parsed_response: {}, server_error?: false, too_many_requests?: false
@@ -29,7 +29,7 @@ def dummy_method(response_caller)
         response_caller.call
       end
 
-      include ::Gitlab::Llm::Concerns::ExponentialBackoff
+      include Gitlab::Llm::Concerns::ExponentialBackoff
       retry_methods_with_exponential_backoff :dummy_method
     end
   end

ee/spec/lib/gitlab/llm/open_ai/templates/explain_vulnerability_spec.rb

@@ -34,6 +34,26 @@
 
   subject { described_class.new(vulnerability) }
 
+  describe '#options' do
+    context 'for OpenAI' do
+      let(:client) { ::Gitlab::Llm::OpenAi::Client }
+
+      it 'returns max tokens' do
+        expect(subject.options(client)).to match(hash_including({
+          max_tokens: described_class::MAX_TOKENS
+        }))
+      end
+    end
+
+    context 'for VertexAI' do
+      let(:client) { ::Gitlab::Llm::VertexAi::Client }
+
+      it 'returns max tokens' do
+        expect(subject.options(client)).to be_empty
+      end
+    end
+  end
+
   describe '#to_prompt' do
     let(:identifiers) { vulnerability.finding.identifiers.pluck(:name).join(", ") }
 

ee/spec/lib/gitlab/llm/vertex_ai/client_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::VertexAi::Client, feature_category: :not_owned do # rubocop: disable  RSpec/InvalidFeatureCategory
+  let_it_be(:user) { create(:user) }
+
+  let(:access_token) { SecureRandom.uuid }
+
+  let(:configuration) do
+    instance_double(
+      ::Gitlab::Llm::VertexAi::Configuration,
+      access_token: access_token,
+      host: "example.com",
+      url: "https://example.com/api"
+    )
+  end
+
+  subject(:client) { described_class.new(user, configuration) }
+
+  describe "#chat" do
+    subject(:response) { client.chat(content: "Hello, world!") }
+
+    let(:successful_response) do
+      {
+        predictions: [
+          candidates: [
+            {
+              content: "Sure, ..."
+            }
+          ]
+        ]
+      }
+    end
+
+    context 'when a successful response is returned from the API' do
+      before do
+        request_body = {
+          instances: [
+            {
+              messages: [
+                {
+                  author: "content",
+                  content: "Hello, world!"
+                }
+              ]
+            }
+          ],
+          parameters: {
+            temperature: described_class::DEFAULT_TEMPERATURE
+          }
+        }
+
+        stub_request(:post, "https://example.com/api").with(
+          headers: {
+            accept: 'application/json',
+            'Authorization' => "Bearer #{access_token}",
+            'Content-Type' => 'application/json',
+            'Host' => 'example.com'
+          },
+          body: request_body.to_json
+        ).to_return(status: 200, body: successful_response.to_json)
+      end
+
+      it 'returns the response' do
+        expect(response).to be_present
+        expect(::Gitlab::Json.parse(response.body, symbolize_names: true)).to match(hash_including(
+          {
+            predictions: [
+              candidates: [
+                {
+                  content: "Sure, ..."
+                }
+              ]
+            ]
+          }
+        ))
+      end
+    end
+
+    context 'when a failed response is returned from the API' do
+      let(:too_many_requests_response) do
+        {
+          error: {
+            code: 429,
+            message: 'Rate Limit Exceeded',
+            status: 'RATE_LIMIT_EXCEEDED',
+            details: [
+              {
+                "@type": 'type.googleapis.com/google.rpc.ErrorInfo',
+                reason: 'RATE_LIMIT_EXCEEDED',
+                metadata: {
+                  service: 'aiplatform.googleapis.com',
+                  method: 'google.cloud.aiplatform.v1.PredictionService.Predict'
+                }
+              }
+            ]
+          }
+        }
+      end
+
+      before do
+        stub_request(:post, "https://example.com/api")
+          .to_return(status: 429, body: too_many_requests_response.to_json)
+          .then.to_return(status: 429, body: too_many_requests_response.to_json)
+          .then.to_return(status: 200, body: successful_response.to_json)
+
+        allow(client).to receive(:sleep).and_return(nil)
+      end
+
+      it 'retries the request' do
+        expect(response).to be_present
+        expect(response.code).to eq(200)
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/llm/vertex_ai/completions/explain_vulnerability_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::VertexAi::Completions::ExplainVulnerability, feature_category: :vulnerability_management do
+  let(:prompt_class) { Gitlab::Llm::OpenAi::Templates::ExplainVulnerability }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:project) do
+    create(:project, :custom_repo, files: {
+      'main.c' => "#include <stdio.h>\n\nint main() { printf(\"hello, world!\"); }"
+    })
+  end
+
+  let_it_be(:vulnerability) { create(:vulnerability, :with_finding, project: project) }
+
+  subject(:explain) { described_class.new(prompt_class) }
+
+  before do
+    allow(GraphqlTriggers).to receive(:ai_completion_response)
+    vulnerability.finding.location['file'] = 'main.c'
+    vulnerability.finding.location['start_line'] = 1
+  end
+
+  describe '#execute' do
+    context 'when the feature flag is disabled' do
+      before do
+        stub_feature_flags(tofa_experimentation: false)
+      end
+
+      it 'falls back to the OpenAI implementation' do
+        options = {}
+
+        allow_next_instance_of(::Gitlab::Llm::OpenAi::Completions::ExplainVulnerability) do |completion|
+          expect(completion).to receive(:execute).with(user, vulnerability, options)
+        end
+
+        explain.execute(user, vulnerability, options)
+
+        expect(GraphqlTriggers).not_to have_received(:ai_completion_response)
+      end
+    end
+
+    context 'when the chat client returns an unsuccessful response' do
+      before do
+        allow_next_instance_of(Gitlab::Llm::VertexAi::Client) do |client|
+          allow(client).to receive(:chat).and_return(
+            { 'error' => 'Ooops...' }.to_json
+          )
+        end
+      end
+
+      it 'publishes the error to the graphql subscription' do
+        explain.execute(user, vulnerability, {})
+
+        expect(GraphqlTriggers).to have_received(:ai_completion_response)
+          .with(user.to_global_id, vulnerability.to_global_id, hash_including({
+            id: anything,
+            model_name: vulnerability.class.name,
+            response_body: '',
+            errors: [{ message: described_class::DEFAULT_ERROR }]
+          }))
+      end
+    end
+
+    context 'when the chat client returns a successful response' do
+      let(:example_answer) { "Sure, ..." }
+
+      let(:example_response) do
+        {
+          "predictions" => [
+            {
+              "candidates" => [
+                {
+                  "author" => "",
+                  "content" => example_answer
+                }
+              ],
+              "safetyAttributes" => {
+                "categories" => ["Violent"],
+                "scores" => [0.4000000059604645],
+                "blocked" => false
+              }
+            }
+          ],
+          "deployedModelId" => "1",
+          "model" => "projects/1/locations/us-central1/models/codechat-bison-001",
+          "modelDisplayName" => "codechat-bison-001",
+          "modelVersionId" => "1"
+        }
+      end
+
+      before do
+        allow_next_instance_of(Gitlab::Llm::VertexAi::Client) do |client|
+          allow(client).to receive(:chat).and_return(example_response.to_json)
+        end
+      end
+
+      it 'publishes the content from the AI response' do
+        explain.execute(user, vulnerability, {})
+
+        expect(GraphqlTriggers).to have_received(:ai_completion_response)
+          .with(user.to_global_id, vulnerability.to_global_id, hash_including({
+            id: anything,
+            model_name: vulnerability.class.name,
+            response_body: example_answer,
+            errors: []
+          }))
+      end
+
+      context 'when an unexpected error is raised' do
+        let(:error) { StandardError.new("Ooops...") }
+
+        before do
+          allow_next_instance_of(Gitlab::Llm::VertexAi::Client) do |client|
+            allow(client).to receive(:chat).and_raise(error)
+          end
+          allow(Gitlab::ErrorTracking).to receive(:track_exception)
+        end
+
+        it 'records the error' do
+          explain.execute(user, vulnerability, {})
+          expect(Gitlab::ErrorTracking).to have_received(:track_exception).with(error)
+        end
+
+        it 'publishes a generic error to the graphql subscription' do
+          explain.execute(user, vulnerability, {})
+
+          expect(GraphqlTriggers).to have_received(:ai_completion_response)
+            .with(user.to_global_id, vulnerability.to_global_id, hash_including({
+              id: anything,
+              model_name: vulnerability.class.name,
+              response_body: '',
+              errors: [{ message: 'An unexpected error has occurred.' }]
+            }))
+        end
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/llm/vertex_ai/configuration_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::VertexAi::Configuration, feature_category: :not_owned do # rubocop: disable  RSpec/InvalidFeatureCategory
+  subject(:configuration) { described_class.new }
+
+  describe '#access_token', :clean_gitlab_redis_cache do
+    context 'when the token is cached', :use_clean_rails_redis_caching do
+      let(:cached_token) { SecureRandom.uuid }
+
+      before do
+        Rails.cache.write(:tofa_access_token, cached_token)
+      end
+
+      it 'returns the cached token' do
+        expect(configuration.access_token).to eq(cached_token)
+      end
+    end
+
+    context 'when an access token has not been minted yet' do
+      let(:access_token) { "x.#{SecureRandom.uuid}.z" }
+      let(:private_key) { OpenSSL::PKey::RSA.new(4096) }
+      let(:credentials) do
+        {
+          type: "service_account",
+          project_id: SecureRandom.uuid,
+          private_key_id: SecureRandom.hex(20),
+          private_key: private_key.to_pem,
+          client_email: "vertex-ai@#{SecureRandom.hex(4)}.iam.gserviceaccount.com",
+          client_id: "1",
+          auth_uri: "https://accounts.google.com/o/oauth2/auth",
+          token_uri: "https://oauth2.googleapis.com/token",
+          auth_provider_x509_cert_url: "https://www.googleapis.com/oauth2/v1/certs",
+          client_x509_cert_url: "https://www.googleapis.com/robot/v1/metadata/x509/vertex-ai.iam.gserviceaccount.com"
+        }
+      end
+
+      before do
+        stub_application_setting(tofa_credentials: credentials.to_json)
+
+        stub_request(:post, "https://www.googleapis.com/oauth2/v4/token").to_return(
+          status: 200,
+          headers: { 'content-type' => 'application/json; charset=utf-8' },
+          body: {
+            access_token: access_token,
+            expires_in: 3600,
+            scope: "https://www.googleapis.com/auth/cloud-platform",
+            token_type: "Bearer"
+          }.to_json
+        ).times(1)
+      end
+
+      it 'generates a new token' do
+        expect(subject.access_token).to eql(access_token)
+      end
+    end
+  end
+
+  describe '#host' do
+    let(:host) { "example-#{SecureRandom.hex(8)}.com" }
+
+    before do
+      stub_application_setting(tofa_host: host)
+    end
+
+    it { expect(configuration.host).to eql(host) }
+  end
+
+  describe '#url' do
+    let(:host) { "example-#{SecureRandom.hex(8)}.com" }
+    let(:url) { "https://#{host}/api" }
+
+    before do
+      stub_application_setting(tofa_url: url)
+    end
+
+    it { expect(configuration.url).to eql(url) }
+  end
+end

ee/spec/requests/api/ai/experimentation/tofa_spec.rb → ee/spec/requests/api/ai/experimentation/vertex_ai_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe API::Ai::Experimentation::Tofa, feature_category: :shared do
+RSpec.describe API::Ai::Experimentation::VertexAi, feature_category: :shared do
   let_it_be(:current_user) { create(:user) }
   let(:body) { { 'test' => 'test' } }
   let(:token) { create(:personal_access_token, user: current_user) }
@@ -17,11 +17,8 @@
   end
 
   before do
-    allow_next_instance_of(Gitlab::Llm::Tofa::Configuration) do |configuration|
+    allow_next_instance_of(Gitlab::Llm::VertexAi::Configuration) do |configuration|
       allow(configuration).to receive(:access_token).and_return(token)
-      allow(configuration).to receive(:tofa_request_payload)
-        .and_return('{"a":[{"b": [{"c": "{{CONTENT}}"}]}], "parameters": {}}')
-      allow(configuration).to receive(:tofa_request_json_keys).and_return('a b c')
     end
 
     stub_feature_flags(tofa_experimentation: true)