Rename Llm::Cache classes

Renames `Llm::Cache` to `ChatStorage + ChatMessage` classes. 

ee/app/graphql/resolvers/ai/messages_resolver.rb

@@ -16,7 +16,7 @@ class MessagesResolver < BaseResolver
       def resolve(**args)
         return [] unless current_user
-        ::Gitlab::Llm::Cache.new(current_user).find_all(args)
+        ::Gitlab::Llm::ChatStorage.new(current_user).messages(args)
       end
     end
   end

ee/app/graphql/types/ai/chat_message_role_enum.rb

@@ -6,7 +6,7 @@ class ChatMessageRoleEnum < BaseEnum
       graphql_name 'AiChatMessageRole'
       description 'Roles to filter in chat message.'
-      ::Gitlab::Llm::Cache::ALLOWED_ROLES.each do |role|
+      ::Gitlab::Llm::ChatMessage::ALLOWED_ROLES.each do |role|
         value role.upcase, description: "Filter only #{role} messages.", value: role
       end
     end

ee/app/services/llm/base_service.rb

@@ -40,12 +40,12 @@ def worker_perform(user, resource, action_name, options)
       message = content(action_name)
       payload = {
         request_id: request_id,
-        role: ::Gitlab::Llm::Cache::ROLE_USER,
+        role: ::Gitlab::Llm::ChatMessage::ROLE_USER,
         content: message,
         timestamp: Time.current
       }
-      ::Gitlab::Llm::Cache.new(user).add(payload) if cache_response?(options)
+      ::Gitlab::Llm::ChatStorage.new(user).add(payload) if cache_response?(options)
       if emit_response?(options)
         # We do not add the `client_subscription_id` here on purpose for now.
@@ -106,7 +106,7 @@ def content(action_name)
     end
     def no_worker_message?(content)
-      content == ::Gitlab::Llm::CachedMessage::RESET_MESSAGE
+      content == ::Gitlab::Llm::ChatMessage::RESET_MESSAGE
     end
     def cache_response?(options)

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

@@ -92,7 +92,7 @@ def execute_streamed_request
                     response: Gitlab::Llm::Chain::PlainResponseModifier.new(content),
                     options: {
                       cache_response: false,
-                      role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+                      role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
                       chunk_id: chunk[:id]
                     }
                   )
@@ -134,7 +134,8 @@ def picked_tool_action(tool_class)
               response_handler.execute(
                 response: Gitlab::Llm::Chain::ToolResponseModifier.new(tool_class),
-                options: { cache_response: false, role: ::Gitlab::Llm::Cache::ROLE_SYSTEM, type: RESPONSE_TYPE_TOOL }
+                options: { cache_response: false, role: ::Gitlab::Llm::ChatMessage::ROLE_SYSTEM,
+                           type: RESPONSE_TYPE_TOOL }
               )
             end
@@ -143,7 +144,7 @@ def prompt_version
             end
             def last_conversation
-              Cache.new(context.current_user).last_conversation
+              ChatStorage.new(context.current_user).last_conversation
             end
             strong_memoize_attr :last_conversation

ee/lib/gitlab/llm/chain/agents/zero_shot/prompts/anthropic.rb

@@ -7,14 +7,14 @@ module Agents
         module ZeroShot
           module Prompts
             ROLE_NAMES = {
-              Llm::Cache::ROLE_USER => 'Human',
-              Llm::Cache::ROLE_ASSISTANT => 'Assistant'
+              Llm::ChatMessage::ROLE_USER => 'Human',
+              Llm::ChatMessage::ROLE_ASSISTANT => 'Assistant'
             }.freeze
             class Anthropic < Base
               def self.prompt(options)
                 text = <<~PROMPT
-                  #{ROLE_NAMES[Llm::Cache::ROLE_USER]}: #{base_prompt(options)}
+                  #{ROLE_NAMES[Llm::ChatMessage::ROLE_USER]}: #{base_prompt(options)}
                 PROMPT
                 history = truncated_conversation(options[:conversation], Requests::Anthropic::PROMPT_SIZE - text.size)

ee/lib/gitlab/llm/cached_message.rb → ee/lib/gitlab/llm/chat_message.rb

@@ -2,7 +2,12 @@
 module Gitlab
   module Llm
-    class CachedMessage
+    class ChatMessage
+      ROLE_USER = 'user'
+      ROLE_ASSISTANT = 'assistant'
+      ROLE_SYSTEM = 'system'
+      ALLOWED_ROLES = [ROLE_USER, ROLE_ASSISTANT, ROLE_SYSTEM].freeze
+
       attr_reader :id, :request_id, :content, :role, :timestamp, :error
       RESET_MESSAGE = '/reset'

