From c6fd0791f134ce126e35949176836167226d6a38 Mon Sep 17 00:00:00 2001
From: Lesley Razzaghian <lrazzaghian@gitlab.com>
Date: Mon, 3 Jun 2024 11:48:07 +0000
Subject: [PATCH] Add ability to add experimental feature flag prompt changes
This MR makes it easier for anyone contributing to Duo Chat to make changes under FF.
---
.../beta/prevent_issue_epic_search.yml | 9 +++
ee/app/models/ai/ai_resource/epic.rb | 6 ++
ee/app/models/ai/ai_resource/issue.rb | 6 ++
.../llm/chain/agents/zero_shot/executor.rb | 60 ++++++++++++++++++-
ee/lib/gitlab/llm/chain/gitlab_context.rb | 4 ++
.../llm/chain/tools/epic_reader/executor.rb | 17 ++++++
.../llm/chain/tools/issue_reader/executor.rb | 17 ++++++
ee/lib/gitlab/llm/chain/tools/tool.rb | 10 ++--
.../chain/agents/zero_shot/executor_spec.rb | 50 ++++++++++++++--
.../lib/gitlab/llm/chain/tools/tool_spec.rb | 52 ++++++++++++++--
ee/spec/models/ai/ai_resource/epic_spec.rb | 10 ++++
ee/spec/models/ai/ai_resource/issue_spec.rb | 10 ++++
12 files changed, 236 insertions(+), 15 deletions(-)
create mode 100644 config/feature_flags/beta/prevent_issue_epic_search.yml
diff --git a/config/feature_flags/beta/prevent_issue_epic_search.yml b/config/feature_flags/beta/prevent_issue_epic_search.yml
new file mode 100644
index 0000000000000..d6bdc28ac663c
--- /dev/null
+++ b/config/feature_flags/beta/prevent_issue_epic_search.yml
@@ -0,0 +1,9 @@
+---
+name: prevent_issue_epic_search
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/457756
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/153668
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/463698
+milestone: '17.1'
+group: group::duo chat
+type: beta
+default_enabled: false
diff --git a/ee/app/models/ai/ai_resource/epic.rb b/ee/app/models/ai/ai_resource/epic.rb
index f176749e9f5b4..7eaef1cbabffa 100644
--- a/ee/app/models/ai/ai_resource/epic.rb
+++ b/ee/app/models/ai/ai_resource/epic.rb
@@ -26,6 +26,12 @@ def current_page_short_description
The user is currently on a page that displays an epic with a description, comments, etc., which the user might refer to, for example, as 'current', 'this' or 'that'. The title of the epic is '#{resource.title}'. Remember to use the 'EpicReader' tool if they ask a question about the epic.
SENTENCE
end
+
+ def current_page_experimental_short_description
+ <<~SENTENCE
+ The user is currently on a page that displays an epic with a description, comments, etc., which the user might refer to, for example, as 'current', 'this' or 'that'. The title of the epic is '#{resource.title}'.
+ SENTENCE
+ end
end
end
end
diff --git a/ee/app/models/ai/ai_resource/issue.rb b/ee/app/models/ai/ai_resource/issue.rb
index 77d209be3b70c..abe078ff90b92 100644
--- a/ee/app/models/ai/ai_resource/issue.rb
+++ b/ee/app/models/ai/ai_resource/issue.rb
@@ -26,6 +26,12 @@ def current_page_short_description
The user is currently on a page that displays an issue with a description, comments, etc., which the user might refer to, for example, as 'current', 'this' or 'that'. The title of the issue is '#{resource.title}'. Remember to use the 'IssueReader' tool if they ask a question about the issue.
SENTENCE
end
+
+ def current_page_experimental_short_description
+ <<~SENTENCE
+ The user is currently on a page that displays an issue with a description, comments, etc., which the user might refer to, for example, as 'current', 'this' or 'that'. The title of the issue is '#{resource.title}'.
+ SENTENCE
+ end
end
end
end
diff --git a/ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb b/ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb
index e4eff63854528..4d17025361d63 100644
--- a/ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb
+++ b/ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb
@@ -102,7 +102,7 @@ def options
@options ||= {
tool_names: tools.map { |tool_class| tool_class::Executor::NAME }.join(', '),
tools_definitions: tools.map do |tool_class|
- tool_class::Executor.full_definition
+ tool_class::Executor.full_definition(use_experimental_prompt: use_experimental_prompt?)
end.join("\n"),
user_input: user_input,
agent_scratchpad: +"",
@@ -154,7 +154,11 @@ def prompt_version
end
def zero_shot_prompt
- ZERO_SHOT_PROMPT
+ use_experimental_prompt? ? ZERO_SHOT_EXPERIMENTAL_PROMPT : ZERO_SHOT_PROMPT
+ end
+
+ def use_experimental_prompt?
+ Feature.enabled?(:prevent_issue_epic_search, context.current_user)
end
def last_conversation
@@ -194,7 +198,11 @@ def prompt_options
end
def current_resource
- context.current_page_short_description
+ if use_experimental_prompt?
+ context.current_page_experimental_short_description
+ else
+ context.current_page_short_description
+ end
rescue ArgumentError
""
end
@@ -255,6 +263,52 @@ def source_template
Begin!
PROMPT
+ ZERO_SHOT_EXPERIMENTAL_PROMPT = <<~PROMPT.freeze
+ Answer the question as accurate as you can.
+
+ You have access only to the following tools:
+ <tool_list>
+ %<tools_definitions>s
+ </tool_list>
+ Consider every tool before making a decision.
+ Ensure that your answer is accurate and contain only information directly supported by the information retrieved using provided tools.
+
+ When you can answer the question directly you must use this response format:
+ Thought: you should always think about how to answer the question
+ Action: DirectAnswer
+ Final Answer: the final answer to the original input question if you have a direct answer to the user's question.
+
+ You must always use the following format when using a tool:
+ Question: the input question you must answer
+ Thought: you should always think about what to do
+ Action: the action to take, should be one tool from this list: [%<tool_names>s]
+ Action Input: the input to the action needs to be provided for every action that uses a tool.
+ Observation: the result of the tool actions. But remember that you're still #{AGENT_NAME}.
+
+
+ ... (this Thought/Action/Action Input/Observation sequence can repeat N times)
+
+ Thought: I know the final answer.
+ Final Answer: the final answer to the original input question.
+
+ When concluding your response, provide the final answer as "Final Answer:". It should contain everything that user needs to see, including answer from "Observation" section.
+ %<current_code>s
+
+ You have access to the following GitLab resources: %<resources>s.
+ You also have access to all information that can be helpful to someone working in software development of any kind.
+ At the moment, you do not have access to the following GitLab resources: Merge Requests, Pipelines, Vulnerabilities.
+ At the moment, you do not have the ability to search Issues or Epics based on a description or keywords. You can only read information about a specific issue/epic IF the user is on the specific issue/epic's page, or provides a URL or ID.
+ Do not use the IssueReader or EpicReader tool if you do not have these specified identifiers.
+
+ %<source_template>s
+
+ Ask user to leave feedback.
+
+ %<current_resource>s
+
+ Begin!
+ PROMPT
+
PROMPT_TEMPLATE = [
Utils::Prompt.as_system(ZERO_SHOT_PROMPT),
Utils::Prompt.as_user("Question: %<user_input>s"),
diff --git a/ee/lib/gitlab/llm/chain/gitlab_context.rb b/ee/lib/gitlab/llm/chain/gitlab_context.rb
index e5970f0e81efb..241aa468ef180 100644
--- a/ee/lib/gitlab/llm/chain/gitlab_context.rb
+++ b/ee/lib/gitlab/llm/chain/gitlab_context.rb
@@ -30,6 +30,10 @@ def current_page_short_description
authorized_resource&.current_page_short_description
end
+ def current_page_experimental_short_description
+ authorized_resource&.current_page_experimental_short_description
+ end
+
def resource_serialized(content_limit:)
return '' unless authorized_resource
diff --git a/ee/lib/gitlab/llm/chain/tools/epic_reader/executor.rb b/ee/lib/gitlab/llm/chain/tools/epic_reader/executor.rb
index 535d7ee945f0c..5e6696561f140 100644
--- a/ee/lib/gitlab/llm/chain/tools/epic_reader/executor.rb
+++ b/ee/lib/gitlab/llm/chain/tools/epic_reader/executor.rb
@@ -16,6 +16,23 @@ class Executor < Identifier
'high-level plans and discussions. Epic can contain multiple issues. ' \
'Action Input for this tool should be the original question or epic identifier.'
+ EXPERIMENTAL_TOOL_DESCRIPTION = <<~PROMPT
+ This tool retrieves the content of a specific epic
+ ONLY if the user question fulfills the strict usage conditions below.
+
+ **Strict Usage Conditions:**
+ * **Condition 1: epic ID Provided:** This tool MUST be used ONLY when the user provides a valid epic ID.
+ * **Condition 2: epic URL Context:** This tool MUST be used ONLY when the user is actively viewing a specific epic URL or a specific URL is provided by the user.
+
+ **Do NOT** attempt to search for or identify epics based on descriptions, keywords, or user questions.
+
+ **Action Input:**
+ * The original question asked by the user.
+
+ **Important:** Reject any input that does not strictly adhere to the usage conditions above.
+ Return a message stating you are unable to search for epics without a valid identifier.
+ PROMPT
+
EXAMPLE =
<<~PROMPT
Question: Please identify the author of &123 epic.
diff --git a/ee/lib/gitlab/llm/chain/tools/issue_reader/executor.rb b/ee/lib/gitlab/llm/chain/tools/issue_reader/executor.rb
index c3a706526b56b..5b63630dcfec0 100644
--- a/ee/lib/gitlab/llm/chain/tools/issue_reader/executor.rb
+++ b/ee/lib/gitlab/llm/chain/tools/issue_reader/executor.rb
@@ -17,6 +17,23 @@ class Executor < Identifier
'collaboration, discussions, planning and tracking of work.' \
'Action Input for this tool should be the original question or issue identifier.'
+ EXPERIMENTAL_TOOL_DESCRIPTION = <<~PROMPT
+ This tool retrieves the content of a specific issue
+ ONLY if the user question fulfills the strict usage conditions below.
+
+ **Strict Usage Conditions:**
+ * **Condition 1: Issue ID Provided:** This tool MUST be used ONLY when the user provides a valid issue ID.
+ * **Condition 2: Issue URL Context:** This tool MUST be used ONLY when the user is actively viewing a specific issue URL or a specific URL is provided by the user.
+
+ **Do NOT** attempt to search for or identify issues based on descriptions, keywords, or user questions.
+
+ **Action Input:**
+ * The original question asked by the user.
+
+ **Important:** Reject any input that does not strictly adhere to the usage conditions above.
+ Return a message stating you are unable to search for issues without a valid identifier.
+ PROMPT
+
EXAMPLE =
<<~PROMPT
Question: Please identify the author of #123 issue
diff --git a/ee/lib/gitlab/llm/chain/tools/tool.rb b/ee/lib/gitlab/llm/chain/tools/tool.rb
index 4e23d40ddd4f4..3d331a660b69c 100644
--- a/ee/lib/gitlab/llm/chain/tools/tool.rb
+++ b/ee/lib/gitlab/llm/chain/tools/tool.rb
@@ -16,12 +16,12 @@ class Tool
delegate :resource, :resource=, to: :context
- def self.full_definition
+ def self.full_definition(use_experimental_prompt: false)
[
"<tool>",
"<tool_name>#{self::NAME}</tool_name>",
"<description>",
- description,
+ description(use_experimental_prompt),
"</description>",
"<example>",
self::EXAMPLE,
@@ -85,8 +85,10 @@ def group_from_context
attr_reader :logger, :stream_response_handler
- def self.description
- self::DESCRIPTION
+ def self.description(use_experimental_prompt)
+ experiment = use_experimental_prompt && defined?(self::EXPERIMENTAL_TOOL_DESCRIPTION)
+
+ experiment ? self::EXPERIMENTAL_TOOL_DESCRIPTION : self::DESCRIPTION
end
def not_found
diff --git a/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/executor_spec.rb b/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/executor_spec.rb
index 54b5cae44164f..84f762595691e 100644
--- a/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/executor_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/executor_spec.rb
@@ -233,11 +233,36 @@
agent.prompt
end
- it 'includes prompt in the options' do
- expect(Gitlab::Llm::Chain::Agents::ZeroShot::Prompts::Anthropic)
- .to receive(:prompt).once.with(a_hash_including(prompt_options))
+ context 'with the `prevent_issue_epic_search` feature flag' do
+ before do
+ stub_const("#{described_class.name}::ZERO_SHOT_EXPERIMENTAL_PROMPT", 'I am an experimental prompt.')
+ end
- agent.prompt
+ context 'when experimental feature flag is enabled' do
+ let(:prompt_options) { { zero_shot_prompt: described_class::ZERO_SHOT_EXPERIMENTAL_PROMPT } }
+
+ it 'includes the experimental prompt in the prompt options' do
+ expect(Gitlab::Llm::Chain::Agents::ZeroShot::Prompts::Anthropic)
+ .to receive(:prompt).once.with(a_hash_including(prompt_options))
+
+ agent.prompt
+ end
+ end
+
+ context 'when experimental feature flag is not enabled' do
+ let(:prompt_options) { { zero_shot_prompt: described_class::ZERO_SHOT_PROMPT } }
+
+ before do
+ stub_feature_flags(prevent_issue_epic_search: false)
+ end
+
+ it 'includes the default prompt options' do
+ expect(Gitlab::Llm::Chain::Agents::ZeroShot::Prompts::Anthropic)
+ .to receive(:prompt).once.with(a_hash_including(prompt_options))
+
+ agent.prompt
+ end
+ end
end
context 'when duo_chat_display_source feature flag is enabled' do
@@ -304,6 +329,10 @@
XML
end
+ before do
+ stub_feature_flags(prevent_issue_epic_search: false)
+ end
+
let(:prompt_resource) do
<<~CONTEXT
<resource>
@@ -313,6 +342,7 @@
end
let(:short_description) { 'short description' }
+ let(:experimental_short_description) { 'experimental short description' }
it 'does not include the current resource metadata' do
expect(context).not_to receive(:resource_serialized)
@@ -323,6 +353,18 @@
expect(context).to receive(:current_page_short_description).and_return(short_description)
expect(system_prompt(agent)).to include(short_description)
end
+
+ context 'when the `prevent_issue_epic_search` is enabled' do
+ before do
+ stub_feature_flags(prevent_issue_epic_search: true)
+ end
+
+ it 'returns experimental short description' do
+ expect(context).to receive(:current_page_experimental_short_description)
+ .and_return(experimental_short_description)
+ expect(system_prompt(agent)).to include(experimental_short_description)
+ end
+ end
end
context 'when the resource is an issue' do
diff --git a/ee/spec/lib/gitlab/llm/chain/tools/tool_spec.rb b/ee/spec/lib/gitlab/llm/chain/tools/tool_spec.rb
index e9fbaf3bce286..18c30150959be 100644
--- a/ee/spec/lib/gitlab/llm/chain/tools/tool_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/tools/tool_spec.rb
@@ -136,7 +136,21 @@
XML
end
- let(:expected_description) { 'TEST' }
+ let(:experimental_definition) do
+ <<~XML.chomp
+ <tool>
+ <tool_name>TEST_TOOL</tool_name>
+ <description>
+ #{experimental_description}
+ </description>
+ <example>
+ EXAMPLE
+ </example>
+ </tool>
+ XML
+ end
+
+ let(:expected_description) { 'No feature flag description' }
before do
stub_const("#{described_class.name}::NAME", 'TEST_TOOL')
@@ -144,9 +158,39 @@
stub_const("#{described_class.name}::EXAMPLE", 'EXAMPLE')
end
- context 'when description is defined' do
- it 'returns detailed description of the tool' do
- expect(described_class.full_definition).to eq(definition)
+ context 'when experimental description constant is not defined' do
+ context 'when experimental prompt is enabled' do
+ it 'returns default description of the tool' do
+ expect(described_class.full_definition(use_experimental_prompt: true)).to eq(definition)
+ end
+ end
+
+ context 'when experimental prompt is not enabled' do
+ stub_feature_flags(prevent_issue_epic_search: false)
+
+ it 'returns default description of the tool' do
+ expect(described_class.full_definition).to eq(definition)
+ end
+ end
+ end
+
+ context 'when experimental description constant is defined' do
+ let(:experimental_description) { 'Experimental description' }
+
+ before do
+ stub_const("#{described_class.name}::EXPERIMENTAL_TOOL_DESCRIPTION", experimental_description)
+ end
+
+ context 'when experimental prompt is enabled' do
+ it 'returns experimental description of the tool' do
+ expect(described_class.full_definition(use_experimental_prompt: true)).to eq(experimental_definition)
+ end
+ end
+
+ context 'when experimental prompt is not enabled' do
+ it 'returns default description of the tool' do
+ expect(described_class.full_definition).to eq(definition)
+ end
end
end
end
diff --git a/ee/spec/models/ai/ai_resource/epic_spec.rb b/ee/spec/models/ai/ai_resource/epic_spec.rb
index 79e29df095656..7883b62cee0a8 100644
--- a/ee/spec/models/ai/ai_resource/epic_spec.rb
+++ b/ee/spec/models/ai/ai_resource/epic_spec.rb
@@ -32,6 +32,16 @@
describe '#current_page_short_description' do
it 'returns prompt' do
expect(wrapped_epic.current_page_short_description).to include("The title of the epic is '#{epic.title}'.")
+ expect(wrapped_epic.current_page_short_description).to include("Remember to use the 'EpicReader' tool")
+ end
+ end
+
+ describe '#current_page_experimental_short_description' do
+ it 'returns experimental short description' do
+ expect(wrapped_epic.current_page_experimental_short_description)
+ .to include("The title of the epic is '#{epic.title}'.")
+ expect(wrapped_epic.current_page_experimental_short_description)
+ .not_to include("Remember to use the 'EpicReader' tool")
end
end
end
diff --git a/ee/spec/models/ai/ai_resource/issue_spec.rb b/ee/spec/models/ai/ai_resource/issue_spec.rb
index 5e3b6da73b0cb..d9861b9fc8254 100644
--- a/ee/spec/models/ai/ai_resource/issue_spec.rb
+++ b/ee/spec/models/ai/ai_resource/issue_spec.rb
@@ -31,6 +31,16 @@
describe '#current_page_short_description' do
it 'returns prompt' do
expect(wrapped_issue.current_page_short_description).to include("The title of the issue is '#{issue.title}'.")
+ expect(wrapped_issue.current_page_short_description).to include("Remember to use the 'IssueReader' tool")
+ end
+ end
+
+ describe '#current_page_experimental_short_description' do
+ it 'returns experimental short description' do
+ expect(wrapped_issue.current_page_experimental_short_description)
+ .to include("The title of the issue is '#{issue.title}'.")
+ expect(wrapped_issue.current_page_experimental_short_description)
+ .not_to include("Remember to use the 'IssueReader' tool")
end
end
end
--
GitLab