diff --git a/ee/lib/code_suggestions/model_details/code_completion.rb b/ee/lib/code_suggestions/model_details/code_completion.rb
index 527c9d72df33aef439afcd48dd94b876a45da4c7..162cbcd4297082edf528e262a4599c09c2f15e0f 100644
--- a/ee/lib/code_suggestions/model_details/code_completion.rb
+++ b/ee/lib/code_suggestions/model_details/code_completion.rb
@@ -15,6 +15,8 @@ def current_model
return fireworks_qwen_2_5_model_details if use_fireworks_qwen_for_code_completions?
+ return codestral_2502_vertex_model_details if code_gecko_disabled? && !user_opted_in_to_qwen?
+
# the default behavior is returning an empty hash
# AI Gateway will fall back to the code-gecko model if model details are not provided
{}
@@ -25,6 +27,10 @@ def saas_primary_model_class
return CodeSuggestions::Prompts::CodeCompletion::FireworksQwen if use_fireworks_qwen_for_code_completions?
+ if code_gecko_disabled? && !user_opted_in_to_qwen?
+ return CodeSuggestions::Prompts::CodeCompletion::CodestralVertex
+ end
+
CodeSuggestions::Prompts::CodeCompletion::VertexAi
end
@@ -37,6 +43,13 @@ def fireworks_qwen_2_5_model_details
}
end
+ def codestral_2502_vertex_model_details
+ {
+ model_provider: CodeSuggestions::Prompts::CodeCompletion::CodestralVertex::MODEL_PROVIDER,
+ model_name: CodeSuggestions::Prompts::CodeCompletion::CodestralVertex::MODEL_NAME
+ }
+ end
+
# We introduced an ops FF to allow organizations to opt out of Fireworks/Qwen.
# On GitLab SaaS, Duo access is managed by top-level group,
# so we are checking the FF by the user's top-level group
@@ -48,7 +61,11 @@ def use_fireworks_qwen_for_code_completions?
return false if Feature.disabled?(:fireworks_qwen_code_completion, current_user, type: :beta)
# if the beta FF is enabled, proceed to check the ops FF
+ user_opted_in_to_qwen?
+ end
+ # check if user's groups or instance has opted out of Fireworks/Qwen
+ def user_opted_in_to_qwen?
# on saas, check the user's groups
return all_user_groups_opted_in_to_fireworks_qwen? if Gitlab.org_or_com? # rubocop: disable Gitlab/AvoidGitlabInstanceChecks -- see comment above method definition
@@ -56,6 +73,10 @@ def use_fireworks_qwen_for_code_completions?
instance_opted_in_to_fireworks_qwen?
end
+ def code_gecko_disabled?
+ Feature.enabled?(:disable_code_gecko_default, current_user)
+ end
+
def all_user_groups_opted_in_to_fireworks_qwen?
# while this has potential to be expensive in terms of Feature Flag checking,
# we are expecting that for most if not all users, only 1 group is providing them the Duo Access
diff --git a/ee/lib/code_suggestions/prompts/code_completion/codestral_vertex.rb b/ee/lib/code_suggestions/prompts/code_completion/codestral_vertex.rb
new file mode 100644
index 0000000000000000000000000000000000000000..ac9c50a1a4e15bc31a750f6a6067361433afd5e8
--- /dev/null
+++ b/ee/lib/code_suggestions/prompts/code_completion/codestral_vertex.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module CodeSuggestions
+ module Prompts
+ module CodeCompletion
+ class CodestralVertex < CodeSuggestions::Prompts::Base
+ GATEWAY_PROMPT_VERSION = 1
+ MODEL_NAME = 'codestral-2501'
+ MODEL_PROVIDER = 'vertex-ai'
+
+ def request_params
+ {
+ prompt_version: GATEWAY_PROMPT_VERSION,
+ model_name: MODEL_NAME,
+ model_provider: MODEL_PROVIDER
+ }
+ end
+ end
+ end
+ end
+end
diff --git a/ee/spec/lib/code_suggestions/model_details/code_completion_spec.rb b/ee/spec/lib/code_suggestions/model_details/code_completion_spec.rb
index 8403ea98bfd3c380fb42e27ca96598c267e33342..dee50492093dcac0df1b709d609874b41ec85a0c 100644
--- a/ee/spec/lib/code_suggestions/model_details/code_completion_spec.rb
+++ b/ee/spec/lib/code_suggestions/model_details/code_completion_spec.rb
@@ -13,6 +13,7 @@
before do
stub_feature_flags(fireworks_qwen_code_completion: true)
stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: false)
+ stub_feature_flags(disable_code_gecko_default: false)
end
context 'on GitLab self-managed' do
@@ -28,7 +29,16 @@
it 'returns the codegecko model' do
stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
- expect(actual_result).to eq(expected_codegecko_result)
+ expect(actual_result).to eq(expected_default_result)
+ end
+
+ context 'when code gecko is disabled' do
+ it 'returns codestral on vertex' do
+ stub_feature_flags(disable_code_gecko_default: true)
+ stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
+
+ expect(actual_result).to eq(expected_codestral_result)
+ end
end
end
end
@@ -61,7 +71,15 @@
end
it 'returns the code gecko model' do
- expect(actual_result).to eq(expected_codegecko_result)
+ expect(actual_result).to eq(expected_default_result)
+ end
+
+ context 'when code gecko is disabled' do
+ it 'returns codestral on vertex' do
+ stub_feature_flags(disable_code_gecko_default: true)
+
+ expect(actual_result).to eq(expected_codestral_result)
+ end
end
end
@@ -80,10 +98,11 @@
context 'when Fireworks/Qwen beta FF is disabled' do
before do
stub_feature_flags(fireworks_qwen_code_completion: false)
+ stub_feature_flags(disable_code_gecko_default: false)
end
- it 'returns the codegecko model' do
- expect(actual_result).to eq(expected_codegecko_result)
+ it 'returns the default model' do
+ expect(actual_result).to eq(expected_default_result)
end
end
@@ -111,7 +130,14 @@
}
end
- let(:expected_codegecko_result) { {} }
+ let(:expected_codestral_result) do
+ {
+ model_provider: 'vertex-ai',
+ model_name: 'codestral-2501'
+ }
+ end
+
+ let(:expected_default_result) { {} }
let(:expected_self_hosted_model_result) { {} }
end
@@ -125,7 +151,11 @@
CodeSuggestions::Prompts::CodeCompletion::FireworksQwen
end
- let(:expected_codegecko_result) do
+ let(:expected_codestral_result) do
+ CodeSuggestions::Prompts::CodeCompletion::CodestralVertex
+ end
+
+ let(:expected_default_result) do
CodeSuggestions::Prompts::CodeCompletion::VertexAi
end
diff --git a/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb b/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb
index 64d659e91be7e62bf140283108de7badffefa531..7c6c9857c69af2e45b3e530cd30997f58911806f 100644
--- a/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb
+++ b/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb
@@ -128,6 +128,13 @@
)
end
+ let(:request_body_for_vertrex_codestral) do
+ request_body_without_model_details.merge(
+ "model_name" => "codestral-2501",
+ "model_provider" => "vertex-ai"
+ )
+ end
+
context 'on GitLab self-managed' do
before do
allow(Gitlab).to receive(:org_or_com?).and_return(false)
@@ -141,10 +148,22 @@
it_behaves_like 'code suggestion task' do
before do
stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
+ stub_feature_flags(disable_code_gecko_default: false)
end
let(:expected_body) { request_body_without_model_details }
end
+
+ context "when code gecko is disabled" do
+ before do
+ stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
+ stub_feature_flags(disable_code_gecko_default: true)
+ end
+
+ it_behaves_like 'code suggestion task' do
+ let(:expected_body) { request_body_for_vertrex_codestral }
+ end
+ end
end
end
@@ -173,11 +192,22 @@
before do
# opt out for group2
stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: group2)
+ stub_feature_flags(disable_code_gecko_default: false)
end
it_behaves_like 'code suggestion task' do
let(:expected_body) { request_body_without_model_details }
end
+
+ context "when code gecko is disabled" do
+ before do
+ stub_feature_flags(disable_code_gecko_default: true)
+ end
+
+ it_behaves_like 'code suggestion task' do
+ let(:expected_body) { request_body_for_vertrex_codestral }
+ end
+ end
end
end
end
@@ -294,6 +324,7 @@
context 'when amazon_q_chat_and_code_suggestions is disabled' do
before do
stub_feature_flags(amazon_q_chat_and_code_suggestions: false)
+ stub_feature_flags(disable_code_gecko_default: false)
end
it_behaves_like 'code suggestion task' do