ee/lib/gitlab/llm/cache.rb → ee/lib/gitlab/llm/chat_storage.rb

@@ -2,7 +2,7 @@
 module Gitlab
   module Llm
-    class Cache
+    class ChatStorage
       # 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
@@ -19,17 +19,14 @@ class Cache
       # sufficient.
       MAX_TEXT_LIMIT = 20_000
-      ROLE_USER = 'user'
-      ROLE_ASSISTANT = 'assistant'
-      ROLE_SYSTEM = 'system'
-      ALLOWED_ROLES = [ROLE_USER, ROLE_ASSISTANT, ROLE_SYSTEM].freeze
-
       def initialize(user)
         @user = user
       end
       def add(payload)
-        raise ArgumentError, "Invalid role '#{payload[:role]}'" unless ALLOWED_ROLES.include?(payload[:role])
+        unless ChatMessage::ALLOWED_ROLES.include?(payload[:role])
+          raise ArgumentError, "Invalid role '#{payload[:role]}'"
+        end
         data = {
           id: SecureRandom.uuid,
@@ -43,16 +40,16 @@ def add(payload)
         cache_data(data)
       end
-      def find_all(filters = {})
+      def messages(filters = {})
         with_redis do |redis|
           redis.xrange(key).filter_map do |_id, data|
-            CachedMessage.new(data) if matches_filters?(data, filters)
+            ChatMessage.new(data) if matches_filters?(data, filters)
           end
         end
       end
       def last_conversation
-        all = find_all
+        all = messages
         idx = all.rindex(&:conversation_reset?)
         return all unless idx
         return [] unless idx + 1 < all.size

ee/lib/gitlab/llm/graphql_subscription_response_service.rb

@@ -20,7 +20,7 @@ def execute
           model_name: resource&.class&.name,
           content: response_modifier.response_body,
           errors: response_modifier.errors,
-          role: options[:role] || Cache::ROLE_ASSISTANT,
+          role: options[:role] || ChatMessage::ROLE_ASSISTANT,
           timestamp: Time.current,
           type: options.fetch(:type, nil),
           chunk_id: options.fetch(:chunk_id, nil)
@@ -35,7 +35,7 @@ def execute
         response_data = data.slice(:request_id, :errors, :role, :content, :timestamp)
         unless options[:internal_request]
-          Gitlab::Llm::Cache.new(user).add(response_data) if options[:cache_response]
+          Gitlab::Llm::ChatStorage.new(user).add(response_data) if options[:cache_response]
           subscription_arguments = { user_id: user.to_global_id, resource_id: resource&.to_global_id }
           if options[:client_subscription_id]

ee/spec/graphql/types/ai/chat_message_type_spec.rb

@@ -16,10 +16,10 @@
   describe '#content_html' do
     let_it_be(:current_user) { create(:user) }
-    let(:cached_message) { Gitlab::Llm::CachedMessage.new('content' => "Hello, **World**!", 'timestamp' => '') }
+    let(:message) { Gitlab::Llm::ChatMessage.new('content' => "Hello, **World**!", 'timestamp' => '') }
     it 'renders html through Banzai' do
-      allow(Banzai).to receive(:render_and_post_process).with(cached_message.content, {
+      allow(Banzai).to receive(:render_and_post_process).with(message.content, {
         current_user: current_user,
         only_path: false,
         pipeline: :full,
@@ -27,7 +27,7 @@
         skip_project_check: true
       }).and_return('banzai_content')
-      resolved_field = resolve_field(:content_html, cached_message, current_user: current_user)
+      resolved_field = resolve_field(:content_html, message, current_user: current_user)
       expect(resolved_field).to eq('banzai_content')
     end

ee/spec/lib/gitlab/llm/chain/agents/zero_shot/executor_spec.rb

@@ -151,11 +151,11 @@
         expect(stream_response_service_double).to receive(:execute).with(
           response: first_response_double,
-          options: { cache_response: false, role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT, chunk_id: 1 }
+          options: { cache_response: false, role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT, chunk_id: 1 }
         )
         expect(stream_response_service_double).to receive(:execute).with(
           response: second_response_double,
-          options: { cache_response: false, role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT, chunk_id: 2 }
+          options: { cache_response: false, role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT, chunk_id: 2 }
         )
         agent.execute
@@ -168,26 +168,26 @@
       allow(agent).to receive(:provider_prompt_class)
                         .and_return(Gitlab::Llm::Chain::Agents::ZeroShot::Prompts::Anthropic)
-      Gitlab::Llm::Cache.new(user).add(request_id: 'uuid1', role: 'user', content: 'question 1')
-      Gitlab::Llm::Cache.new(user).add(request_id: 'uuid1', role: 'assistant', content: 'response 1')
+      Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid1', role: 'user', content: 'question 1')
+      Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid1', role: 'assistant', content: 'response 1')
       # this should be ignored because response contains an error
-      Gitlab::Llm::Cache.new(user).add(request_id: 'uuid2', role: 'user', content: 'question 2')
-      Gitlab::Llm::Cache.new(user)
+      Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid2', role: 'user', content: 'question 2')
+      Gitlab::Llm::ChatStorage.new(user)
                         .add(request_id: 'uuid2', role: 'assistant', content: 'response 2', errors: ['error'])
       # this should be ignored because it doesn't contain response
-      Gitlab::Llm::Cache.new(user).add(request_id: 'uuid3', role: 'user', content: 'question 3')
+      Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid3', role: 'user', content: 'question 3')
       travel(2.minutes) do
-        Gitlab::Llm::Cache.new(user).add(request_id: 'uuid4', role: 'user', content: 'question 4')
+        Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid4', role: 'user', content: 'question 4')
       end
       travel(2.minutes) do
-        Gitlab::Llm::Cache.new(user).add(request_id: 'uuid5', role: 'user', content: 'question 5')
+        Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid5', role: 'user', content: 'question 5')
       end
       travel(3.minutes) do
-        Gitlab::Llm::Cache.new(user).add(request_id: 'uuid4', role: 'assistant', content: 'response 4')
+        Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid4', role: 'assistant', content: 'response 4')
       end
       travel(4.minutes) do
-        Gitlab::Llm::Cache.new(user).add(request_id: 'uuid5', role: 'assistant', content: 'response 5')
+        Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid5', role: 'assistant', content: 'response 5')
       end
     end
@@ -414,7 +414,8 @@
           uuid = SecureRandom.uuid
           history.each do |message|
-            Gitlab::Llm::Cache.new(user).add({ request_id: uuid, role: message[:role], content: message[:content] })
+            Gitlab::Llm::ChatStorage.new(user).add({ request_id: uuid, role: message[:role],
+content: message[:content] })
           end
         end
@@ -542,7 +543,8 @@
             uuid = SecureRandom.uuid
             history.each do |message|
-              Gitlab::Llm::Cache.new(user).add({ request_id: uuid, role: message[:role], content: message[:content] })
+              Gitlab::Llm::ChatStorage.new(user).add({ request_id: uuid, role: message[:role],
+content: message[:content] })
             end
           end

ee/spec/lib/gitlab/llm/chain/agents/zero_shot/prompts/anthropic_spec.rb

@@ -12,16 +12,16 @@
         user_input: 'foo?',
         agent_scratchpad: "some observation",
         conversation: [
-          Gitlab::Llm::CachedMessage.new(
+          Gitlab::Llm::ChatMessage.new(
             'request_id' => 'uuid1', 'role' => 'user', 'content' => 'question 1', 'timestamp' => Time.current.to_s
           ),
-          Gitlab::Llm::CachedMessage.new(
+          Gitlab::Llm::ChatMessage.new(
             'request_id' => 'uuid1', 'role' => 'assistant', 'content' => 'response 1', 'timestamp' => Time.current.to_s
           ),
-          Gitlab::Llm::CachedMessage.new(
+          Gitlab::Llm::ChatMessage.new(
             'request_id' => 'uuid1', 'role' => 'user', 'content' => 'question 2', 'timestamp' => Time.current.to_s
           ),
-          Gitlab::Llm::CachedMessage.new(
+          Gitlab::Llm::ChatMessage.new(
             'request_id' => 'uuid1', 'role' => 'assistant', 'content' => 'response 2', 'timestamp' => Time.current.to_s
           )
         ],

ee/spec/lib/gitlab/llm/cached_message_spec.rb → ee/spec/lib/gitlab/llm/chat_message_spec.rb

@@ -2,7 +2,7 @@
 require 'spec_helper'
-RSpec.describe Gitlab::Llm::CachedMessage, feature_category: :duo_chat do
+RSpec.describe Gitlab::Llm::ChatMessage, feature_category: :duo_chat do
   let(:timestamp) { Time.current }
   let(:data) do
     {
@@ -19,7 +19,7 @@
   describe '#to_global_id' do
     it 'returns global ID' do
-      expect(subject.to_global_id.to_s).to eq('gid://gitlab/Gitlab::Llm::CachedMessage/uuid')
+      expect(subject.to_global_id.to_s).to eq('gid://gitlab/Gitlab::Llm::ChatMessage/uuid')
     end
   end

ee/spec/lib/gitlab/llm/cache_spec.rb → ee/spec/lib/gitlab/llm/chat_storage_spec.rb

@@ -2,7 +2,7 @@
 require 'spec_helper'
-RSpec.describe Gitlab::Llm::Cache, :clean_gitlab_redis_chat, feature_category: :duo_chat do
+RSpec.describe Gitlab::Llm::ChatStorage, :clean_gitlab_redis_chat, feature_category: :duo_chat do
   let_it_be(:user) { create(:user) }
   let(:request_id) { 'uuid' }
   let(:timestamp) { Time.current.to_s }
@@ -30,11 +30,11 @@
       uuid = 'unique_id'
       expect(SecureRandom).to receive(:uuid).once.and_return(uuid)
-      expect(subject.find_all).to be_empty
+      expect(subject.messages).to be_empty
       subject.add(payload)
-      last = subject.find_all.last
+      last = subject.messages.last
       expect(last.id).to eq(uuid)
       expect(last.request_id).to eq(request_id)
       expect(last.errors).to eq(['some error1. another error'])
@@ -48,7 +48,7 @@
       subject.add(payload)
-      last = subject.find_all.last
+      last = subject.messages.last
       expect(last.errors).to eq([])
     end
@@ -66,23 +66,23 @@
     context 'with MAX_MESSAGES limit' do
       before do
-        stub_const('Gitlab::Llm::Cache::MAX_MESSAGES', 2)
+        stub_const('Gitlab::Llm::ChatStorage::MAX_MESSAGES', 2)
       end
       it 'removes oldes messages if we reach maximum message limit' do
         subject.add(payload.merge(content: 'msg1'))
         subject.add(payload.merge(content: 'msg2'))
-        expect(subject.find_all.map(&:content)).to eq(%w[msg1 msg2])
+        expect(subject.messages.map(&:content)).to eq(%w[msg1 msg2])
         subject.add(payload.merge(content: 'msg3'))
-        expect(subject.find_all.map(&:content)).to eq(%w[msg2 msg3])
+        expect(subject.messages.map(&:content)).to eq(%w[msg2 msg3])
       end
     end
   end
-  describe '#find_all' do
+  describe '#messages' do
     let(:filters) { {} }
     before do
@@ -92,14 +92,14 @@
     end
     it 'returns all records for this user' do
-      expect(subject.find_all(filters).map(&:content)).to eq(%w[msg1 msg2 msg3])
+      expect(subject.messages(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.find_all(filters).map(&:content)).to eq(%w[msg1])
+        expect(subject.messages(filters).map(&:content)).to eq(%w[msg1])
       end
     end
@@ -107,7 +107,7 @@
       let(:filters) { { request_ids: %w[2 3] } }
       it 'returns only records with the same request_id' do
-        expect(subject.find_all(filters).map(&:content)).to eq(%w[msg2 msg3])
+        expect(subject.messages(filters).map(&:content)).to eq(%w[msg2 msg3])
       end
     end
   end

ee/spec/lib/gitlab/llm/completions/explain_vulnerability_spec.rb

@@ -78,7 +78,7 @@
           id: anything,
           model_name: vulnerability.class.name,
           content: '',
-          role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+          role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
           request_id: 'uuid',
           errors: [described_class::NULL_PROMPT_ERROR]
         }))
@@ -103,7 +103,7 @@
             id: anything,
             model_name: vulnerability.class.name,
             content: '',
-            role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+            role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
             request_id: 'uuid',
             errors: errors[llm_client]
           }))
@@ -133,7 +133,7 @@
             id: anything,
             model_name: vulnerability.class.name,
             content: example_answer,
-            role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+            role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
             request_id: 'uuid',
             errors: []
           }))
