diff --git a/ee/lib/gitlab/llm/completions/chat.rb b/ee/lib/gitlab/llm/completions/chat.rb
index 3997b55f696dc49b388e861414c87a66ea421a0b..59e157b3026af834ce652683e2a438bb9667d43c 100644
--- a/ee/lib/gitlab/llm/completions/chat.rb
+++ b/ee/lib/gitlab/llm/completions/chat.rb
@@ -70,7 +70,8 @@ def execute
# Send full message to custom clientSubscriptionId at the end of streaming.
if response_options[:client_subscription_id]
- ::Gitlab::Llm::ResponseService.new(context, response_options).execute(response: response_modifier)
+ ::Gitlab::Llm::ResponseService.new(context, response_options)
+ .execute(response: response_modifier, save_message: false)
end
response_handler.execute(response: response_modifier)
diff --git a/ee/lib/gitlab/llm/graphql_subscription_response_service.rb b/ee/lib/gitlab/llm/graphql_subscription_response_service.rb
index cac52ab5af9ad32c0a5108385abcefd1243d1443..13906175d81f8f2f66dd6769313c738571e0dff8 100644
--- a/ee/lib/gitlab/llm/graphql_subscription_response_service.rb
+++ b/ee/lib/gitlab/llm/graphql_subscription_response_service.rb
@@ -3,11 +3,12 @@
module Gitlab
module Llm
class GraphqlSubscriptionResponseService < BaseService
- def initialize(user, resource, response_modifier, options:)
+ def initialize(user, resource, response_modifier, options:, save_message: true)
@user = user
@resource = resource
@response_modifier = response_modifier
@options = options
+ @save_message = save_message
@logger = Gitlab::Llm::Logger.build
end
@@ -39,10 +40,11 @@ def execute
private
- attr_reader :user, :resource, :response_modifier, :options, :logger
+ attr_reader :user, :resource, :response_modifier, :options, :save_message, :logger
def save_message?
- response_message.is_a?(ChatMessage) &&
+ save_message &&
+ response_message.is_a?(ChatMessage) &&
!response_message.type &&
!response_message.chunk_id
end
diff --git a/ee/lib/gitlab/llm/response_service.rb b/ee/lib/gitlab/llm/response_service.rb
index b701937d8f721c6c954fed2b23372629e59bae42..709dc0443e678889fedcc25a4664a700b4f8a7aa 100644
--- a/ee/lib/gitlab/llm/response_service.rb
+++ b/ee/lib/gitlab/llm/response_service.rb
@@ -9,10 +9,11 @@ def initialize(context, basic_options)
@basic_options = basic_options
end
- def execute(response:, options: {})
+ def execute(response:, options: {}, save_message: true)
::Gitlab::Llm::GraphqlSubscriptionResponseService
.new(user, resource, response,
- options: basic_options.merge(options))
+ options: basic_options.merge(options),
+ save_message: save_message)
.execute
end
diff --git a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
index 6ffe9c80a727c56104c808544050f25fbb820141..3de78b5ac26f276a5963b5bc1ac7fdc51c99647c 100644
--- a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
+++ b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
@@ -140,7 +140,7 @@
an_instance_of(Gitlab::Llm::Chain::GitlabContext), { request_id: 'uuid', ai_action: :chat,
client_subscription_id: 'someid' }
).and_return(stream_response_handler).twice
- expect(stream_response_handler).to receive(:execute)
+ expect(stream_response_handler).to receive(:execute).with(response: anything, save_message: false)
expect(categorize_service).to receive(:execute)
expect(::Llm::ExecuteMethodService).to receive(:new)
.with(user, user, :categorize_question, categorize_service_params)
diff --git a/ee/spec/lib/gitlab/llm/graphql_subscription_response_service_spec.rb b/ee/spec/lib/gitlab/llm/graphql_subscription_response_service_spec.rb
index a3ff19fb21c64b579349b7d8d00b5d28c8b238f4..6fb61ddb062e402d41b7beb8a5d51c68917d05b2 100644
--- a/ee/spec/lib/gitlab/llm/graphql_subscription_response_service_spec.rb
+++ b/ee/spec/lib/gitlab/llm/graphql_subscription_response_service_spec.rb
@@ -75,12 +75,26 @@
end
describe '#execute' do
- let(:service) { described_class.new(user, resource, response_modifier, options: options) }
+ let(:save_message) { true }
+ let(:service) do
+ described_class.new(user, resource, response_modifier, options: options, save_message: save_message)
+ end
+
let_it_be(:resource) { project }
subject { service.execute }
context 'when message is chat' do
+ shared_examples 'not saving the message' do
+ it 'does not save the message' do
+ expect_next_instance_of(::Gitlab::Llm::AiMessage) do |instance|
+ expect(instance).not_to receive(:save!)
+ end
+
+ subject
+ end
+ end
+
let(:ai_action) { 'chat' }
it 'saves the message' do
@@ -89,28 +103,22 @@
subject
end
+ context 'when save_message is false' do
+ let(:save_message) { false }
+
+ it_behaves_like 'not saving the message'
+ end
+
context 'when message is stream chunk' do
let(:options) { super().merge(chunk_id: 1) }
- it 'does not save the message' do
- expect_next_instance_of(::Gitlab::Llm::AiMessage) do |instance|
- expect(instance).not_to receive(:save!)
- end
-
- subject
- end
+ it_behaves_like 'not saving the message'
end
context 'when message has special type' do
let(:options) { super().merge(type: 'tool') }
- it 'does not save the message' do
- expect_next_instance_of(::Gitlab::Llm::AiMessage) do |instance|
- expect(instance).not_to receive(:save!)
- end
-
- subject
- end
+ it_behaves_like 'not saving the message'
end
end
diff --git a/ee/spec/lib/gitlab/llm/response_service_spec.rb b/ee/spec/lib/gitlab/llm/response_service_spec.rb
index 40a294733d9e3e79eaaa3c24c3e0d60f27778e58..27125c4b68e147445ab135b2b110efe011d61757 100644
--- a/ee/spec/lib/gitlab/llm/response_service_spec.rb
+++ b/ee/spec/lib/gitlab/llm/response_service_spec.rb
@@ -11,15 +11,18 @@
let(:basic_options) { { cache_request: true } }
let(:options) { { cache_request: false } }
+ let(:save_message) { false }
let(:graphql_subscription_double) { instance_double(::Gitlab::Llm::GraphqlSubscriptionResponseService) }
describe '#execute' do
it 'calls GraphQL subscription service with the right params' do
expect(graphql_subscription_double).to receive(:execute)
expect(::Gitlab::Llm::GraphqlSubscriptionResponseService).to receive(:new)
- .with(user, issue, 'response', options: { cache_request: false }).and_return(graphql_subscription_double)
+ .with(user, issue, 'response', options: options, save_message: save_message)
+ .and_return(graphql_subscription_double)
- described_class.new(context, basic_options).execute(response: 'response', options: options)
+ described_class.new(context, basic_options)
+ .execute(response: 'response', options: options, save_message: save_message)
end
end
end