diff --git a/ee/app/models/ai/conversation/message.rb b/ee/app/models/ai/conversation/message.rb
index 76acf37bc1bc7cad570a371103ccb0a239fe30cb..c9d9ceee576f83396792541ba99a610255bd9801 100644
--- a/ee/app/models/ai/conversation/message.rb
+++ b/ee/app/models/ai/conversation/message.rb
@@ -11,6 +11,7 @@ class Message < ApplicationRecord
validates :content, :role, :thread_id, presence: true
scope :for_thread, ->(thread) { where(thread: thread) }
+ scope :for_user, ->(user) { joins(:thread).where(ai_conversation_threads: { user_id: user.id }) }
# This message_xid is a secure random ID that is generated in runtime.
# https://gitlab.com/gitlab-org/gitlab/-/blob/master/ee/lib/gitlab/llm/ai_message.rb#L47
scope :for_message_xid, ->(message_xid) { where(message_xid: message_xid) }
diff --git a/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb b/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb
index 09425b65c18f2d6811c3568dffd9c1af06e3ac14..81980b563287a496ee0d69c2d6913ded5063e427 100644
--- a/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb
@@ -71,8 +71,15 @@ def attributes(response)
data.merge(Gitlab::Llm::ChatMessageAnalyzer.new(messages).execute)
end
+ def valid?
+ messages.present?
+ rescue ActiveRecord::RecordNotFound
+ false
+ end
+
def messages
- ::Gitlab::Llm::ChatStorage.new(user).messages_up_to(options[:message_id])
+ message = ::Ai::Conversation::Message.for_user(user).for_message_xid(options[:message_id]).first! # rubocop:disable CodeReuse/ActiveRecord -- not sure why first is allowed but not first!
+ ::Gitlab::Llm::ChatStorage.new(user, nil, message.thread).messages_up_to(options[:message_id])
end
strong_memoize_attr :messages
diff --git a/ee/lib/gitlab/llm/ai_message.rb b/ee/lib/gitlab/llm/ai_message.rb
index d21ff13141c47d068a9ffdf355cc805664bd2f31..a873765a35f6d35cc921b82f4833cee5557b0ee7 100644
--- a/ee/lib/gitlab/llm/ai_message.rb
+++ b/ee/lib/gitlab/llm/ai_message.rb
@@ -111,6 +111,10 @@ def ==(other)
@id == other.id
)
end
+
+ def thread_id
+ thread&.id
+ end
end
end
end
diff --git a/ee/lib/gitlab/llm/chat_message.rb b/ee/lib/gitlab/llm/chat_message.rb
index 715d4c9a645fada5d2d2879f1784805a5806f0f4..7a6cad696ece613d0f9ea1ab62bff1a3b0eb0328 100644
--- a/ee/lib/gitlab/llm/chat_message.rb
+++ b/ee/lib/gitlab/llm/chat_message.rb
@@ -6,13 +6,15 @@ class ChatMessage < AiMessage
RESET_MESSAGE = '/reset'
CLEAR_HISTORY_MESSAGE = '/clear'
+ attr_writer :active_record
+
def save!
storage = ChatStorage.new(user, agent_version_id, thread)
if content == CLEAR_HISTORY_MESSAGE
storage.clear!
else
- storage.add(self)
+ @active_record = storage.add(self)
end
self.thread = storage.current_thread
@@ -30,6 +32,10 @@ def question?
user? && !conversation_reset? && !clear_history?
end
+ def active_record
+ @active_record ||= ::Ai::Conversation::Message.for_user(user).for_message_xid(id).first
+ end
+
def chat?
true
end
diff --git a/ee/lib/gitlab/llm/chat_storage.rb b/ee/lib/gitlab/llm/chat_storage.rb
index a943caac84d9ae06d9f5e990d48e061ac88ed2a8..ef44ad534bf0a68b27b0afd2c5c0410441b1b8ef 100644
--- a/ee/lib/gitlab/llm/chat_storage.rb
+++ b/ee/lib/gitlab/llm/chat_storage.rb
@@ -18,8 +18,8 @@ def initialize(user, agent_version_id = nil, thread = nil)
end
def add(message)
- postgres_storage.add(message)
redis_storage.add(message) if ::Feature.disabled?(:duo_chat_drop_redis_storage, user)
+ postgres_storage.add(message)
end
def update_message_extras(request_id, key, value)
diff --git a/ee/lib/gitlab/llm/chat_storage/postgresql.rb b/ee/lib/gitlab/llm/chat_storage/postgresql.rb
index 2de46f6192e3b92a58580bc7d21dfdb42076af02..b44e819406806cf8447352b02c1144d1eb33ad4b 100644
--- a/ee/lib/gitlab/llm/chat_storage/postgresql.rb
+++ b/ee/lib/gitlab/llm/chat_storage/postgresql.rb
@@ -15,9 +15,11 @@ def add(message)
data['request_xid'] = data.delete('request_id') if data['request_id']
data.delete('timestamp') if data['timestamp']
- current_thread.messages.create!(**data)
+ result = current_thread.messages.create!(**data)
current_thread.update_column(:last_updated_at, Time.current)
clear_memoization(:messages)
+
+ result
end
def set_has_feedback(message)
@@ -37,6 +39,8 @@ def messages
msg = load_message(data)
msg.extras['has_feedback'] = data.delete('has_feedback') if data['has_feedback']
+ msg.active_record = message
+ msg.thread = current_thread
msg
end
diff --git a/ee/spec/lib/gitlab/llm/ai_gateway/completions/categorize_question_spec.rb b/ee/spec/lib/gitlab/llm/ai_gateway/completions/categorize_question_spec.rb
index c4ced9ce7808f8781c9f4746c892eaff7f274a65..c2b20d3f91ee04ba2bd151ebe3d576dd05907997 100644
--- a/ee/spec/lib/gitlab/llm/ai_gateway/completions/categorize_question_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_gateway/completions/categorize_question_spec.rb
@@ -3,17 +3,26 @@
require 'spec_helper'
RSpec.describe Gitlab::Llm::AiGateway::Completions::CategorizeQuestion, feature_category: :duo_chat do
+ let_it_be(:organization) { create(:organization) }
let_it_be(:user) { create(:user) }
+ let_it_be(:ai_conversation_thread) { create(:ai_conversation_thread, user: user) }
+ let_it_be(:ai_conversation_message) { create(:ai_conversation_message, thread: ai_conversation_thread) }
+ let_it_be(:question) { 'What is the pipeline?' }
+ let(:chat_message) do
+ message = build(:ai_chat_message, content: question, id: ai_conversation_message.message_xid)
+ message.active_record = ai_conversation_message
+ message.thread = ai_conversation_thread
+ message
+ end
let(:template_class) { ::Gitlab::Llm::Templates::CategorizeQuestion }
let(:ai_client) { instance_double(Gitlab::Llm::AiGateway::Client) }
let(:ai_response) { instance_double(HTTParty::Response, body: llm_analysis_response.to_json, success?: true) }
let(:uuid) { SecureRandom.uuid }
let(:tracking_context) { { action: :categorize_question, request_id: uuid } }
- let(:question) { 'What is the pipeline?' }
- let(:chat_message) { build(:ai_chat_message, content: question) }
let(:messages) { [chat_message] }
- let(:ai_options) { { question: chat_message.content, message_id: chat_message.id } }
+ let(:message_id) { chat_message.id }
+ let(:ai_options) { { question: chat_message.content, message_id: message_id } }
let(:prompt_message) { build(:ai_message, :categorize_question, user: user, request_id: uuid) }
let(:llm_analysis_response) do
{
@@ -26,7 +35,7 @@
end
before do
- allow_next_instance_of(::Gitlab::Llm::ChatStorage, user) do |storage|
+ allow_next_instance_of(::Gitlab::Llm::ChatStorage, user, nil, chat_message.thread) do |storage|
allow(storage).to receive(:messages_up_to).with(chat_message.id).and_return(messages)
end
end
@@ -150,5 +159,14 @@ def expect_client
)
end
end
+
+ context 'when message_id no longer exists' do
+ let(:message_id) { nil }
+
+ it 'raises error' do
+ expect(Gitlab::Llm::AiGateway::Client).not_to receive(:new)
+ expect(execute).to be_nil
+ end
+ end
end
end
diff --git a/ee/spec/lib/gitlab/llm/ai_message_spec.rb b/ee/spec/lib/gitlab/llm/ai_message_spec.rb
index ff183746521c9fc9cc4e78a21a0ab0d330ca027d..55ce581ec153d64323b39df82b92d72767b328a4 100644
--- a/ee/spec/lib/gitlab/llm/ai_message_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_message_spec.rb
@@ -229,4 +229,20 @@
expect(subject).not_to be_chat
end
end
+
+ describe '#thread_id' do
+ context 'when thread is present' do
+ it 'returns the thread id' do
+ expect(subject.thread_id).to eq(thread.id)
+ end
+ end
+
+ context 'when thread is nil' do
+ it 'returns nil' do
+ subject.thread = nil
+
+ expect(subject.thread_id).to be_nil
+ end
+ end
+ end
end
diff --git a/ee/spec/lib/gitlab/llm/chat_message_spec.rb b/ee/spec/lib/gitlab/llm/chat_message_spec.rb
index d70404f7b8ed734c6e8986ee1d2fad9f4e278254..19ef767f6ecd75cd0ade7e996857ba0b87c8120f 100644
--- a/ee/spec/lib/gitlab/llm/chat_message_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chat_message_spec.rb
@@ -115,4 +115,24 @@
expect(subject).to be_chat
end
end
+
+ describe '#active_record' do
+ it 'returns the active record if assigned' do
+ subject.active_record = build_stubbed(:ai_conversation_message)
+
+ expect(subject.active_record).to be_a(Ai::Conversation::Message)
+ end
+
+ it 'returns the active record if saved' do
+ subject.save!
+
+ expect(subject.active_record).to be_a(Ai::Conversation::Message)
+ end
+
+ it 'returns nil if not saved' do
+ subject
+
+ expect(subject.active_record).to be_nil
+ end
+ end
end
diff --git a/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb b/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb
index 070d3daebf39ab7fd12452fab400d5834d9ade36..c6981c905c9f9bdad8d7671991c98629cc15d490 100644
--- a/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb
@@ -36,7 +36,7 @@
allow(SecureRandom).to receive(:uuid).and_return(uuid)
expect(storage.messages).to be_empty
- storage.add(message)
+ result = storage.add(message)
last = storage.messages.last
expect(last.id).to eq(uuid)
@@ -50,6 +50,9 @@
expect(last.timestamp).not_to be_nil
expect(last.referer_url).to eq('http://127.0.0.1:3000')
expect(last.extras['additional_context']).to eq(payload[:additional_context].to_a)
+
+ expect(result).to be_a(Ai::Conversation::Message)
+ expect(result.message_xid).to eq(last.id)
end
context 'when the content exceeds the text limit' do
diff --git a/ee/spec/lib/gitlab/llm/chat_storage_spec.rb b/ee/spec/lib/gitlab/llm/chat_storage_spec.rb
index 357fef3297812ae8b7b98b980c09b3916dcc96db..ad3c9c9b97d338fe5f77142a5437dd285ef155ba 100644
--- a/ee/spec/lib/gitlab/llm/chat_storage_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chat_storage_spec.rb
@@ -34,10 +34,11 @@
describe '#add' do
it 'stores the message in PostgreSQL' do
- subject.add(message)
+ result = subject.add(message)
expect(postgres_storage.messages).to include(message)
expect(redis_storage.messages).to be_empty
+ expect(result).to eq(Ai::Conversation::Message.last)
end
context 'when feature flag duo_chat_drop_redis_storage is disabled' do
@@ -46,9 +47,10 @@
end
it 'updates Redis storage as well' do
- subject.add(message)
+ result = subject.add(message)
expect(redis_storage.messages).to include(message)
+ expect(result).to eq(Ai::Conversation::Message.last)
end
end
end
diff --git a/ee/spec/models/ai/conversation/message_spec.rb b/ee/spec/models/ai/conversation/message_spec.rb
index b026413d810eb1128b79a055473283edbc9feee1..70b73b78596170916a82905cb0d58424a1a20e00 100644
--- a/ee/spec/models/ai/conversation/message_spec.rb
+++ b/ee/spec/models/ai/conversation/message_spec.rb
@@ -53,6 +53,20 @@
expect(messages).to eq([message1, message2])
end
end
+
+ describe '.for_user' do
+ let_it_be(:user) { create(:user) }
+ let_it_be(:thread) { create(:ai_conversation_thread, user: user) }
+
+ let_it_be(:message) { create(:ai_conversation_message, thread: thread, role: :user) }
+ let_it_be(:message_from_other_user) { create(:ai_conversation_message, role: :user) }
+
+ it 'returns messages readable by the user' do
+ messages = described_class.for_user(user)
+
+ expect(messages).to contain_exactly(message)
+ end
+ end
end
describe '.recent' do