@@ -162,7 +162,7 @@
               id: anything,
               model_name: vulnerability.class.name,
               content: '',
-              role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+              role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
               request_id: 'uuid',
               errors: [described_class::DEFAULT_ERROR]
             }))
@@ -192,7 +192,7 @@
               id: anything,
               model_name: vulnerability.class.name,
               content: '',
-              role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+              role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
               request_id: 'uuid',
               errors: [described_class::CLIENT_TIMEOUT_ERROR]
             }))
@@ -219,7 +219,7 @@
                 id: anything,
                 model_name: vulnerability.class.name,
                 content: example_answer,
-                role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+                role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
                 request_id: 'uuid',
                 errors: []
               })).twice
@@ -251,7 +251,7 @@
                   id: anything,
                   model_name: vulnerability.class.name,
                   content: example_answer,
-                  role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+                  role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
                   request_id: 'uuid',
                   errors: []
                 })

ee/spec/lib/gitlab/llm/graphql_subscription_response_service_spec.rb

@@ -86,7 +86,7 @@
       let(:cache_response) { true }
       it 'caches response' do
-        expect_next_instance_of(::Gitlab::Llm::Cache) do |cache|
+        expect_next_instance_of(::Gitlab::Llm::ChatStorage) do |cache|
           expect(cache).to receive(:add)
             .with(payload.slice(:request_id, :errors, :role, :timestamp).merge(content: payload[:content]))
         end
