Stream Duo Chat slash commands responses

ee/lib/gitlab/llm/chain/agents/zero_shot/executor.rb

@@ -80,22 +80,7 @@ def execute
             private
 
             def execute_streamed_request
-              streamed_answer = StreamedZeroShotAnswer.new
-
-              request do |content|
-                next unless stream_response_handler
-
-                chunk = streamed_answer.next_chunk(content)
-
-                if chunk
-                  stream_response_handler.execute(
-                    response: Gitlab::Llm::Chain::PlainResponseModifier.new(content),
-                    options: {
-                      chunk_id: chunk[:id]
-                    }
-                  )
-                end
-              end
+              request(&streamed_request_handler(StreamedZeroShotAnswer.new))
             end
 
             attr_reader :logger, :stream_response_handler

ee/lib/gitlab/llm/chain/concerns/ai_dependent.rb

@@ -19,6 +19,21 @@ def request(&block)
             ai_request.request(prompt_str, &block)
           end
 
+          def streamed_request_handler(streamed_answer)
+            proc do |content|
+              next unless stream_response_handler
+
+              chunk = streamed_answer.next_chunk(content)
+
+              if chunk
+                stream_response_handler.execute(
+                  response: Gitlab::Llm::Chain::PlainResponseModifier.new(content),
+                  options: { chunk_id: chunk[:id] }
+                )
+              end
+            end
+          end
+
           private
 
           def ai_request

ee/lib/gitlab/llm/chain/tools/explain_code/executor.rb

@@ -57,12 +57,6 @@ def self.slash_commands
               SLASH_COMMANDS
             end
 
-            def perform
-              Answer.new(status: :ok, context: context, content: request, tool: nil)
-            rescue StandardError
-              Answer.error_answer(context: context, content: _("Unexpected error"))
-            end
-
             private
 
             def authorize

ee/lib/gitlab/llm/chain/tools/refactor_code/executor.rb

@@ -64,12 +64,6 @@ def self.slash_commands
               SLASH_COMMANDS
             end
 
-            def perform
-              Answer.new(status: :ok, context: context, content: request, tool: nil)
-            rescue StandardError
-              Answer.error_answer(context: context, content: _("Unexpected error"))
-            end
-
             private
 
             def selected_text_options

ee/lib/gitlab/llm/chain/tools/slash_command_tool.rb

@@ -7,6 +7,14 @@ module Tools
         class SlashCommandTool < Tool
           extend ::Gitlab::Utils::Override
 
+          def perform
+            content = request(&streamed_request_handler(StreamedAnswer.new))
+
+            Answer.new(status: :ok, context: context, content: content, tool: nil)
+          rescue StandardError
+            Answer.error_answer(context: context, content: _("Unexpected error"))
+          end
+
           private
 
           attr_reader :command

ee/lib/gitlab/llm/chain/tools/write_tests/executor.rb

@@ -63,12 +63,6 @@ def self.slash_commands
               SLASH_COMMANDS
             end
 
-            def perform
-              Answer.new(status: :ok, context: context, content: request, tool: nil)
-            rescue StandardError
-              Answer.error_answer(context: context, content: _("Unexpected error"))
-            end
-
             private
 
             def authorize

ee/spec/lib/gitlab/llm/chain/tools/explain_code/executor_spec.rb

@@ -8,6 +8,7 @@
   let(:ai_request_double) { instance_double(Gitlab::Llm::Chain::Requests::Anthropic) }
   let(:input) { 'input' }
   let(:options) { { input: input } }
+  let(:stream_response_handler) { nil }
   let(:command) { nil }
 
   let(:context) do
@@ -22,7 +23,11 @@
     )
   end
 
-  subject(:tool) { described_class.new(context: context, options: options, command: command) }
+  subject(:tool) do
+    described_class.new(
+      context: context, options: options, stream_response_handler: stream_response_handler, command: command
+    )
+  end
 
   describe '#name' do
     it 'returns tool name' do

ee/spec/lib/gitlab/llm/chain/tools/refactor_code/executor_spec.rb

@@ -8,6 +8,7 @@
   let(:ai_request_double) { instance_double(Gitlab::Llm::Chain::Requests::Anthropic) }
   let(:input) { 'input' }
   let(:options) { { input: input } }
+  let(:stream_response_handler) { nil }
   let(:command) { nil }
 
   let(:context) do
@@ -22,7 +23,11 @@
     )
   end
 
-  subject(:tool) { described_class.new(context: context, options: options, command: command) }
+  subject(:tool) do
+    described_class.new(
+      context: context, options: options, stream_response_handler: stream_response_handler, command: command
+    )
+  end
 
   describe '#name' do
     it 'returns tool name' do

ee/spec/lib/gitlab/llm/chain/tools/write_tests/executor_spec.rb

@@ -8,6 +8,7 @@
   let(:ai_request_double) { instance_double(Gitlab::Llm::Chain::Requests::Anthropic) }
   let(:input) { 'input' }
   let(:options) { { input: input } }
+  let(:stream_response_handler) { nil }
   let(:command) { nil }
 
   let(:context) do
@@ -17,7 +18,11 @@
     )
   end
 
-  subject(:tool) { described_class.new(context: context, options: options, command: command) }
+  subject(:tool) do
+    described_class.new(
+      context: context, options: options, stream_response_handler: stream_response_handler, command: command
+    )
+  end
 
   describe '#name' do
     it 'returns tool name' do

ee/spec/support/shared_examples/lib/gitlab/llm/chain/slash_command_tool_shared_examples.rb

@@ -84,4 +84,34 @@
       tool.execute
     end
   end
+
+  context 'when stream_response_service is set' do
+    let(:stream_response_handler) { instance_double(::Gitlab::Llm::ResponseService) }
+
+    before do
+      allow(ai_request_double).to receive(:request).and_yield("Hello").and_yield(" World")
+    end
+
+    it 'streams the final answer' do
+      first_response_double = double
+      second_response_double = double
+
+      allow(Gitlab::Llm::Chain::PlainResponseModifier).to receive(:new).with("Hello")
+        .and_return(first_response_double)
+
+      allow(Gitlab::Llm::Chain::PlainResponseModifier).to receive(:new).with(" World")
+        .and_return(second_response_double)
+
+      expect(stream_response_handler).to receive(:execute).with(
+        response: first_response_double,
+        options: { chunk_id: 1 }
+      )
+      expect(stream_response_handler).to receive(:execute).with(
+        response: second_response_double,
+        options: { chunk_id: 2 }
+      )
+
+      tool.execute
+    end
+  end
 end