diff --git a/ee/app/models/ai/ai_resource/epic.rb b/ee/app/models/ai/ai_resource/epic.rb index dd5ab3b7b08e0eaeacfbef8aecb9cfe1e24af1fb..89153e894931bfdbcf97b6a1db7fcf0095672336 100644 --- a/ee/app/models/ai/ai_resource/epic.rb +++ b/ee/app/models/ai/ai_resource/epic.rb @@ -15,6 +15,10 @@ def serialize_for_ai(user:, content_limit:) }) end + def current_page_type + "epic" + end + def current_page_sentence <<~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 data is provided in tags, and if it is sufficient in answering the question, utilize it instead of using the 'EpicReader' tool. diff --git a/ee/app/models/ai/ai_resource/issue.rb b/ee/app/models/ai/ai_resource/issue.rb index f91380d5635ee28b6f3fdf57bced2b860397521c..0c653175ffb2f5141948416a43ee0a49a02f9029 100644 --- a/ee/app/models/ai/ai_resource/issue.rb +++ b/ee/app/models/ai/ai_resource/issue.rb @@ -15,6 +15,10 @@ def serialize_for_ai(user:, content_limit:) }) end + def current_page_type + "issue" + end + def current_page_sentence <<~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 data is provided in tags, and if it is sufficient in answering the question, utilize it instead of using the 'IssueReader' tool. diff --git a/ee/config/feature_flags/wip/v2_chat_agent_integration.yml b/ee/config/feature_flags/wip/v2_chat_agent_integration.yml new file mode 100644 index 0000000000000000000000000000000000000000..3d2c99a82ea5690fc479a2e479a70010b471036b --- /dev/null +++ b/ee/config/feature_flags/wip/v2_chat_agent_integration.yml @@ -0,0 +1,9 @@ +--- +name: v2_chat_agent_integration +feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/456258 +introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/150529 +rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/466910 +milestone: '17.2' +group: group::duo chat +type: wip +default_enabled: false diff --git a/ee/lib/gitlab/llm/chain/agents/single_action_executor.rb b/ee/lib/gitlab/llm/chain/agents/single_action_executor.rb new file mode 100644 index 0000000000000000000000000000000000000000..eb4e0b1a8d84a89e672b6506e5e8d2923355adf4 --- /dev/null +++ b/ee/lib/gitlab/llm/chain/agents/single_action_executor.rb @@ -0,0 +1,172 @@ +# frozen_string_literal: true + +module Gitlab + module Llm + module Chain + module Agents + class SingleActionExecutor + include Gitlab::Utils::StrongMemoize + include Concerns::AiDependent + include Langsmith::RunHelpers + + attr_reader :tools, :user_input, :context, :response_handler + attr_accessor :iterations + + MAX_ITERATIONS = 10 + RESPONSE_TYPE_TOOL = 'tool' + + # @param [String] user_input - a question from a user + # @param [Array] tools - an array of Tools defined in the tools module. + # @param [GitlabContext] context - Gitlab context containing useful context information + # @param [ResponseService] response_handler - Handles returning the response to the client + # @param [ResponseService] stream_response_handler - Handles streaming chunks to the client + def initialize(user_input:, tools:, context:, response_handler:, stream_response_handler: nil) + @user_input = user_input + @tools = tools + @context = context + @iterations = 0 + @logger = Gitlab::Llm::Logger.build + @response_handler = response_handler + @stream_response_handler = stream_response_handler + end + + def execute + @agent_scratchpad = [] + MAX_ITERATIONS.times do + step = {} + thoughts = execute_streamed_request + + answer = Answer.from_response( + response_body: thoughts, + tools: tools, + context: context, + parser_klass: Parsers::SingleActionParser + ) + + return answer if answer.is_final? + + step[:thought] = answer.suggestions + step[:tool] = answer.tool + step[:tool_input] = user_input + + tool_class = answer.tool + + picked_tool_action(tool_class) + + tool = tool_class.new( + context: context, + options: { + input: user_input, + suggestions: answer.suggestions + }, + stream_response_handler: stream_response_handler + ) + + tool_answer = tool.execute + + return tool_answer if tool_answer.is_final? + + step[:observation] = tool_answer.content.strip + @agent_scratchpad.push(step) + end + + Answer.default_final_answer(context: context) + rescue Net::ReadTimeout => error + Gitlab::ErrorTracking.track_exception(error) + Answer.error_answer( + context: context, + content: _("I'm sorry, I couldn't respond in time. Please try again."), + error_code: "A1000" + ) + rescue Gitlab::Llm::AiGateway::Client::ConnectionError => error + Gitlab::ErrorTracking.track_exception(error) + Answer.error_answer( + context: context, + error_code: "A1001" + ) + end + traceable :execute, name: 'Run ReAct' + + private + + def streamed_content(_content, chunk) + chunk[:content] + end + + def execute_streamed_request + request(&streamed_request_handler(Answers::StreamedJson.new)) + end + + attr_reader :logger, :stream_response_handler + + # This method should not be memoized because the input variables change over time + def prompt + { prompt: user_input, options: prompt_options } + end + + def prompt_options + @options = { + agent_scratchpad: @agent_scratchpad, + conversation: conversation, + current_resource_type: current_resource_type, + current_resource_content: current_resource_content, + single_action_agent: true + } + end + + def picked_tool_action(tool_class) + logger.info(message: "Picked tool", tool: tool_class.to_s) + + stream_response_handler.execute( + response: Gitlab::Llm::Chain::ToolResponseModifier.new(tool_class), + options: { + role: ::Gitlab::Llm::ChatMessage::ROLE_SYSTEM, + type: RESPONSE_TYPE_TOOL + } + ) + end + + # agent_version is deprecated, Chat conversation doesn't have this param anymore + def last_conversation + ChatStorage.new(context.current_user, nil).last_conversation + end + strong_memoize_attr :last_conversation + + def conversation + # include only messages with successful response and reorder + # messages so each question is followed by its answer + by_request = last_conversation + .reject { |message| message.errors.present? } + .group_by(&:request_id) + .select { |_uuid, messages| messages.size > 1 } + + c = by_request.values.sort_by { |messages| messages.first.timestamp }.flatten + + return [] if c.blank? + + c = c.last(50).map do |message, _| + { role: message.role.to_sym, content: message.content } + end + + c.to_s + end + strong_memoize_attr :conversation + + # TODO: remove issue condition when next issue is implemented + # https://gitlab.com/gitlab-org/gitlab/-/issues/468905 + def current_resource_type + context.current_page_type + rescue ArgumentError + nil + end + + def current_resource_content + context.current_page_short_description + rescue ArgumentError + nil + end + end + 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 cca37eec78e919b4a78176c676a6478beddb14a9..a1a707dded293dc0fafb431eed200babbec24948 100644 --- a/ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb +++ b/ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb @@ -1,5 +1,8 @@ # frozen_string_literal: true +# Deprecation: this executor will be removed in favor of SingleActionExecutor +# see https://gitlab.com/gitlab-org/gitlab/-/issues/469087 + module Gitlab module Llm module Chain diff --git a/ee/lib/gitlab/llm/chain/answer.rb b/ee/lib/gitlab/llm/chain/answer.rb index 10c53e3a6e4ad1d539466eba0e6d986c85700266..e1ee0d8694911921ba7998a71641059ef0fdcb7d 100644 --- a/ee/lib/gitlab/llm/chain/answer.rb +++ b/ee/lib/gitlab/llm/chain/answer.rb @@ -9,8 +9,8 @@ class Answer attr_accessor :status, :content, :context, :tool, :suggestions, :is_final, :extras, :error_code alias_method :is_final?, :is_final - def self.from_response(response_body:, tools:, context:) - parser = Parsers::ChainOfThoughtParser.new(output: response_body) + def self.from_response(response_body:, tools:, context:, parser_klass: Parsers::ChainOfThoughtParser) + parser = parser_klass.new(output: response_body) parser.parse return final_answer(context: context, content: parser.final_answer) if parser.final_answer diff --git a/ee/lib/gitlab/llm/chain/answers/streamed_json.rb b/ee/lib/gitlab/llm/chain/answers/streamed_json.rb new file mode 100644 index 0000000000000000000000000000000000000000..36d5006b5af9b593c044050b509b70efdf87e177 --- /dev/null +++ b/ee/lib/gitlab/llm/chain/answers/streamed_json.rb @@ -0,0 +1,47 @@ +# frozen_string_literal: true + +module Gitlab + module Llm + module Chain + module Answers + class StreamedJson < StreamedAnswer + def initialize + @final_answer_started = false + @full_message = '' + + super + end + + def next_chunk(content) + return if content.empty? + + content = ::Gitlab::Json.parse(content) + answer_chunk = final_answer_chunk(content) + + return unless answer_chunk + return payload(answer_chunk) if final_answer_started + + @full_message += answer_chunk + + return unless final_answer_start(content) + + @final_answer_started = true + payload(answer_chunk) + end + + private + + attr_accessor :full_message, :final_answer_started + + def final_answer_start(content) + 'final_answer_delta' == content['type'] + end + + def final_answer_chunk(content) + content.dig('data', 'text') + end + end + end + end + end +end diff --git a/ee/lib/gitlab/llm/chain/concerns/ai_dependent.rb b/ee/lib/gitlab/llm/chain/concerns/ai_dependent.rb index a645bc9a1fcc6647c8abe3948dfdaf25494b04c6..92a436ce9782b4130e2676dd4d815b06f79fcaa6 100644 --- a/ee/lib/gitlab/llm/chain/concerns/ai_dependent.rb +++ b/ee/lib/gitlab/llm/chain/concerns/ai_dependent.rb @@ -34,7 +34,8 @@ def streamed_request_handler(streamed_answer) if chunk stream_response_handler.execute( - response: Gitlab::Llm::Chain::StreamedResponseModifier.new(content, chunk_id: chunk[:id]), + response: Gitlab::Llm::Chain::StreamedResponseModifier + .new(streamed_content(content, chunk), chunk_id: chunk[:id]), options: { chunk_id: chunk[:id] } ) end @@ -56,6 +57,11 @@ def provider_prompt_class def unit_primitive nil end + + # This method is modified in SingleActionExecutor for Duo Chat + def streamed_content(content, _chunk) + content + end end end end diff --git a/ee/lib/gitlab/llm/chain/gitlab_context.rb b/ee/lib/gitlab/llm/chain/gitlab_context.rb index e5970f0e81efb9273f074075be51ec118f5fabf3..9a9cfb94da94dedd35a20aa0e0fbaebe3429e909 100644 --- a/ee/lib/gitlab/llm/chain/gitlab_context.rb +++ b/ee/lib/gitlab/llm/chain/gitlab_context.rb @@ -7,6 +7,9 @@ class GitlabContext attr_accessor :current_user, :container, :resource, :ai_request, :tools_used, :extra_resource, :request_id, :current_file, :agent_version + delegate :current_page_type, :current_page_sentence, :current_page_short_description, + to: :authorized_resource, allow_nil: true + def initialize( current_user:, container:, resource:, ai_request:, extra_resource: {}, request_id: nil, current_file: {}, agent_version: nil @@ -22,14 +25,6 @@ def initialize( @agent_version = agent_version end - def current_page_sentence - authorized_resource&.current_page_sentence - end - - def current_page_short_description - authorized_resource&.current_page_short_description - end - def resource_serialized(content_limit:) return '' unless authorized_resource diff --git a/ee/lib/gitlab/llm/chain/parsers/single_action_parser.rb b/ee/lib/gitlab/llm/chain/parsers/single_action_parser.rb new file mode 100644 index 0000000000000000000000000000000000000000..d89e3900488cfca08b8f3f57ac934fd7488c9644 --- /dev/null +++ b/ee/lib/gitlab/llm/chain/parsers/single_action_parser.rb @@ -0,0 +1,60 @@ +# frozen_string_literal: true + +module Gitlab + module Llm + module Chain + module Parsers + class SingleActionParser < OutputParser + attr_reader :action, :action_input, :thought, :final_answer + + def parse + return unless @output + + @parsed_thoughts = parse_json_objects + + return unless @parsed_thoughts.present? + + parse_final_answer + parse_action + end + + private + + def final_answer? + @parsed_thoughts.first[:type] == 'final_answer_delta' + end + + def parse_final_answer + return unless final_answer? + + @final_answer = '' + + @parsed_thoughts.each do |t| + @final_answer += t[:data][:text] + end + + @final_answer + end + + def parse_action + response = @parsed_thoughts.first + + return unless response[:type] == 'action' + + @thought = response[:data][:thought] + @action = response[:data][:tool].camelcase + @action_input = response[:data][:tool_input] + end + + def parse_json_objects + json_strings = @output.split("\n") + + json_strings.map do |str| + Gitlab::Json.parse(str).with_indifferent_access + end + end + end + end + end + end +end diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb index 33b5dd56e93d49a3bc43add5242e067a5e6687c1..b64e5f336d64b93df35753ec36b7fa66c9f027e1 100644 --- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb +++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb @@ -15,6 +15,7 @@ class AiGateway < Base ENDPOINT = '/v1/chat/agent' BASE_ENDPOINT = '/v1/chat' + CHAT_V2_ENDPOINT = '/v2/chat/agent' DEFAULT_TYPE = 'prompt' DEFAULT_SOURCE = 'GitLab EE' TEMPERATURE = 0.1 @@ -33,10 +34,16 @@ def request(prompt, unit_primitive: nil) options = default_options.merge(prompt.fetch(:options, {})) return unless model_provider_valid?(options) - body = request_body(prompt: prompt[:prompt], options: options) + v2_chat_schema = Feature.enabled?(:v2_chat_agent_integration, user) && options.delete(:single_action_agent) + + body = if v2_chat_schema + request_body_chat_2(prompt: prompt[:prompt], options: options) + else + request_body(prompt: prompt[:prompt], options: options) + end response = ai_client.stream( - endpoint: endpoint(unit_primitive), + endpoint: endpoint(unit_primitive, v2_chat_schema), body: body ) do |data| yield data if block_given? @@ -76,9 +83,11 @@ def model_provider_valid?(options) provider(options) end - def endpoint(unit_primitive) + def endpoint(unit_primitive, v2_chat_schema) if unit_primitive.present? "#{BASE_ENDPOINT}/#{unit_primitive}" + elsif v2_chat_schema + CHAT_V2_ENDPOINT else ENDPOINT end @@ -118,6 +127,28 @@ def model_params(options) end end + def request_body_chat_2(prompt:, options: {}) + option_params = { + chat_history: "", + agent_scratchpad: { + agent_type: "react", + steps: options[:agent_scratchpad] + } + } + + if options[:current_resource_type] + option_params[:context] = { + type: options[:current_resource_type], + content: options[:current_resource_content] + } + end + + { + prompt: prompt, + options: option_params + } + end + def payload_params(options) allowed_params = ALLOWED_PARAMS.fetch(provider(options)) params = options.slice(*allowed_params) diff --git a/ee/lib/gitlab/llm/completions/chat.rb b/ee/lib/gitlab/llm/completions/chat.rb index 066b770835586a778a9526ac1e2a4cdeeeee080d..0d98c7d7d1ec961a90d9c5fbe87da51d1d9bf1ee 100644 --- a/ee/lib/gitlab/llm/completions/chat.rb +++ b/ee/lib/gitlab/llm/completions/chat.rb @@ -23,6 +23,9 @@ class Chat < Base ::Gitlab::Llm::Chain::Tools::ExplainVulnerability ].freeze + # @param [Gitlab::Llm::AiMessage] prompt_message - user question + # @param [NilClass] ai_prompt_class - not used for chat + # @param [Hash] options - additional context def initialize(prompt_message, ai_prompt_class, options = {}) super @@ -109,13 +112,23 @@ def agent_or_tool_response(response_handler) end def execute_with_tool_chosen_by_ai(response_handler, stream_response_handler) - Gitlab::Llm::Chain::Agents::ZeroShot::Executor.new( - user_input: prompt_message.content, - tools: tools, - context: context, - response_handler: response_handler, - stream_response_handler: stream_response_handler - ).execute + if Feature.enabled?(:v2_chat_agent_integration, user) + Gitlab::Llm::Chain::Agents::SingleActionExecutor.new( + user_input: prompt_message.content, + tools: tools, + context: context, + response_handler: response_handler, + stream_response_handler: stream_response_handler + ).execute + else + Gitlab::Llm::Chain::Agents::ZeroShot::Executor.new( + user_input: prompt_message.content, + tools: tools, + context: context, + response_handler: response_handler, + stream_response_handler: stream_response_handler + ).execute + end end def execute_with_slash_command_tool(stream_response_handler) diff --git a/ee/spec/factories/llm/chain/answers.rb b/ee/spec/factories/llm/chain/answers.rb new file mode 100644 index 0000000000000000000000000000000000000000..3b22dde1593622c4cb446661196d9be148c8f4fa --- /dev/null +++ b/ee/spec/factories/llm/chain/answers.rb @@ -0,0 +1,36 @@ +# frozen_string_literal: true + +FactoryBot.define do + factory :answer, class: '::Gitlab::Llm::Chain::Answer' do + status { :ok } + is_final { false } + gitlab_context { 'context' } + content { 'content' } + tool { nil } + suggestion { nil } + extras { nil } + + trait :final do + is_final { true } + end + + trait :tool do + tool { Gitlab::Llm::Chain::Tools::IssueReader } + suggestion { 'suggestion' } + end + + initialize_with do + new( + status: status, + context: gitlab_context, + content: content, + tool: tool, + suggestions: suggestion, + is_final: is_final, + extras: extras + ) + end + + skip_create + end +end diff --git a/ee/spec/features/duo_chat_spec.rb b/ee/spec/features/duo_chat_spec.rb index 24ae5a84aabc77bebb542df43cef39b5ff519976..59bd023a57e4c4af5a00ddd60f63135d8a360795 100644 --- a/ee/spec/features/duo_chat_spec.rb +++ b/ee/spec/features/duo_chat_spec.rb @@ -34,13 +34,8 @@ let(:chat_response) { "Final Answer: #{answer}" } before do - # TODO: Switch to AI Gateway - # See https://gitlab.com/gitlab-org/gitlab/-/issues/431563 - stub_request(:post, "https://api.anthropic.com/v1/complete") - .to_return( - status: 200, body: { completion: "question_category" }.to_json, - headers: { 'Content-Type' => 'application/json' } - ) + # TODO: remove with https://gitlab.com/gitlab-org/gitlab/-/issues/456258 + stub_feature_flags(v2_chat_agent_integration: false) stub_request(:post, "#{Gitlab::AiGateway.url}/v1/chat/agent") .with(body: hash_including({ "stream" => true })) diff --git a/ee/spec/lib/gitlab/llm/chain/agents/single_action_executor_spec.rb b/ee/spec/lib/gitlab/llm/chain/agents/single_action_executor_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..1a6dbcd834c68b6cdf6c8616573e44bce9b1150e --- /dev/null +++ b/ee/spec/lib/gitlab/llm/chain/agents/single_action_executor_spec.rb @@ -0,0 +1,218 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Llm::Chain::Agents::SingleActionExecutor, feature_category: :duo_chat do + describe "#execute" do + subject(:answer) { agent.execute } + + let(:agent) do + described_class.new( + user_input: user_input, + tools: tools, + context: context, + response_handler: response_service_double, + stream_response_handler: stream_response_service_double + ) + end + + let_it_be(:issue) { build_stubbed(:issue) } + let_it_be(:resource) { issue } + let_it_be(:user) { issue.author } + + let(:user_input) { 'question?' } + let(:tools) { [Gitlab::Llm::Chain::Tools::IssueReader] } + let(:tool_double) { instance_double(Gitlab::Llm::Chain::Tools::IssueReader::Executor) } + let(:response_service_double) { instance_double(::Gitlab::Llm::ResponseService) } + let(:stream_response_service_double) { instance_double(::Gitlab::Llm::ResponseService) } + + let(:ai_request_double) { instance_double(Gitlab::Llm::Chain::Requests::AiGateway) } + + let(:context) do + Gitlab::Llm::Chain::GitlabContext.new( + current_user: user, container: nil, resource: resource, ai_request: ai_request_double, + extra_resource: nil, current_file: nil, agent_version: nil + ) + end + + let(:answer_chunk) do + "{\"type\":\"final_answer_delta\",\"data\":{\"thought\":\"Thought: direct answer.\",\"text\":\"Ans\"}}" + end + + before do + allow(context).to receive(:ai_request).and_return(ai_request_double) + allow(ai_request_double).to receive(:request).and_return(answer_chunk) + end + + context "when answer is final" do + let(:another_chunk) do + "{\"type\":\"final_answer_delta\",\"data\":{\"thought\":\"\",\"text\":\"wer\"}}" + end + + let(:response_double) do + "#{answer_chunk}\n#{another_chunk}" + end + + let(:first_response_double) { double } + let(:second_response_double) { double } + + before do + allow(ai_request_double).to receive(:request).and_yield(answer_chunk) + .and_yield(another_chunk) + .and_return(response_double) + allow(Gitlab::Llm::Chain::StreamedResponseModifier).to receive(:new).with("Ans", { chunk_id: 1 }) + .and_return(first_response_double) + allow(Gitlab::Llm::Chain::StreamedResponseModifier).to receive(:new).with("wer", { chunk_id: 2 }) + .and_return(second_response_double) + + allow(context).to receive(:current_page_type).and_return("issue") + allow(context).to receive(:current_page_short_description).and_return("issue description") + end + + it "streams final answer" do + expect(stream_response_service_double).to receive(:execute).with( + response: first_response_double, + options: { chunk_id: 1 } + ) + expect(stream_response_service_double).to receive(:execute).with( + response: second_response_double, + options: { chunk_id: 2 } + ) + + expect(ai_request_double).to receive(:request).with( + { + prompt: user_input, + options: { + agent_scratchpad: [], + conversation: [], + single_action_agent: true, + current_resource_type: "issue", + current_resource_content: "issue description" + } + }, + { unit_primitive: nil } + ) + + expect(answer.is_final?).to be_truthy + expect(answer.content).to include("Answer") + end + end + + context "when tool answer if final" do + let(:llm_answer) { create(:answer, :tool, tool: Gitlab::Llm::Chain::Tools::IssueReader::Executor) } + let(:tool_answer) { create(:answer, :final, content: 'tool answer') } + + before do + allow(::Gitlab::Llm::Chain::Answer).to receive(:from_response).and_return(llm_answer) + + allow_next_instance_of(Gitlab::Llm::Chain::Tools::IssueReader::Executor) do |issue_tool| + allow(issue_tool).to receive(:execute).and_return(tool_answer) + end + end + + it "returns tool answer" do + expect(stream_response_service_double).to receive(:execute) + expect(answer.is_final?).to be_truthy + expect(answer.content).to include("tool answer") + end + end + + context "when max iteration reached" do + let(:llm_answer) { create(:answer, :tool, tool: Gitlab::Llm::Chain::Tools::IssueReader::Executor) } + + before do + stub_const("#{described_class.name}::MAX_ITERATIONS", 2) + allow(stream_response_service_double).to receive(:execute) + allow(::Gitlab::Llm::Chain::Answer).to receive(:from_response).and_return(llm_answer) + + allow_next_instance_of(Gitlab::Llm::Chain::Tools::IssueReader::Executor) do |issue_tool| + allow(issue_tool).to receive(:execute).and_return(llm_answer) + end + end + + it "returns default answer" do + expect(answer.is_final?).to eq(true) + expect(answer.content).to include(Gitlab::Llm::Chain::Answer.default_final_message) + end + end + + context "when resource is not authorized" do + let(:resource) { user } + + it "sends request without context" do + expect(ai_request_double).to receive(:request).with( + { + prompt: user_input, + options: { + agent_scratchpad: [], + conversation: [], + single_action_agent: true, + current_resource_type: nil, + current_resource_content: nil + } + }, + { unit_primitive: nil } + ) + + agent.execute + end + end + + context "when times out error is raised" do + let(:error) { Net::ReadTimeout.new } + + before do + allow(Gitlab::ErrorTracking).to receive(:track_exception) + end + + shared_examples "time out error" do + it "returns an error" do + expect(answer.is_final?).to eq(true) + expect(answer.content).to include("I'm sorry, I couldn't respond in time. Please try again.") + expect(answer.error_code).to include("A1000") + expect(Gitlab::ErrorTracking).to have_received(:track_exception).with(error) + end + end + + context "when streamed request times out" do + before do + allow(ai_request_double).to receive(:request).and_raise(error) + end + + it_behaves_like "time out error" + end + + context "when tool times out out" do + let(:llm_answer) { create(:answer, :tool, tool: Gitlab::Llm::Chain::Tools::IssueReader::Executor) } + + before do + allow(ai_request_double).to receive(:request) + allow(::Gitlab::Llm::Chain::Answer).to receive(:from_response).and_return(llm_answer) + allow_next_instance_of(Gitlab::Llm::Chain::Tools::IssueReader::Executor) do |issue_tool| + allow(issue_tool).to receive(:execute).and_raise(error) + end + + allow(stream_response_service_double).to receive(:execute) + end + + it_behaves_like "time out error" + end + end + + context "when connection error is raised" do + let(:error) { ::Gitlab::Llm::AiGateway::Client::ConnectionError.new } + + before do + allow(Gitlab::ErrorTracking).to receive(:track_exception) + allow(ai_request_double).to receive(:request).and_raise(error) + end + + it "returns an error" do + expect(answer.is_final).to eq(true) + expect(answer.content).to include("I'm sorry, I can't generate a response. Please try again.") + expect(answer.error_code).to include("A1001") + expect(Gitlab::ErrorTracking).to have_received(:track_exception).with(error) + end + end + end +end diff --git a/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/qa_evaluation_spec.rb b/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/qa_evaluation_spec.rb index e2b017ccc2356696ea0841d8444a5526bf50ab1d..387f9c0ebde5fd53e49223d5a3f8e70943bb4cee 100644 --- a/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/qa_evaluation_spec.rb +++ b/ee/spec/lib/gitlab/llm/chain/agents/zero_shot/qa_evaluation_spec.rb @@ -14,6 +14,7 @@ let_it_be(:issue_fixtures) { load_fixture('issues') } before_all do + stub_feature_flags(v2_chat_agent_integration: false) # link_reference_pattern is memoized for Issue # and stubbed url (gitlab.com) is not used to derive the link reference pattern. Issue.instance_variable_set(:@link_reference_pattern, nil) diff --git a/ee/spec/lib/gitlab/llm/chain/answer_spec.rb b/ee/spec/lib/gitlab/llm/chain/answer_spec.rb index f758247c2475b3d4c2b8d3da7ce41fbed342d4d2..845c0b8e2469a4c11528bd412cc78c09bf7dc70c 100644 --- a/ee/spec/lib/gitlab/llm/chain/answer_spec.rb +++ b/ee/spec/lib/gitlab/llm/chain/answer_spec.rb @@ -100,6 +100,33 @@ expect(answer.content).to eq(input) end end + + context 'with different parser' do + subject(:answer) do + described_class.from_response( + response_body: input, + tools: tools, + context: context, + parser_klass: Gitlab::Llm::Chain::Parsers::SingleActionParser + ) + end + + let(:input) do + { + type: "action", + data: { + thought: "Thought: I need to retrieve the issue content using the \"issue_reader\" tool.", + tool: "issue_reader", + tool_input: "what is the title of this issue" + } + }.to_json + end + + it 'returns intermediate answer with parsed values and a tool' do + expect(answer.is_final?).to eq(false) + expect(answer.tool::NAME).to eq('IssueReader') + end + end end describe '.final_answer' do diff --git a/ee/spec/lib/gitlab/llm/chain/answers/streamed_json_spec.rb b/ee/spec/lib/gitlab/llm/chain/answers/streamed_json_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..7dfa14bab3857574df15872774572debb5e185bd --- /dev/null +++ b/ee/spec/lib/gitlab/llm/chain/answers/streamed_json_spec.rb @@ -0,0 +1,50 @@ +# frozen_string_literal: true + +require 'fast_spec_helper' + +RSpec.describe Gitlab::Llm::Chain::Answers::StreamedJson, feature_category: :duo_chat do + describe "#next_chunk" do + subject { described_class.new.next_chunk(chunk) } + + context "when stream is empty" do + let(:chunk) { "" } + + it 'returns nil' do + is_expected.to be_nil + end + end + + context "when stream does not contain the final answer" do + let(:chunk) do + { + type: "action", + data: { + thought: "Thought: I need to retrieve the issue content using the \"issue_reader\" tool.", + tool: "issue_reader", + tool_input: "what is the title of this issue" + } + }.to_json + end + + it 'returns nil' do + is_expected.to be_nil + end + end + + context "when streaming beginning of the answer" do + let(:chunk) do + { + type: "final_answer_delta", + data: { + thought: "Thought: I should provide a direct response.", + text: "I" + } + }.to_json + end + + it 'returns stream payload' do + is_expected.to eq({ id: 1, content: "I" }) + end + end + end +end diff --git a/ee/spec/lib/gitlab/llm/chain/gitlab_context_spec.rb b/ee/spec/lib/gitlab/llm/chain/gitlab_context_spec.rb index 349acab2fbd4b6d27fffb050644e24b1cb9b5608..0100995a78359305c669d153c4f30bd54e67290c 100644 --- a/ee/spec/lib/gitlab/llm/chain/gitlab_context_spec.rb +++ b/ee/spec/lib/gitlab/llm/chain/gitlab_context_spec.rb @@ -24,6 +24,14 @@ group.namespace_settings.update!(experiment_features_enabled: true) end + describe '#current_page_type' do + let(:resource) { create(:issue, project: project) } + + it 'delegates to ai resource' do + expect(context.current_page_type).to eq("issue") + end + end + describe '#resource_serialized' do let(:content_limit) { 500 } diff --git a/ee/spec/lib/gitlab/llm/chain/parsers/single_action_parser_spec.rb b/ee/spec/lib/gitlab/llm/chain/parsers/single_action_parser_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..526e5afd1ba128de1b3a993b62f080bd18e34c05 --- /dev/null +++ b/ee/spec/lib/gitlab/llm/chain/parsers/single_action_parser_spec.rb @@ -0,0 +1,81 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Llm::Chain::Parsers::SingleActionParser, feature_category: :duo_chat do + describe "#parse" do + let(:parser) { described_class.new(output: output) } + let(:output) { chunks.map(&:to_json).join("\n") } + + before do + parser.parse + end + + context "with final answer" do + let(:chunks) do + [ + { + type: "final_answer_delta", + data: { + thought: "Thought: I don't need any specific GitLab resources to answer this.", + text: "To" + } + }, + { type: "final_answer_delta", data: { thought: "", text: " perform" } }, + { type: "final_answer_delta", data: { thought: "", text: " a" } }, + { type: "final_answer_delta", data: { thought: "", text: " Git" } }, + { type: "final_answer_delta", data: { thought: "", text: " re" } }, + { type: "final_answer_delta", data: { thought: "", text: "base" } }, + { type: "final_answer_delta", data: { thought: "", text: "," } } + ] + end + + it "returns only the final answer" do + expect(parser.action).to be_nil + expect(parser.action_input).to be_nil + expect(parser.thought).to be_nil + expect(parser.final_answer).to eq("To perform a Git rebase,") + end + end + + context "with chosen action" do + let(:chunks) do + [ + { + type: "action", + data: { + thought: "Thought: I need to retrieve the issue details using the \"issue_reader\" tool.", + tool: "issue_reader", + tool_input: "What is the title of this issue?" + } + } + ] + end + + it "returns the action" do + expect(parser.action).to eq("IssueReader") + expect(parser.action_input).to eq("What is the title of this issue?") + expect(parser.thought).to eq("Thought: I need to retrieve the issue details using the \"issue_reader\" tool.") + expect(parser.final_answer).to be_nil + end + end + + context "with no output" do + let(:output) { nil } + + it "returns nil" do + expect(parser.action).to be_nil + expect(parser.final_answer).to be_nil + end + end + + context "with empty output" do + let(:output) { "" } + + it "returns nil" do + expect(parser.action).to be_nil + expect(parser.final_answer).to be_nil + end + end + end +end diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb index 4c2378deed66ba443a7e1c9b3216a1fe0d416a03..c414e1d75c179a49d051b29f9a891ef31ac372e0 100644 --- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb +++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb @@ -165,5 +165,67 @@ it_behaves_like 'performing request to the AI Gateway' end + + context 'when request is sent for a new ReAct Duo Chat prompt' do + let(:endpoint) { described_class::CHAT_V2_ENDPOINT } + + let(:prompt) { { prompt: user_prompt, options: options } } + + let(:options) do + { + agent_scratchpad: [], + single_action_agent: true, + current_resource_type: "issue", + current_resource_content: "string" + } + end + + let(:body) do + { + prompt: user_prompt, + options: { + chat_history: "", + agent_scratchpad: { + agent_type: "react", + steps: [] + }, + context: { + type: "issue", + content: "string" + } + } + } + end + + it_behaves_like 'performing request to the AI Gateway' + end + + context 'when request is sent for a new ReAct Duo Chat prompt withouth context params' do + let(:endpoint) { described_class::CHAT_V2_ENDPOINT } + + let(:prompt) { { prompt: user_prompt, options: options } } + + let(:options) do + { + agent_scratchpad: [], + single_action_agent: true + } + end + + let(:body) do + { + prompt: user_prompt, + options: { + chat_history: "", + agent_scratchpad: { + agent_type: "react", + steps: [] + } + } + } + end + + it_behaves_like 'performing request to the AI Gateway' + end end end diff --git a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb index 00e649ccb50007f48d354f645b5bb24b9d828415..6aab877abf12d413945914e12651309cb368bf73 100644 --- a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb +++ b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb @@ -71,7 +71,7 @@ subject { described_class.new(prompt_message, nil, **options).execute } shared_examples 'success' do - xit 'calls the ZeroShot Agent with the right parameters', :snowplow do + xit 'calls the SingleAction Agent with the right parameters', :snowplow do expected_params = [ user_input: content, tools: match_array(tools), @@ -80,7 +80,7 @@ stream_response_handler: stream_response_handler ] - expect_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor, *expected_params) do |instance| + expect_next_instance_of(::Gitlab::Llm::Chain::Agents::SingleActionExecutor, *expected_params) do |instance| expect(instance).to receive(:execute).and_return(answer) end @@ -126,7 +126,7 @@ stream_response_handler: stream_response_handler ] - expect_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor, *expected_params) do |instance| + expect_next_instance_of(::Gitlab::Llm::Chain::Agents::SingleActionExecutor, *expected_params) do |instance| expect(instance).to receive(:execute).and_return(answer) end @@ -157,7 +157,7 @@ end xit 'sends process_gitlab_duo_question snowplow event with value eql 0' do - allow_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor) do |instance| + allow_next_instance_of(::Gitlab::Llm::Chain::Agents::SingleActionExecutor) do |instance| expect(instance).to receive(:execute).and_return(answer) end @@ -218,7 +218,7 @@ stream_response_handler: stream_response_handler ] - expect_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor, *expected_params) do |instance| + expect_next_instance_of(::Gitlab::Llm::Chain::Agents::SingleActionExecutor, *expected_params) do |instance| expect(instance).to receive(:execute).and_return(answer) end expect(response_handler).to receive(:execute) @@ -254,7 +254,7 @@ command: an_instance_of(::Gitlab::Llm::Chain::SlashCommand) } - expect(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor).not_to receive(:new) + expect(::Gitlab::Llm::Chain::Agents::SingleActionExecutor).not_to receive(:new) expect(expected_tool) .to receive(:new).with(expected_params).and_return(executor) @@ -308,7 +308,7 @@ let(:command) { '/explain2' } it 'process the message with zero shot agent' do - expect_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor) do |instance| + expect_next_instance_of(::Gitlab::Llm::Chain::Agents::SingleActionExecutor) do |instance| expect(instance).to receive(:execute).and_return(answer) end expect(::Gitlab::Llm::Chain::Tools::ExplainCode::Executor).not_to receive(:new) @@ -332,7 +332,7 @@ stream_response_handler: stream_response_handler ] - allow_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor, *expected_params) do |instance| + allow_next_instance_of(::Gitlab::Llm::Chain::Agents::SingleActionExecutor, *expected_params) do |instance| allow(instance).to receive(:execute).and_return(answer) end @@ -350,5 +350,50 @@ subject end end + + context 'with disabled v2_chat_agent_integration flag' do + before do + stub_feature_flags(v2_chat_agent_integration: false) + end + + xit 'calls the ZeroShot Agent with the right parameters', :snowplow do + expected_params = [ + user_input: content, + tools: match_array(tools), + context: context, + response_handler: response_handler, + stream_response_handler: stream_response_handler + ] + + expect_next_instance_of(::Gitlab::Llm::Chain::Agents::ZeroShot::Executor, *expected_params) do |instance| + expect(instance).to receive(:execute).and_return(answer) + end + + expect(response_handler).to receive(:execute) + expect(::Gitlab::Llm::ResponseService).to receive(:new).with(context, { request_id: 'uuid', ai_action: :chat }) + .and_return(response_handler) + expect(::Gitlab::Llm::Chain::GitlabContext).to receive(:new) + .with(current_user: user, container: expected_container, resource: resource, ai_request: ai_request, + extra_resource: extra_resource, request_id: 'uuid', current_file: current_file, + agent_version: agent_version) + .and_return(context) + expect(categorize_service).to receive(:execute) + expect(::Llm::ExecuteMethodService).to receive(:new) + .with(user, user, :categorize_question, categorize_service_params) + .and_return(categorize_service) + + subject + + expect_snowplow_event( + category: described_class.to_s, + label: "IssueReader", + action: 'process_gitlab_duo_question', + property: 'uuid', + namespace: container, + user: user, + value: 1 + ) + 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 4113331e6e16bb57902432a12db566bf193fcdeb..fa85be72bbd6df02d2a01bead2da212075b2cd96 100644 --- a/ee/spec/models/ai/ai_resource/epic_spec.rb +++ b/ee/spec/models/ai/ai_resource/epic_spec.rb @@ -36,4 +36,10 @@ .not_to include("utilize it instead of using the 'EpicReader' tool") end end + + describe '#current_page_type' do + it 'returns type' do + expect(wrapped_epic.current_page_type).to eq('epic') + 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 e4f6947c1ec23753072697a46214395ea1bfa113..91e653a9d1f3ea189b6dad281aa6288135979c5b 100644 --- a/ee/spec/models/ai/ai_resource/issue_spec.rb +++ b/ee/spec/models/ai/ai_resource/issue_spec.rb @@ -35,4 +35,10 @@ .not_to include("utilize it instead of using the 'IssueReader' tool") end end + + describe '#current_page_type' do + it 'returns type' do + expect(wrapped_issue.current_page_type).to eq('issue') + end + end end