@@ -99,7 +99,7 @@
       let(:cache_response) { false }
       it 'does not cache the response' do
-        expect(Gitlab::Llm::Cache).not_to receive(:new)
+        expect(Gitlab::Llm::ChatStorage).not_to receive(:new)
         subject
       end
@@ -151,7 +151,7 @@
       it 'returns response but does not cache or broadcast' do
         expect(GraphqlTriggers).not_to receive(:ai_completion_response)
-        expect(Gitlab::Llm::Cache).not_to receive(:new)
+        expect(Gitlab::Llm::ChatStorage).not_to receive(:new)
         expect(subject[:content]).to eq(response_body)
       end

ee/spec/lib/gitlab/metrics/llm_spec.rb

@@ -30,7 +30,7 @@
     end
     it 'returns :unknwon for other classes' do
-      expect(described_class.client_label(Gitlab::Llm::Cache)).to eq(:unknown)
+      expect(described_class.client_label(Gitlab::Llm::ChatStorage)).to eq(:unknown)
     end
   end
 end

ee/spec/requests/api/graphql/ai_messages_spec.rb

@@ -36,10 +36,10 @@
   subject { graphql_data.dig('aiMessages', 'nodes') }
   before do
-    ::Gitlab::Llm::Cache.new(user).add(request_id: 'uuid1', role: 'user', content: 'question 1')
-    ::Gitlab::Llm::Cache.new(user).add(request_id: 'uuid1', role: 'assistant', content: response_content)
+    ::Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid1', role: 'user', content: 'question 1')
+    ::Gitlab::Llm::ChatStorage.new(user).add(request_id: 'uuid1', role: 'assistant', content: response_content)
     # should not be included in response because it's for other user
