From bc9b79df103e6decbee0daa9328f2149b865ab9b Mon Sep 17 00:00:00 2001
From: Tetiana Chupryna <tchupryna@gitlab.com>
Date: Wed, 5 Feb 2025 11:09:53 +0000
Subject: [PATCH] Send current page params as set of params
We are going to deprecate current page prompts and move them
to AI Gateway
---
.../models/ai/ai_resource/base_ai_resource.rb | 11 +++
ee/app/models/ai/ai_resource/ci/build.rb | 6 ++
.../current_page_context_prompt_in_aigw.yml | 9 +++
ee/lib/gitlab/duo/chat/react_executor.rb | 8 +++
ee/lib/gitlab/llm/chain/gitlab_context.rb | 2 +-
.../gitlab/duo/chat/react_executor_spec.rb | 67 +++++++++++++++++--
.../ai/ai_resource/base_ai_resource_spec.rb | 7 ++
.../models/ai/ai_resource/ci/build_spec.rb | 6 ++
8 files changed, 110 insertions(+), 6 deletions(-)
create mode 100644 ee/config/feature_flags/wip/current_page_context_prompt_in_aigw.yml
diff --git a/ee/app/models/ai/ai_resource/base_ai_resource.rb b/ee/app/models/ai/ai_resource/base_ai_resource.rb
index 6935c700493b3..09baa3a911d9d 100644
--- a/ee/app/models/ai/ai_resource/base_ai_resource.rb
+++ b/ee/app/models/ai/ai_resource/base_ai_resource.rb
@@ -13,6 +13,17 @@ def initialize(user, resource)
def serialize_for_ai(_content_limit:)
raise NotImplementedError
end
+
+ def current_page_params
+ {
+ type: current_page_type,
+ title: resource.title
+ }
+ end
+
+ def current_page_type
+ raise NotImplementedError
+ end
end
end
end
diff --git a/ee/app/models/ai/ai_resource/ci/build.rb b/ee/app/models/ai/ai_resource/ci/build.rb
index c4006103c59ab..2f61cc223411e 100644
--- a/ee/app/models/ai/ai_resource/ci/build.rb
+++ b/ee/app/models/ai/ai_resource/ci/build.rb
@@ -25,6 +25,12 @@ def current_page_short_description
The user is currently on a page that displays a ci build which the user might refer to, for example, as 'current', 'this' or 'that'.
SENTENCE
end
+
+ def current_page_params
+ {
+ type: current_page_type
+ }
+ end
end
end
end
diff --git a/ee/config/feature_flags/wip/current_page_context_prompt_in_aigw.yml b/ee/config/feature_flags/wip/current_page_context_prompt_in_aigw.yml
new file mode 100644
index 0000000000000..8929ee377ed47
--- /dev/null
+++ b/ee/config/feature_flags/wip/current_page_context_prompt_in_aigw.yml
@@ -0,0 +1,9 @@
+---
+name: current_page_context_prompt_in_aigw
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/508317
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/177622
+rollout_issue_url:
+milestone: '17.9'
+group: group::duo chat
+type: wip
+default_enabled: false
diff --git a/ee/lib/gitlab/duo/chat/react_executor.rb b/ee/lib/gitlab/duo/chat/react_executor.rb
index 446c95c038c8a..bb309c38ca76d 100644
--- a/ee/lib/gitlab/duo/chat/react_executor.rb
+++ b/ee/lib/gitlab/duo/chat/react_executor.rb
@@ -313,6 +313,7 @@ def conversation
end
def current_resource_params
+ return current_page_params if Feature.enabled?(:current_page_context_prompt_in_aigw, context.current_user)
return unless current_resource_type
{
@@ -335,6 +336,13 @@ def current_resource_content
end
strong_memoize_attr :current_resource_content
+ def current_page_params
+ context.current_page_params
+ rescue ArgumentError
+ nil
+ end
+ strong_memoize_attr :current_page_params
+
def current_file_params
return unless current_selection || current_blob
diff --git a/ee/lib/gitlab/llm/chain/gitlab_context.rb b/ee/lib/gitlab/llm/chain/gitlab_context.rb
index e42a2ccaf4adf..9d20aad78648a 100644
--- a/ee/lib/gitlab/llm/chain/gitlab_context.rb
+++ b/ee/lib/gitlab/llm/chain/gitlab_context.rb
@@ -13,7 +13,7 @@ class GitlabContext
attr_reader :project
- delegate :current_page_type, :current_page_short_description,
+ delegate :current_page_type, :current_page_short_description, :current_page_params,
to: :authorized_resource, allow_nil: true
# rubocop:disable Metrics/ParameterLists -- we probably need to rethink this initializer
diff --git a/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb b/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
index 6876fa755d64e..37388dd499d08 100644
--- a/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
+++ b/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
@@ -53,6 +53,8 @@
end
let(:issue_resource) { Ai::AiResource::Issue.new(user, resource) }
+ let(:issue_page_params) { { type: issue_resource.current_page_type, title: resource.title } }
+
let(:answer_chunk) { create(:final_answer_chunk, chunk: "Ans") }
let(:step_params) do
@@ -60,10 +62,7 @@
messages: [{
role: "user",
content: user_input,
- context: {
- type: issue_resource.current_page_type,
- content: issue_resource.current_page_short_description
- },
+ context: issue_page_params,
current_file: nil,
additional_context: context.additional_context
}],
@@ -367,7 +366,7 @@ def expect_sli_error(failed)
context "when error event received and it's prompt length error" do
let(:message) do
<<~MESSAGE
- Error code: 400 - {'type': 'error', 'error': {'type': 'invalid_request_error', 'message': 'prompt is too long: 200082 tokens > 199999 maximum'}}
+ Error code: 400 - {'type': 'error', 'error': {'type': 'invalid_request_error', 'message': 'prompt is too long: 200082 tokens > 199999 maximum'}}
MESSAGE
end
@@ -435,6 +434,29 @@ def expect_sli_error(failed)
end
end
+ context "when there is no resource" do
+ let(:context) do
+ Gitlab::Llm::Chain::GitlabContext.new(
+ current_user: user,
+ container: nil,
+ resource: nil,
+ ai_request: nil
+ )
+ end
+
+ it "sends request without context" do
+ params = step_params
+ params[:messages].first[:context] = nil
+
+ expect_next_instance_of(Gitlab::Duo::Chat::StepExecutor) do |react_agent|
+ expect(react_agent).to receive(:step).with(hash_including(params))
+ .and_yield(action_event).and_return([action_event])
+ end
+
+ agent.execute
+ end
+ end
+
context "when code is selected" do
let(:selected_text) { 'code selection' }
let(:current_file) do
@@ -485,6 +507,41 @@ def expect_sli_error(failed)
end
end
+ context "when current page is included in context" do
+ it "pass current page params" do
+ params = step_params
+ params[:messages].first[:context] = issue_page_params
+
+ expect_next_instance_of(Gitlab::Duo::Chat::StepExecutor) do |react_agent|
+ expect(react_agent).to receive(:step).with(params)
+ .and_yield(action_event).and_return([action_event])
+ end
+
+ agent.execute
+ end
+
+ context "with feature flag disabled" do
+ before do
+ stub_feature_flags(current_page_context_prompt_in_aigw: false)
+ end
+
+ it "pass current page description" do
+ params = step_params
+ params[:messages].first[:context] = {
+ type: issue_resource.current_page_type,
+ content: issue_resource.current_page_short_description
+ }
+
+ expect_next_instance_of(Gitlab::Duo::Chat::StepExecutor) do |react_agent|
+ expect(react_agent).to receive(:step).with(params)
+ .and_yield(action_event).and_return([action_event])
+ end
+
+ agent.execute
+ end
+ end
+ end
+
context 'when Duo chat is self-hosted' do
let_it_be(:self_hosted_model) { create(:ai_self_hosted_model, api_token: 'test_token') }
let_it_be(:ai_feature) { create(:ai_feature_setting, self_hosted_model: self_hosted_model, feature: :duo_chat) }
diff --git a/ee/spec/models/ai/ai_resource/base_ai_resource_spec.rb b/ee/spec/models/ai/ai_resource/base_ai_resource_spec.rb
index c140ce3d0e620..fa2fb3029212a 100644
--- a/ee/spec/models/ai/ai_resource/base_ai_resource_spec.rb
+++ b/ee/spec/models/ai/ai_resource/base_ai_resource_spec.rb
@@ -9,4 +9,11 @@
.to raise_error(NotImplementedError)
end
end
+
+ describe '#current_page_params' do
+ it 'returns params to construct prompt' do
+ expect { described_class.new(nil, nil).current_page_params }
+ .to raise_error(NotImplementedError)
+ end
+ end
end
diff --git a/ee/spec/models/ai/ai_resource/ci/build_spec.rb b/ee/spec/models/ai/ai_resource/ci/build_spec.rb
index eccfe2219d54a..ab4e9c520baef 100644
--- a/ee/spec/models/ai/ai_resource/ci/build_spec.rb
+++ b/ee/spec/models/ai/ai_resource/ci/build_spec.rb
@@ -37,4 +37,10 @@
.to include("The user is currently on a page that displays a ci build")
end
end
+
+ describe '#current_page_params' do
+ it 'returns params to construct prompt' do
+ expect(wrapped_build.current_page_params.keys).to eq([:type])
+ end
+ end
end
--
GitLab