diff --git a/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb b/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb
index 0e4e69d83e55b0adf54412e522d39e48b28bee41..45bd49ec8a4e3eef4d7a67dffe2b21ec9ab6add7 100644
--- a/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb
+++ b/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb
@@ -16,7 +16,7 @@ class ChatMessagesResolver < BaseResolver
def resolve(**args)
return [] unless current_user
- ::Gitlab::Llm::ChatStorage.new(current_user).messages(args).map(&:to_h)
+ ::Gitlab::Llm::ChatStorage.new(current_user).messages_by(args).map(&:to_h)
end
end
end
diff --git a/ee/lib/gitlab/llm/ai_message.rb b/ee/lib/gitlab/llm/ai_message.rb
index 6897dbdcc31523788551ebf4c7c67412e008bb4a..0bf1d8e4471e5b948452b53bac3b88f3a9281422 100644
--- a/ee/lib/gitlab/llm/ai_message.rb
+++ b/ee/lib/gitlab/llm/ai_message.rb
@@ -63,6 +63,10 @@ def user?
role == ROLE_USER
end
+ def assistant?
+ role == ROLE_ASSISTANT
+ end
+
def slash_command?
content.to_s.match?(%r{\A/\w})
end
@@ -72,6 +76,14 @@ def slash_command_and_input
content.split(' ', 2)
end
+
+ def ==(other)
+ super || (
+ self.class == other.class &&
+ !@id.nil? &&
+ @id == other.id
+ )
+ end
end
end
end
diff --git a/ee/lib/gitlab/llm/chat_message.rb b/ee/lib/gitlab/llm/chat_message.rb
index c362d9d6c9f91315e191ef547d6ac672fa2d641b..b4c138dc8ee4cb507ec31bc4d437aff8e2d2d21d 100644
--- a/ee/lib/gitlab/llm/chat_message.rb
+++ b/ee/lib/gitlab/llm/chat_message.rb
@@ -24,6 +24,10 @@ def conversation_reset?
def clean_history?
content == CLEAN_HISTORY_MESSAGE
end
+
+ def question?
+ user? && !conversation_reset? && !clean_history?
+ end
end
end
end
diff --git a/ee/lib/gitlab/llm/chat_storage.rb b/ee/lib/gitlab/llm/chat_storage.rb
index 0355ab00d96bf7e877fcd1a037663d241b0c7f8e..275fa4732419d3691473aa20d82b66437ce44d55 100644
--- a/ee/lib/gitlab/llm/chat_storage.rb
+++ b/ee/lib/gitlab/llm/chat_storage.rb
@@ -3,6 +3,8 @@
module Gitlab
module Llm
class ChatStorage
+ include Gitlab::Utils::StrongMemoize
+
# Expiration time of user messages should not be more than 90 days.
# EXPIRE_TIME sets expiration time for the whole chat history stream (not
# for individual messages) - so the stream is deleted after 3 days since
@@ -25,15 +27,23 @@ def initialize(user)
def add(message)
cache_data(dump_message(message))
+ clear_memoization(:messages)
end
- def messages(filters = {})
+ def messages
with_redis do |redis|
- redis.xrange(key).filter_map do |_id, data|
- load_message(data) if matches_filters?(data, filters)
+ redis.xrange(key).map do |_id, data|
+ load_message(data)
end
end
end
+ strong_memoize_attr :messages
+
+ def messages_by(filters = {})
+ messages.select do |message|
+ matches_filters?(message, filters)
+ end
+ end
def last_conversation
all = messages
@@ -44,10 +54,17 @@ def last_conversation
all[idx + 1..]
end
+ def messages_up_to(message_id)
+ all = messages
+ idx = all.rindex { |m| m.id == message_id }
+ idx ? all.first(idx + 1) : []
+ end
+
def clean!
with_redis do |redis|
redis.xtrim(key, 0)
end
+ clear_memoization(:messages)
end
private
@@ -69,11 +86,11 @@ def with_redis(&block)
Gitlab::Redis::Chat.with(&block) # rubocop: disable CodeReuse/ActiveRecord
end
- def matches_filters?(data, filters)
- return false if filters[:roles]&.exclude?(data['role'])
- return false if filters[:request_ids]&.exclude?(data['request_id'])
+ def matches_filters?(message, filters)
+ return false if filters[:roles]&.exclude?(message.role)
+ return false if filters[:request_ids]&.exclude?(message.request_id)
- data
+ true
end
def dump_message(message)
diff --git a/ee/spec/lib/gitlab/llm/ai_message_spec.rb b/ee/spec/lib/gitlab/llm/ai_message_spec.rb
index a3288b0a5c3d1aa3410e878d33d8b3304f1f3cb1..3ac7f6b13d48baff46b05407dd83978551fdfb92 100644
--- a/ee/spec/lib/gitlab/llm/ai_message_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_message_spec.rb
@@ -135,9 +135,56 @@
end
end
+ describe '#assistant?' do
+ context 'when role is assistant' do
+ before do
+ data[:role] = 'assistant'
+ end
+
+ it { is_expected.to be_assistant }
+ end
+
+ context 'when role is not assistant' do
+ it { is_expected.not_to be_assistant }
+ end
+ end
+
describe '#resource' do
it 'delegates to context' do
expect(subject.resource).to eq(data[:context].resource)
end
end
+
+ describe '#==' do
+ it 'returns true if comparing to self' do
+ m1 = build(:ai_message)
+
+ expect(m1).to eq(m1)
+ end
+
+ it 'compares id only' do
+ m1 = build(:ai_chat_message, content: 'foobarbaz')
+ m2 = build(:ai_chat_message, content: 'foobar...')
+
+ expect(m1).not_to eq(m2)
+
+ m2.id = m1.id
+
+ expect(m1).to eq(m2)
+ end
+
+ it 'returns false if class is different' do
+ m1 = build(:ai_message, content: 'foobar')
+ m2 = build(:ai_chat_message, content: 'foobar', id: m1.id)
+
+ expect(m1).not_to eq(m2)
+ end
+
+ it 'returns false if id is nil' do
+ m1 = build(:ai_message, id: nil)
+ m2 = build(:ai_message, id: nil)
+
+ expect(m1).not_to eq(m2)
+ 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 c9ae0376178f8b1bd4eabc45d679d88da96d02e7..fbe1a852bef72b3a148b19191231620e1f91cd72 100644
--- a/ee/spec/lib/gitlab/llm/chat_message_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chat_message_spec.rb
@@ -25,6 +25,25 @@
end
end
+ describe '#question?' do
+ where(:role, :content, :expectation) do
+ [
+ ['user', 'foo?', true],
+ ['user', '/reset', false],
+ ['user', '/clean', false],
+ ['assistant', 'foo?', false]
+ ]
+ end
+
+ with_them do
+ it "returns expectation" do
+ subject.assign_attributes(role: role, content: content)
+
+ expect(subject.question?).to eq(expectation)
+ end
+ end
+ end
+
describe '#save!' do
it 'saves the message to chat storage' do
expect_next_instance_of(Gitlab::Llm::ChatStorage, subject.user) do |instance|
diff --git a/ee/spec/lib/gitlab/llm/chat_storage_spec.rb b/ee/spec/lib/gitlab/llm/chat_storage_spec.rb
index 47b7102d41d6d66be040882ef24da5d5754773bb..beddd20784d6f046de6efad8df913d6f56890640 100644
--- a/ee/spec/lib/gitlab/llm/chat_storage_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chat_storage_spec.rb
@@ -66,6 +66,18 @@
end
describe '#messages' do
+ before do
+ subject.add(build(:ai_chat_message, payload.merge(content: 'msg1', role: 'user', request_id: '1')))
+ subject.add(build(:ai_chat_message, payload.merge(content: 'msg2', role: 'assistant', request_id: '2')))
+ subject.add(build(:ai_chat_message, payload.merge(content: 'msg3', role: 'assistant', request_id: '3')))
+ end
+
+ it 'returns all records for this user' do
+ expect(subject.messages.map(&:content)).to eq(%w[msg1 msg2 msg3])
+ end
+ end
+
+ describe '#messages_by' do
let(:filters) { {} }
before do
@@ -75,14 +87,14 @@
end
it 'returns all records for this user' do
- expect(subject.messages(filters).map(&:content)).to eq(%w[msg1 msg2 msg3])
+ expect(subject.messages_by(filters).map(&:content)).to eq(%w[msg1 msg2 msg3])
end
context 'when filtering by role' do
let(:filters) { { roles: ['user'] } }
it 'returns only records for this role' do
- expect(subject.messages(filters).map(&:content)).to eq(%w[msg1])
+ expect(subject.messages_by(filters).map(&:content)).to eq(%w[msg1])
end
end
@@ -90,7 +102,7 @@
let(:filters) { { request_ids: %w[2 3] } }
it 'returns only records with the same request_id' do
- expect(subject.messages(filters).map(&:content)).to eq(%w[msg2 msg3])
+ expect(subject.messages_by(filters).map(&:content)).to eq(%w[msg2 msg3])
end
end
end
@@ -150,4 +162,26 @@
expect(subject.messages).to be_empty
end
end
+
+ describe '#messages_up_to' do
+ let(:messages) do
+ [
+ build(:ai_chat_message, payload.merge(content: 'msg1', role: 'assistant')),
+ build(:ai_chat_message, payload.merge(content: 'msg2', role: 'user')),
+ build(:ai_chat_message, payload.merge(content: 'msg3', role: 'assistant'))
+ ]
+ end
+
+ before do
+ messages.each { |m| subject.add(m) }
+ end
+
+ it 'returns first n messages up to one with matching message id' do
+ expect(subject.messages_up_to(messages[1].id)).to eq(messages.first(2))
+ end
+
+ it 'returns [] if message id is not found' do
+ expect(subject.messages_up_to('missing id')).to eq([])
+ end
+ end
end