-    ::Gitlab::Llm::Cache.new(other_user).add(request_id: 'uuid1', role: 'user', content: 'question 2')
+    ::Gitlab::Llm::ChatStorage.new(other_user).add(request_id: 'uuid1', role: 'user', content: 'question 2')
   end
   context 'when user is not logged in' do

ee/spec/requests/api/graphql/subscriptions/ai_completion_response_spec.rb

@@ -41,7 +41,7 @@
         model_name: resource.class.name,
         request_id: request_id,
         content: content,
-        role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
+        role: ::Gitlab::Llm::ChatMessage::ROLE_ASSISTANT,
         errors: [],
         chunk_id: nil
       }

ee/spec/support/shared_examples/services/llm/llm_cache_shared_examples.rb

@@ -15,7 +15,7 @@
   end
   it 'caches response' do
-    expect_next_instance_of(::Gitlab::Llm::Cache) do |cache|
+    expect_next_instance_of(::Gitlab::Llm::ChatStorage) do |cache|
       expect(cache).to receive(:add).with(expected_cache_payload)
     end
@@ -32,7 +32,7 @@
     it 'only stores the message in cache' do
       expect(::Llm::CompletionWorker).not_to receive(:perform_async)
-      expect_next_instance_of(::Gitlab::Llm::Cache) do |cache|
+      expect_next_instance_of(::Gitlab::Llm::ChatStorage) do |cache|
         expect(cache).to receive(:add).with(expected_cache_payload)
       end
@@ -43,7 +43,7 @@
 RSpec.shared_examples 'llm service does not cache user request' do
   it 'does not cache the request' do
-    expect(::Gitlab::Llm::Cache).not_to receive(:new)
+    expect(::Gitlab::Llm::ChatStorage).not_to receive(:new)
     subject.execute
   end