From 90d1e8d2dbe3ec7a8c455d77cac9d09ba16f5329 Mon Sep 17 00:00:00 2001
From: Igor Drozdov <idrozdov@gitlab.com>
Date: Fri, 27 Sep 2024 06:45:27 +0200
Subject: [PATCH] Migrate FixCode tool
---
.../experiment/prompt_migration_fix_code.yml | 9 +++++++++
ee/lib/gitlab/llm/chain/tools/fix_code/executor.rb | 1 +
.../lib/gitlab/llm/chain/tools/fix_code/executor_spec.rb | 5 +++++
3 files changed, 15 insertions(+)
create mode 100644 config/feature_flags/experiment/prompt_migration_fix_code.yml
diff --git a/config/feature_flags/experiment/prompt_migration_fix_code.yml b/config/feature_flags/experiment/prompt_migration_fix_code.yml
new file mode 100644
index 000000000000..e67ed9536887
--- /dev/null
+++ b/config/feature_flags/experiment/prompt_migration_fix_code.yml
@@ -0,0 +1,9 @@
+---
+name: prompt_migration_fix_code
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/475047
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/167399
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/496750
+milestone: '17.5'
+group: group::custom models
+type: experiment
+default_enabled: false
diff --git a/ee/lib/gitlab/llm/chain/tools/fix_code/executor.rb b/ee/lib/gitlab/llm/chain/tools/fix_code/executor.rb
index a83bea52aacb..cf343e671f7f 100644
--- a/ee/lib/gitlab/llm/chain/tools/fix_code/executor.rb
+++ b/ee/lib/gitlab/llm/chain/tools/fix_code/executor.rb
@@ -7,6 +7,7 @@ module Tools
module FixCode
class Executor < SlashCommandTool
extend ::Gitlab::Utils::Override
+ prepend Concerns::UseAiGatewayAgentPrompt
NAME = 'FixCode'
HUMAN_NAME = 'Fix Code'
diff --git a/ee/spec/lib/gitlab/llm/chain/tools/fix_code/executor_spec.rb b/ee/spec/lib/gitlab/llm/chain/tools/fix_code/executor_spec.rb
index 3a50aa9a524f..ddc540e66b48 100644
--- a/ee/spec/lib/gitlab/llm/chain/tools/fix_code/executor_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/tools/fix_code/executor_spec.rb
@@ -133,6 +133,11 @@
expect(answer.error_code).to eq("M4000")
end
end
+
+ it_behaves_like 'uses ai gateway agent prompt' do
+ let(:prompt_class) { Gitlab::Llm::Chain::Tools::FixCode::Prompts::Anthropic }
+ let(:unit_primitive) { 'fix_code' }
+ end
end
context 'when context is not authorized' do
--
GitLab