From 8d2318656d5b79448ef5aee6a2ab949bd7f98685 Mon Sep 17 00:00:00 2001
From: Jan Provaznik <jprovaznik@gitlab.com>
Date: Fri, 31 May 2024 11:21:30 +0000
Subject: [PATCH] Automatically test connection on admin page
When admin visits "GitLab Duo Pro" admin page, we automatically run a
simple test completion request to check if code suggestions request work
as expected and show this result as a flash message.
---
.../admin/code_suggestions_controller.rb | 8 +++
.../llm/ai_gateway/code_suggestions_client.rb | 70 +++++++++++++++++++
.../code_suggestions_client_spec.rb | 59 ++++++++++++++++
.../admin/code_suggestions_controller_spec.rb | 33 +++++++--
locale/gitlab.pot | 6 ++
5 files changed, 172 insertions(+), 4 deletions(-)
create mode 100644 ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
create mode 100644 ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
diff --git a/ee/app/controllers/admin/code_suggestions_controller.rb b/ee/app/controllers/admin/code_suggestions_controller.rb
index 02e252160b0d8..af3f1b30c573d 100644
--- a/ee/app/controllers/admin/code_suggestions_controller.rb
+++ b/ee/app/controllers/admin/code_suggestions_controller.rb
@@ -13,6 +13,14 @@ class CodeSuggestionsController < Admin::ApplicationController
before_action :ensure_feature_available!
def index
+ error = ::Gitlab::Llm::AiGateway::CodeSuggestionsClient.new(current_user).test_completion
+
+ if error.blank?
+ flash[:notice] = _('Code completion test was successful')
+ else
+ flash[:alert] = format(_('Code completion test failed: %{error}'), error: error)
+ end
+
@subscription_name = License.current.subscription_name
end
diff --git a/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb b/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
new file mode 100644
index 0000000000000..691181fe13f95
--- /dev/null
+++ b/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module Gitlab
+ module Llm
+ module AiGateway
+ class CodeSuggestionsClient
+ include ::Gitlab::Utils::StrongMemoize
+ include ::API::Helpers::CloudConnector
+
+ COMPLETION_CHECK_TIMEOUT = 3.seconds
+
+ def initialize(user)
+ @user = user
+ end
+
+ def test_completion
+ return 'Access token is missing' unless access_token
+
+ response = Gitlab::HTTP.post(
+ task.endpoint,
+ headers: request_headers,
+ body: task.body,
+ timeout: COMPLETION_CHECK_TIMEOUT,
+ allow_local_requests: true
+ )
+
+ return "AI Gateway returned code #{response.code}: #{response.body}" unless response.code == 200
+ return "Response doesn't contain a completion" unless choice?(response)
+
+ nil
+ rescue StandardError => err
+ err.message
+ end
+
+ private
+
+ attr_reader :user
+
+ def request_headers
+ {
+ 'X-Gitlab-Authentication-Type' => 'oidc',
+ 'Authorization' => "Bearer #{access_token}",
+ 'Content-Type' => 'application/json',
+ 'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id
+ }.merge(cloud_connector_headers(user))
+ end
+
+ def access_token
+ ::CloudConnector::AvailableServices.find_by_name(:code_suggestions).access_token(user)
+ end
+ strong_memoize_attr :access_token
+
+ def choice?(response)
+ response['choices']&.first&.dig('text').present?
+ end
+
+ def task
+ params = {
+ current_file: {
+ file_name: 'test.rb',
+ content_above_cursor: 'def hello_world'
+ }
+ }
+ CodeSuggestions::Tasks::CodeCompletion.new(unsafe_passthrough_params: params)
+ end
+ strong_memoize_attr :task
+ end
+ end
+ end
+end
diff --git a/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb b/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
new file mode 100644
index 0000000000000..f67ff8c615b7a
--- /dev/null
+++ b/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::AiGateway::CodeSuggestionsClient, feature_category: :code_suggestions do
+ let_it_be(:user) { create(:user) }
+
+ describe "#test_completion" do
+ let_it_be(:token) { create(:service_access_token, :active) }
+ let(:body) { { choices: [{ text: "puts \"Hello World!\"\nend", index: 0, finish_reason: "length" }] } }
+ let(:code) { 200 }
+
+ subject(:result) { described_class.new(user).test_completion }
+
+ shared_examples "error response" do |message|
+ it "returns an error" do
+ expect(result).to eq(message)
+ end
+ end
+
+ before do
+ stub_request(:post, /#{Gitlab::AiGateway.url}/)
+ .to_return(status: code, body: body.to_json, headers: { "Content-Type" => "application/json" })
+ allow(CloudConnector::AvailableServices).to receive_message_chain(:find_by_name,
+ :access_token).and_return(token)
+ end
+
+ it 'returns nil if there is no error' do
+ expect(result).to be_nil
+ end
+
+ context 'when there is not valid token' do
+ let(:token) { nil }
+
+ it_behaves_like 'error response', "Access token is missing"
+ end
+
+ context 'when response does not contain a valid choice' do
+ let(:body) { { choices: [] } }
+
+ it_behaves_like 'error response', "Response doesn't contain a completion"
+ end
+
+ context 'when response code is not 200' do
+ let(:code) { 401 }
+ let(:body) { 'an error' }
+
+ it_behaves_like 'error response', 'AI Gateway returned code 401: "an error"'
+ end
+
+ context 'when request raises an error' do
+ before do
+ stub_request(:post, /#{Gitlab::AiGateway.url}/).to_raise(StandardError.new('an error'))
+ end
+
+ it_behaves_like 'error response', 'an error'
+ end
+ end
+end
diff --git a/ee/spec/requests/admin/code_suggestions_controller_spec.rb b/ee/spec/requests/admin/code_suggestions_controller_spec.rb
index c8a420e3429e2..5cbeebd63cc47 100644
--- a/ee/spec/requests/admin/code_suggestions_controller_spec.rb
+++ b/ee/spec/requests/admin/code_suggestions_controller_spec.rb
@@ -15,11 +15,36 @@
end
shared_examples 'renders the activation form' do
- it 'renders the activation form' do
- get admin_code_suggestions_path
+ context 'when connection check succeeds' do
+ before do
+ allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
+ allow(client).to receive(:test_completion).and_return(nil)
+ end
+ end
+
+ it 'renders the activation form' do
+ get admin_code_suggestions_path
+
+ expect(response).to render_template(:index)
+ expect(response.body).to include('js-code-suggestions-page')
+ expect(flash[:notice]).to eq("Code completion test was successful")
+ end
+ end
- expect(response).to render_template(:index)
- expect(response.body).to include('js-code-suggestions-page')
+ context 'when connection check fails' do
+ before do
+ allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
+ allow(client).to receive(:test_completion).and_return('an error')
+ end
+ end
+
+ it 'renders the activation form with alert message' do
+ get admin_code_suggestions_path
+
+ expect(response).to render_template(:index)
+ expect(response.body).to include('js-code-suggestions-page')
+ expect(flash[:alert]).to eq("Code completion test failed: an error")
+ end
end
end
diff --git a/locale/gitlab.pot b/locale/gitlab.pot
index d5409593aff63..d48f0beba2ac3 100644
--- a/locale/gitlab.pot
+++ b/locale/gitlab.pot
@@ -12583,6 +12583,12 @@ msgstr ""
msgid "Code block"
msgstr ""
+msgid "Code completion test failed: %{error}"
+msgstr ""
+
+msgid "Code completion test was successful"
+msgstr ""
+
msgid "Code coverage statistics for %{ref} %{start_date} - %{end_date}"
msgstr ""
--
GitLab