diff --git a/doc/api/graphql/reference/index.md b/doc/api/graphql/reference/index.md
index c27075a9a539ad7445dc67ad7e57c189ad298077..84c6891f372693049fb3faccfb8c3e85c08f2331 100644
--- a/doc/api/graphql/reference/index.md
+++ b/doc/api/graphql/reference/index.md
@@ -36,6 +36,27 @@ in [Removed Items](../removed_items.md).
The `Query` type contains the API's top-level entry points for all executable queries.
+### `Query.aiMessages`
+
+Find AI messages.
+
+WARNING:
+**Introduced** in 16.1.
+This feature is an Experiment. It can be changed or removed at any time.
+
+Returns [`AiCachedMessageTypeConnection!`](#aicachedmessagetypeconnection).
+
+This field returns a [connection](#connections). It accepts the
+four standard [pagination arguments](#connection-pagination-arguments):
+`before: String`, `after: String`, `first: Int`, `last: Int`.
+
+#### Arguments
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="queryaimessagesrequestids"></a>`requestIds` | [`[ID!]`](#id) | Array of request IDs to fetch. |
+| <a id="queryaimessagesroles"></a>`roles` | [`[AiCachedMessageRole!]`](#aicachedmessagerole) | Array of request IDs to fetch. |
+
### `Query.boardList`
Find an issue board list.
@@ -7244,6 +7265,29 @@ The edge type for [`AgentConfiguration`](#agentconfiguration).
| <a id="agentconfigurationedgecursor"></a>`cursor` | [`String!`](#string) | A cursor for use in pagination. |
| <a id="agentconfigurationedgenode"></a>`node` | [`AgentConfiguration`](#agentconfiguration) | The item at the end of the edge. |
+#### `AiCachedMessageTypeConnection`
+
+The connection type for [`AiCachedMessageType`](#aicachedmessagetype).
+
+##### Fields
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="aicachedmessagetypeconnectionedges"></a>`edges` | [`[AiCachedMessageTypeEdge]`](#aicachedmessagetypeedge) | A list of edges. |
+| <a id="aicachedmessagetypeconnectionnodes"></a>`nodes` | [`[AiCachedMessageType]`](#aicachedmessagetype) | A list of nodes. |
+| <a id="aicachedmessagetypeconnectionpageinfo"></a>`pageInfo` | [`PageInfo!`](#pageinfo) | Information to aid in pagination. |
+
+#### `AiCachedMessageTypeEdge`
+
+The edge type for [`AiCachedMessageType`](#aicachedmessagetype).
+
+##### Fields
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="aicachedmessagetypeedgecursor"></a>`cursor` | [`String!`](#string) | A cursor for use in pagination. |
+| <a id="aicachedmessagetypeedgenode"></a>`node` | [`AiCachedMessageType`](#aicachedmessagetype) | The item at the end of the edge. |
+
#### `AiMessageTypeConnection`
The connection type for [`AiMessageType`](#aimessagetype).
@@ -11686,6 +11730,19 @@ Information about a connected Agent.
| <a id="agentmetadatapodnamespace"></a>`podNamespace` | [`String`](#string) | Namespace of the pod running the Agent. |
| <a id="agentmetadataversion"></a>`version` | [`String`](#string) | Agent version tag. |
+### `AiCachedMessageType`
+
+#### Fields
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="aicachedmessagetypecontent"></a>`content` | [`String`](#string) | Content of the message. Can be null for user requests or failed responses. |
+| <a id="aicachedmessagetypeerrors"></a>`errors` | [`[String!]!`](#string) | Errors that occurred while asynchronously fetching an AI(assistant) response. |
+| <a id="aicachedmessagetypeid"></a>`id` | [`ID`](#id) | UUID of the message. |
+| <a id="aicachedmessagetyperequestid"></a>`requestId` | [`ID`](#id) | UUID of the original request message. |
+| <a id="aicachedmessagetyperole"></a>`role` | [`AiCachedMessageRole!`](#aicachedmessagerole) | Message role. |
+| <a id="aicachedmessagetypetimestamp"></a>`timestamp` | [`Time!`](#time) | Message timestamp. |
+
### `AiMessageType`
#### Fields
@@ -23859,6 +23916,15 @@ Agent token statuses.
| <a id="agenttokenstatusactive"></a>`ACTIVE` | Active agent token. |
| <a id="agenttokenstatusrevoked"></a>`REVOKED` | Revoked agent token. |
+### `AiCachedMessageRole`
+
+Roles to filter in chat message.
+
+| Value | Description |
+| ----- | ----------- |
+| <a id="aicachedmessageroleassistant"></a>`ASSISTANT` | Filter only AI responses. |
+| <a id="aicachedmessageroleuser"></a>`USER` | Filter only user messages. |
+
### `AlertManagementAlertSort`
Values for sorting alerts.
diff --git a/ee/app/graphql/ee/types/query_type.rb b/ee/app/graphql/ee/types/query_type.rb
index c50fee320784a83a5162573144e0381492f26848..bf7ac0f4f7b756bc554b99a6f601ed3989ba1886 100644
--- a/ee/app/graphql/ee/types/query_type.rb
+++ b/ee/app/graphql/ee/types/query_type.rb
@@ -97,6 +97,11 @@ module QueryType
null: true,
description: 'Instance level external audit event destinations.',
resolver: ::Resolvers::AuditEvents::InstanceExternalAuditEventDestinationsResolver
+
+ field :ai_messages, ::Types::Ai::CachedMessageType.connection_type,
+ resolver: ::Resolvers::Ai::MessagesResolver,
+ alpha: { milestone: '16.1' },
+ description: 'Find AI messages.'
end
def vulnerability(id:)
diff --git a/ee/app/graphql/resolvers/ai/messages_resolver.rb b/ee/app/graphql/resolvers/ai/messages_resolver.rb
new file mode 100644
index 0000000000000000000000000000000000000000..15b704f31f723f5ca28929c895124fe00dca9ff6
--- /dev/null
+++ b/ee/app/graphql/resolvers/ai/messages_resolver.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Resolvers
+ module Ai
+ class MessagesResolver < BaseResolver
+ type Types::Ai::CachedMessageType, null: false
+
+ argument :request_ids, [GraphQL::Types::ID],
+ required: false,
+ description: 'Array of request IDs to fetch.'
+
+ argument :roles, [Types::Ai::CachedMessageRoleEnum],
+ required: false,
+ description: 'Array of request IDs to fetch.'
+
+ def resolve(**args)
+ return [] unless current_user && Feature.enabled?(:ai_redis_cache, current_user)
+
+ ::Gitlab::Llm::Cache.new(current_user).find_all(args)
+ end
+ end
+ end
+end
diff --git a/ee/app/graphql/types/ai/cached_message_role_enum.rb b/ee/app/graphql/types/ai/cached_message_role_enum.rb
new file mode 100644
index 0000000000000000000000000000000000000000..4d3c264e25bf65e3d2eb4ea1f6c3bb581ff0dcad
--- /dev/null
+++ b/ee/app/graphql/types/ai/cached_message_role_enum.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Types
+ module Ai
+ class CachedMessageRoleEnum < BaseEnum
+ graphql_name 'AiCachedMessageRole'
+ description 'Roles to filter in chat message.'
+
+ value 'USER', 'Filter only user messages.', value: 'user'
+ value 'ASSISTANT', 'Filter only AI responses.', value: 'assistant'
+ end
+ end
+end
diff --git a/ee/app/graphql/types/ai/cached_message_type.rb b/ee/app/graphql/types/ai/cached_message_type.rb
new file mode 100644
index 0000000000000000000000000000000000000000..ac5f2e75c086d7dfa0ea9365d8f71f66884431ca
--- /dev/null
+++ b/ee/app/graphql/types/ai/cached_message_type.rb
@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Types
+ module Ai
+ # rubocop: disable Graphql/AuthorizeTypes
+ class CachedMessageType < Types::BaseObject
+ graphql_name 'AiCachedMessageType'
+
+ field :id,
+ GraphQL::Types::ID,
+ description: 'UUID of the message.'
+
+ field :request_id,
+ GraphQL::Types::ID,
+ description: 'UUID of the original request message.'
+
+ field :content,
+ GraphQL::Types::String,
+ null: true,
+ description: 'Content of the message. Can be null for user requests or failed responses.'
+
+ field :role,
+ Types::Ai::CachedMessageRoleEnum,
+ null: false,
+ description: 'Message role.'
+
+ field :timestamp,
+ Types::TimeType,
+ null: false,
+ description: 'Message timestamp.'
+
+ field :errors,
+ [GraphQL::Types::String],
+ null: false,
+ description: 'Errors that occurred while asynchronously fetching an AI(assistant) response.'
+ end
+ # rubocop: enable Graphql/AuthorizeTypes
+ end
+end
diff --git a/ee/app/services/llm/base_service.rb b/ee/app/services/llm/base_service.rb
index c714b419591785135a642ce9ac366da663024d82..88e7631d7932358ea42b4c888743134219883f27 100644
--- a/ee/app/services/llm/base_service.rb
+++ b/ee/app/services/llm/base_service.rb
@@ -49,7 +49,7 @@ def perform_async(user, resource, action_name, options)
::Llm::CompletionWorker.perform_async(user.id, resource.id, resource.class.name, action_name, options)
- payload = { request_id: request_id }
+ payload = { request_id: request_id, role: 'user' }
::Gitlab::Llm::Cache.new(user).add(payload)
success(payload)
diff --git a/ee/lib/gitlab/llm/cache.rb b/ee/lib/gitlab/llm/cache.rb
index 253ae345cae5f7c9e41b7025896d5d2d4f91e49d..30774c1b17c0b9b21180482c6f0abbb4e1d941dd 100644
--- a/ee/lib/gitlab/llm/cache.rb
+++ b/ee/lib/gitlab/llm/cache.rb
@@ -21,22 +21,22 @@ def add(payload)
return unless Feature.enabled?(:ai_redis_cache, user)
data = {
+ id: SecureRandom.uuid,
request_id: payload[:request_id],
- timestamp: Time.now.to_i
+ timestamp: Time.current.to_s,
+ role: payload[:role] || 'user'
}
- data[:response_body] = payload[:response_body][0, MAX_TEXT_LIMIT] if payload[:response_body]
+ data[:content] = payload[:content][0, MAX_TEXT_LIMIT] if payload[:content]
data[:error] = payload[:errors].join(". ") if payload[:errors]
cache_data(data)
end
- def get(request_id)
- all.find { |data| data['request_id'] == request_id && data['response_body'].present? }
- end
-
- def all
+ def find_all(filters = {})
with_redis do |redis|
- redis.xrange(key).map { |_id, data| data }
+ redis.xrange(key).filter_map do |_id, data|
+ CachedMessage.new(data) if matches_filters?(data, filters)
+ end
end
end
@@ -58,6 +58,13 @@ def key
def with_redis(&block)
Gitlab::Redis::Chat.with(&block) # rubocop: disable CodeReuse/ActiveRecord
end
+
+ def matches_filters?(data, filters)
+ return false if filters[:roles] && filters[:roles].exclude?(data['role'])
+ return false if filters[:request_ids] && filters[:request_ids].exclude?(data['request_id'])
+
+ data
+ end
end
end
end
diff --git a/ee/lib/gitlab/llm/cached_message.rb b/ee/lib/gitlab/llm/cached_message.rb
new file mode 100644
index 0000000000000000000000000000000000000000..8b2f0194143636247db07648bb02737c03c6cc3d
--- /dev/null
+++ b/ee/lib/gitlab/llm/cached_message.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Gitlab
+ module Llm
+ class CachedMessage
+ attr_reader :id, :request_id, :content, :role, :timestamp, :error
+
+ def initialize(data)
+ @id = data['id']
+ @request_id = data['request_id']
+ @content = data['content']
+ @role = data['role']
+ @error = data['error']
+ @timestamp = Time.zone.parse(data['timestamp'])
+ end
+
+ def to_global_id
+ ::Gitlab::GlobalId.build(self)
+ end
+
+ def errors
+ Array.wrap(error)
+ end
+ end
+ end
+end
diff --git a/ee/lib/gitlab/llm/graphql_subscription_response_service.rb b/ee/lib/gitlab/llm/graphql_subscription_response_service.rb
index 1c0d23a25d21f04623c78239df43640f3ca30c4a..c412b5064069db56a139c6ee7c607d137d0dcf38 100644
--- a/ee/lib/gitlab/llm/graphql_subscription_response_service.rb
+++ b/ee/lib/gitlab/llm/graphql_subscription_response_service.rb
@@ -20,7 +20,8 @@ def execute
model_name: resource.class.name,
# todo: do we need to sanitize/refine this response in any ways?
response_body: generate_response_body(response_modifier.response_body),
- errors: response_modifier.errors
+ errors: response_modifier.errors,
+ role: 'assistant'
}
logger.debug(
@@ -28,7 +29,9 @@ def execute
data: data
)
- Gitlab::Llm::Cache.new(user).add(data.slice(:request_id, :response_body, :errors))
+ Gitlab::Llm::Cache.new(user).add(
+ data.slice(:request_id, :errors, :role).merge(content: data[:response_body])
+ )
GraphqlTriggers.ai_completion_response(user.to_global_id, resource.to_global_id, data)
end
diff --git a/ee/spec/graphql/types/query_type_spec.rb b/ee/spec/graphql/types/query_type_spec.rb
index 5e0d128490cbf63949452f7590a154c308396f65..bc4aaa204e7c1fbed1d930c8b5a937f8e9cd5aad 100644
--- a/ee/spec/graphql/types/query_type_spec.rb
+++ b/ee/spec/graphql/types/query_type_spec.rb
@@ -7,6 +7,7 @@
specify do
expected_ee_fields = [
+ :ai_messages,
:ci_catalog_resources,
:ci_minutes_usage,
:current_license,
diff --git a/ee/spec/lib/gitlab/llm/cache_spec.rb b/ee/spec/lib/gitlab/llm/cache_spec.rb
index 133652a7e32a4d832d1a476cb45df95af495e72c..9c1469312adb968adc659c9d45f332ca97a5eb30 100644
--- a/ee/spec/lib/gitlab/llm/cache_spec.rb
+++ b/ee/spec/lib/gitlab/llm/cache_spec.rb
@@ -2,16 +2,17 @@
require 'spec_helper'
-RSpec.describe Gitlab::Llm::Cache, :clean_gitlab_redis_cache, feature_category: :no_category do # rubocop: disable RSpec/InvalidFeatureCategory
+RSpec.describe Gitlab::Llm::Cache, :clean_gitlab_redis_cache, feature_category: :shared do
let_it_be(:user) { create(:user) }
- let(:uuid) { 'uuid' }
- let(:timestamp) { Time.now.to_i.to_s }
+ let(:request_id) { 'uuid' }
+ let(:timestamp) { Time.current.to_s }
let(:payload) do
{
timestamp: timestamp,
- request_id: uuid,
+ request_id: request_id,
errors: ['some error1', 'another error'],
- response_body: 'response'
+ role: 'user',
+ content: 'response'
}
end
@@ -21,20 +22,25 @@
other_user = create(:user)
other_cache = described_class.new(other_user)
- other_cache.add(payload.merge(response_body: 'other user unrelated cache'))
+ other_cache.add(payload.merge(content: 'other user unrelated cache'))
end
describe '#add' do
- it 'adds new message' do
- expect(subject.all).to be_empty
+ it 'adds new message', :aggregate_failures do
+ uuid = 'unique_id'
+
+ expect(SecureRandom).to receive(:uuid).once.and_return(uuid)
+ expect(subject.find_all).to be_empty
subject.add(payload)
- expected_data = payload
- .except(:errors)
- .merge(error: 'some error1. another error')
- .with_indifferent_access
- expect(subject.all).to eq([expected_data.with_indifferent_access])
+ last = subject.find_all.last
+ expect(last.id).to eq(uuid)
+ expect(last.request_id).to eq(request_id)
+ expect(last.errors).to eq(['some error1. another error'])
+ expect(last.content).to eq('response')
+ expect(last.role).to eq('user')
+ expect(last.timestamp).not_to be_nil
end
context 'when ai_redis_cache is disabled' do
@@ -43,11 +49,11 @@
end
it 'does not add new message' do
- expect(subject.all).to be_empty
+ expect(subject.find_all).to be_empty
subject.add(payload)
- expect(subject.all).to be_empty
+ expect(subject.find_all).to be_empty
end
end
@@ -57,69 +63,45 @@
end
it 'removes oldes messages if we reach maximum message limit' do
- subject.add(payload.merge(response_body: 'msg1'))
- subject.add(payload.merge(response_body: 'msg2'))
+ subject.add(payload.merge(content: 'msg1'))
+ subject.add(payload.merge(content: 'msg2'))
- expect(subject.all).to match([
- a_hash_including('response_body' => 'msg1'),
- a_hash_including('response_body' => 'msg2')
- ])
+ expect(subject.find_all.map(&:content)).to eq(%w[msg1 msg2])
- subject.add(payload.merge(response_body: 'msg3'))
+ subject.add(payload.merge(content: 'msg3'))
- expect(subject.all).to match([
- a_hash_including('response_body' => 'msg2'),
- a_hash_including('response_body' => 'msg3')
- ])
+ expect(subject.find_all.map(&:content)).to eq(%w[msg2 msg3])
end
end
end
- describe '#get' do
- context 'when there is both request and response' do
- before do
- subject.add(payload.merge(response_body: nil))
- subject.add(payload.merge(response_body: 'msg'))
- end
-
- it 'gets response by request id' do
- data = subject.get(uuid)
+ describe '#find_all' do
+ let(:filters) { {} }
- expect(data).not_to be_nil
- expect(data['response_body']).to eq('msg')
- end
+ before do
+ subject.add(payload.merge(content: 'msg1', role: 'user', request_id: '1'))
+ subject.add(payload.merge(content: 'msg2', role: 'assistant', request_id: '2'))
+ subject.add(payload.merge(content: 'msg3', role: 'assistant', request_id: '3'))
end
- context 'when there is only request' do
- before do
- subject.add(payload.merge(response_body: nil))
- end
-
- it 'returns nil' do
- data = subject.get(uuid)
-
- expect(data).to be_nil
- end
+ it 'returns all records for this user' do
+ expect(subject.find_all(filters).map(&:content)).to eq(%w[msg1 msg2 msg3])
end
- context 'when there is no record with this request id' do
- it 'returns nil' do
- data = subject.get(uuid)
+ context 'when filtering by role' do
+ let(:filters) { { roles: ['user'] } }
- expect(data).to be_nil
+ it 'returns only records for this role' do
+ expect(subject.find_all(filters).map(&:content)).to eq(%w[msg1])
end
end
- end
- describe '#all' do
- it 'returns all records for this user' do
- subject.add(payload.merge(response_body: 'msg1'))
- subject.add(payload.merge(response_body: 'msg2'))
+ context 'when filtering by request_ids' do
+ let(:filters) { { request_ids: %w[2 3] } }
- expect(subject.all).to match([
- a_hash_including('response_body' => 'msg1'),
- a_hash_including('response_body' => 'msg2')
- ])
+ it 'returns only records with the same request_id' do
+ expect(subject.find_all(filters).map(&:content)).to eq(%w[msg2 msg3])
+ end
end
end
end
diff --git a/ee/spec/lib/gitlab/llm/cached_message_spec.rb b/ee/spec/lib/gitlab/llm/cached_message_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..3348f53d1a3fe08b80d036a53c6ffe539082010a
--- /dev/null
+++ b/ee/spec/lib/gitlab/llm/cached_message_spec.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::CachedMessage, feature_category: :shared do
+ let(:timestamp) { Time.current }
+ let(:data) do
+ {
+ 'timestamp' => timestamp.to_s,
+ 'id' => 'uuid',
+ 'request_id' => 'original_request_id',
+ 'error' => 'some error1. another error',
+ 'role' => 'user',
+ 'content' => 'response'
+ }
+ end
+
+ subject { described_class.new(data) }
+
+ 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')
+ end
+ end
+
+ describe '#errors' do
+ it 'returns message error wrapped as an array' do
+ expect(subject.errors).to eq([data['error']])
+ end
+ end
+end
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 88ae7f21d5a16b7d409ac19d3ecda75a878778bc..3130c1b283e2ef03eb2404848ba5b4b8d4f7523c 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
@@ -41,6 +41,7 @@
model_name: resource.class.name,
response_body: response_body,
request_id: 'uuid',
+ role: 'assistant',
errors: []
}
end
@@ -59,7 +60,8 @@
it 'caches response' do
expect_next_instance_of(::Gitlab::Llm::Cache) do |cache|
- expect(cache).to receive(:add).with(payload.slice(:request_id, :response_body, :errors))
+ expect(cache).to receive(:add)
+ .with(payload.slice(:request_id, :errors, :role).merge(content: payload[:response_body]))
end
subject
diff --git a/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb b/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb
index f5631466783e95bc72261c4669e1a56a74dcb3d3..b95e33134c87f51ea510326d1894c6f3cf2cf3aa 100644
--- a/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb
+++ b/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb
@@ -62,13 +62,14 @@
uuid = 'uuid'
- expect(SecureRandom).to receive(:uuid).and_return(uuid)
+ expect(SecureRandom).to receive(:uuid).twice.and_return(uuid)
data = {
id: uuid,
model_name: 'MergeRequest',
response_body: content,
- request_id: 'uuid',
+ request_id: uuid,
+ role: 'assistant',
errors: []
}
diff --git a/ee/spec/requests/api/graphql/ai_messages_spec.rb b/ee/spec/requests/api/graphql/ai_messages_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..a7a4210cb49a5be491fff126f5b91f8d6823fb38
--- /dev/null
+++ b/ee/spec/requests/api/graphql/ai_messages_spec.rb
@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Querying user AI messages', :clean_gitlab_redis_cache, feature_category: :shared do
+ include GraphqlHelpers
+
+ let_it_be(:user) { create(:user) }
+ let_it_be(:other_user) { create(:user) }
+
+ let(:fields) do
+ <<~GRAPHQL
+ nodes {
+ requestId
+ content
+ role
+ timestamp
+ errors
+ }
+ GRAPHQL
+ end
+
+ let(:arguments) { { requestIds: 'uuid1' } }
+ let(:query) { graphql_query_for('aiMessages', arguments, fields) }
+
+ subject { graphql_data.dig('aiMessages', 'nodes') }
+
+ before do
+ ::Gitlab::Llm::Cache.new(user).add(request_id: 'uuid1', role: 'user')
+ ::Gitlab::Llm::Cache.new(user).add(request_id: 'uuid1', role: 'assistant', content: 'response')
+ # should not be included in response because it's for other user
+ ::Gitlab::Llm::Cache.new(other_user).add(request_id: 'uuid1', role: 'user')
+ end
+
+ context 'when user is not logged in' do
+ let(:current_user) { nil }
+
+ it 'returns an empty array' do
+ post_graphql(query, current_user: current_user)
+
+ expect(subject).to be_empty
+ end
+ end
+
+ context 'when user is logged in' do
+ let(:current_user) { user }
+
+ it 'returns user messages', :freeze_time do
+ post_graphql(query, current_user: current_user)
+
+ expect(subject).to eq([
+ { 'requestId' => 'uuid1', 'content' => nil, 'role' => 'USER', 'errors' => [],
+ 'timestamp' => Time.current.iso8601 },
+ { 'requestId' => 'uuid1', 'content' => 'response', 'role' => 'ASSISTANT', 'errors' => [],
+ 'timestamp' => Time.current.iso8601 }
+ ])
+ end
+
+ context 'when ai_redis_cache is disabled' do
+ before do
+ stub_feature_flags(ai_redis_cache: false)
+ end
+
+ it 'returns an empty array' do
+ post_graphql(query, current_user: current_user)
+
+ expect(subject).to be_empty
+ end
+ end
+ end
+end
diff --git a/ee/spec/support/shared_examples/services/llm/async_service_shared_examples.rb b/ee/spec/support/shared_examples/services/llm/async_service_shared_examples.rb
index 813b3c6b872a1e4d6491dbb007db7c42bc170b59..f039261ed16f46f2120f4a5ff14014e250a33f29 100644
--- a/ee/spec/support/shared_examples/services/llm/async_service_shared_examples.rb
+++ b/ee/spec/support/shared_examples/services/llm/async_service_shared_examples.rb
@@ -4,7 +4,7 @@
it 'executes the service asynchronously' do
expected_options = options.merge(request_id: 'uuid')
- expect(SecureRandom).to receive(:uuid).once.and_return('uuid')
+ expect(SecureRandom).to receive(:uuid).twice.and_return('uuid')
expect(::Llm::CompletionWorker)
.to receive(:perform_async)
.with(user.id, resource.id, resource.class.name, action_name, expected_options)
@@ -15,7 +15,7 @@
it 'caches request' do
expect(SecureRandom).to receive(:uuid).once.and_return('uuid')
expect_next_instance_of(::Gitlab::Llm::Cache) do |cache|
- expect(cache).to receive(:add).with({ request_id: 'uuid' })
+ expect(cache).to receive(:add).with({ request_id: 'uuid', role: 'user' })
end
subject.execute