diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts index 59b627929d163..5024a8b619c24 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts @@ -30,9 +30,6 @@ export const getConversationResponseMock = (): ConversationResponse => ({ model: 'test', provider: 'Azure OpenAI', }, - summary: { - content: 'test', - }, category: 'assistant', users: [ { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts index cc25af9f96782..dd97490e6366b 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts @@ -7,82 +7,10 @@ import { UpdateConversationSchema } from './update_conversation'; -export const getUpdateScript = ({ - conversation, - isPatch, -}: { - conversation: UpdateConversationSchema; - isPatch?: boolean; -}) => { +export const getUpdateScript = ({ conversation }: { conversation: UpdateConversationSchema }) => { + // https://www.elastic.co/docs/reference/elasticsearch/mapping-reference/semantic-text#update-script + // Cannot use script for bulk update of the documents with semantic_text fields return { - script: { - source: ` - if (params.assignEmpty == true || params.containsKey('api_config')) { - if (ctx._source.api_config != null) { - if (params.assignEmpty == true || params.api_config.containsKey('connector_id')) { - ctx._source.api_config.connector_id = params.api_config.connector_id; - ctx._source.api_config.remove('model'); - ctx._source.api_config.remove('provider'); - } - // an update to apiConfig that does not contain defaultSystemPromptId should remove it - if (params.assignEmpty == true || (params.containsKey('api_config') && !params.api_config.containsKey('default_system_prompt_id'))) { - ctx._source.api_config.remove('default_system_prompt_id'); - } - if (params.assignEmpty == true || params.api_config.containsKey('action_type_id')) { - ctx._source.api_config.action_type_id = params.api_config.action_type_id; - } - if (params.assignEmpty == true || params.api_config.containsKey('default_system_prompt_id')) { - ctx._source.api_config.default_system_prompt_id = params.api_config.default_system_prompt_id; - } - if (params.assignEmpty == true || params.api_config.containsKey('model')) { - ctx._source.api_config.model = params.api_config.model; - } - if (params.assignEmpty == true || params.api_config.containsKey('provider')) { - ctx._source.api_config.provider = params.api_config.provider; - } - } else { - ctx._source.api_config = params.api_config; - } - } - if (params.assignEmpty == true || params.containsKey('exclude_from_last_conversation_storage')) { - ctx._source.exclude_from_last_conversation_storage = params.exclude_from_last_conversation_storage; - } - if (params.assignEmpty == true || params.containsKey('replacements')) { - ctx._source.replacements = params.replacements; - } - if (params.assignEmpty == true || params.containsKey('title')) { - ctx._source.title = params.title; - } - if (params.assignEmpty == true || params.containsKey('messages')) { - def messages = []; - for (message in params.messages) { - def newMessage = [:]; - newMessage['@timestamp'] = message['@timestamp']; - newMessage.content = message.content; - newMessage.is_error = message.is_error; - newMessage.reader = message.reader; - newMessage.role = message.role; - if (message.trace_data != null) { - newMessage.trace_data = message.trace_data; - } - if (message.metadata != null) { - newMessage.metadata = [:]; - if (message.metadata.content_references != null) { - newMessage.metadata.content_references = message.metadata.content_references; - } - } - messages.add(newMessage); - } - ctx._source.messages = messages; - } - ctx._source.updated_at = params.updated_at; - `, - lang: 'painless', - params: { - ...conversation, // when assigning undefined in painless, it will remove property and wil set it to null - // for patch we don't want to remove unspecified value in payload - assignEmpty: !(isPatch ?? true), - }, - }, + doc: conversation, }; }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts index f58a06d812f75..3572daec276bb 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts @@ -5,7 +5,7 @@ * 2.0. */ import { elasticsearchServiceMock, loggingSystemMock } from '@kbn/core/server/mocks'; -import type { UpdateByQueryRequest } from '@elastic/elasticsearch/lib/api/types'; +import type { BulkRequest } from '@elastic/elasticsearch/lib/api/types'; import { AIAssistantConversationsDataClient } from '.'; import { getUpdateConversationSchemaMock } from '../../__mocks__/conversations_schema.mock'; import { authenticatedUser } from '../../__mocks__/user'; @@ -146,6 +146,20 @@ describe('AIAssistantConversationsDataClient', () => { }); test('should update conversation with new messages', async () => { + clusterClient.search.mockReturnValue({ + // @ts-ignore + hits: { + total: { value: 1 }, + hits: [ + { + _id: '04128c15-0d1b-4716-a4c5-46997ac7f3bd', + _index: 'test-index', + _source: {}, + }, + ], + }, + }); + const assistantConversationsDataClient = new AIAssistantConversationsDataClient( assistantConversationsDataClientParams ); @@ -156,45 +170,43 @@ describe('AIAssistantConversationsDataClient', () => { ), }); - const params = clusterClient.updateByQuery.mock.calls[0][0] as UpdateByQueryRequest; + const params = clusterClient.bulk.mock.calls[0][0] as BulkRequest; - expect(params.query).toEqual({ - ids: { - values: ['04128c15-0d1b-4716-a4c5-46997ac7f3bd'], - }, - }); - - expect(params.script).toEqual({ - source: expect.anything(), - lang: 'painless', - params: { - api_config: { - action_type_id: '.gen-ai', - connector_id: '2', - default_system_prompt_id: 'Default', - model: 'model', - provider: undefined, + expect(params.refresh).toEqual('wait_for'); + expect(params.body).toEqual([ + { + update: { + _id: '04128c15-0d1b-4716-a4c5-46997ac7f3bd', + _index: 'test-index', + _source: true, + retry_on_conflict: 3, }, - assignEmpty: false, - exclude_from_last_conversation_storage: false, - id: '04128c15-0d1b-4716-a4c5-46997ac7f3bd', - messages: [ - { - '@timestamp': '2019-12-13T16:40:33.400Z', - content: 'test content', - is_error: undefined, - reader: undefined, - role: 'user', - trace_data: { - trace_id: '1', - transaction_id: '2', - }, + }, + { + doc: { + id: '04128c15-0d1b-4716-a4c5-46997ac7f3bd', + updated_at: '2023-03-28T22:27:28.159Z', + title: 'Welcome 2', + api_config: { + action_type_id: '.gen-ai', + connector_id: '2', + default_system_prompt_id: 'Default', + model: 'model', }, - ], - replacements: undefined, - title: 'Welcome 2', - updated_at: '2023-03-28T22:27:28.159Z', + exclude_from_last_conversation_storage: false, + messages: [ + { + '@timestamp': '2019-12-13T16:40:33.400Z', + content: 'test content', + role: 'user', + trace_data: { + trace_id: '1', + transaction_id: '2', + }, + }, + ], + }, }, - }); + ]); }); }); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts index ee7508c1399b2..ce50a2a87bd21 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts @@ -127,19 +127,15 @@ export class AIAssistantConversationsDataClient extends AIAssistantDataClient { public updateConversation = async ({ conversationUpdateProps, authenticatedUser, - isPatch, }: { conversationUpdateProps: ConversationUpdateProps; authenticatedUser?: AuthenticatedUser; - isPatch?: boolean; }): Promise => { - const esClient = await this.options.elasticsearchClientPromise; + const dataWriter = await this.getWriter(); return updateConversation({ - esClient, - logger: this.options.logger, - conversationIndex: this.indexTemplateAndPattern.alias, conversationUpdateProps, - isPatch, + dataWriter, + logger: this.options.logger, user: authenticatedUser ?? this.options.currentUser ?? undefined, }); }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transform.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transform.test.ts new file mode 100644 index 0000000000000..2275ff6fdd67d --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transform.test.ts @@ -0,0 +1,251 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { estypes } from '@elastic/elasticsearch'; +import { + transformESToConversation, + transformESSearchToConversations, + transformESToConversations, + transformFieldNamesToSourceScheme, +} from './transforms'; +import type { EsConversationSchema } from './types'; + +const userAsUser = { + id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', + name: 'elastic', +}; + +const getEsConversationMock = (): EsConversationSchema => { + return { + '@timestamp': '2025-08-19T10:49:52.884Z', + updated_at: '2025-08-19T13:26:01.746Z', + api_config: { + action_type_id: '.gen-ai', + connector_id: 'gpt-4-1', + }, + namespace: 'default', + created_at: '2025-08-19T10:49:52.884Z', + messages: [ + { + '@timestamp': '2025-08-19T10:49:53.799Z', + role: 'user', + content: 'Hello there, how many opened alerts do I have?', + }, + { + metadata: { + content_references: { + oQ5xL: { + id: 'oQ5xL', + type: 'SecurityAlertsPage', + }, + }, + }, + '@timestamp': '2025-08-19T10:49:57.398Z', + role: 'assistant', + is_error: false, + trace_data: { + transaction_id: 'ee432e8be6ad3f9c', + trace_id: 'f44d01b6095d35dce15aa8137df76e29', + }, + content: 'You currently have 61 open alerts in your environment. {reference(oQ5xL)}', + }, + ], + replacements: [], + title: 'Viewing the Number of Open Alerts in Elastic Security', + category: 'assistant', + users: [userAsUser], + id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + }; +}; + +const getEsSearchConversationsMock = (): estypes.SearchResponse => { + return { + took: 0, + timed_out: false, + _shards: { + total: 1, + successful: 1, + skipped: 0, + failed: 0, + }, + hits: { + total: { + value: 1, + relation: 'eq', + }, + max_score: null, + hits: [ + { + _index: '.ds-.kibana-elastic-ai-assistant-conversations-default-2025.08.19-000001', + _id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + _seq_no: 8, + _primary_term: 1, + _score: null, + _source: { + ...getEsConversationMock(), + }, + sort: [1755607491083], + }, + ], + }, + }; +}; + +describe('transforms', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + describe('transformESToConversation', () => { + it('should correctly transform ES conversation', () => { + const esConversation = getEsConversationMock(); + const conversation = transformESToConversation(esConversation); + expect(conversation).toEqual({ + timestamp: '2025-08-19T10:49:52.884Z', + createdAt: '2025-08-19T10:49:52.884Z', + users: [userAsUser], + title: 'Viewing the Number of Open Alerts in Elastic Security', + category: 'assistant', + apiConfig: { actionTypeId: '.gen-ai', connectorId: 'gpt-4-1' }, + messages: [ + { + timestamp: '2025-08-19T10:49:53.799Z', + content: 'Hello there, how many opened alerts do I have?', + role: 'user', + }, + { + timestamp: '2025-08-19T10:49:57.398Z', + content: 'You currently have 61 open alerts in your environment. {reference(oQ5xL)}', + role: 'assistant', + metadata: { contentReferences: { oQ5xL: { id: 'oQ5xL', type: 'SecurityAlertsPage' } } }, + traceData: { + traceId: 'f44d01b6095d35dce15aa8137df76e29', + transactionId: 'ee432e8be6ad3f9c', + }, + }, + ], + updatedAt: '2025-08-19T13:26:01.746Z', + replacements: {}, + namespace: 'default', + id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + }); + }); + }); + + describe('transformESSearchToConversations', () => { + it('should correctly transform conversation hits', () => { + const conversationHits = getEsSearchConversationsMock(); + const conversations = transformESSearchToConversations(conversationHits); + expect(conversations).toEqual([ + { + timestamp: '2025-08-19T10:49:52.884Z', + createdAt: '2025-08-19T10:49:52.884Z', + users: [userAsUser], + title: 'Viewing the Number of Open Alerts in Elastic Security', + category: 'assistant', + apiConfig: { actionTypeId: '.gen-ai', connectorId: 'gpt-4-1' }, + messages: [ + { + timestamp: '2025-08-19T10:49:53.799Z', + content: 'Hello there, how many opened alerts do I have?', + role: 'user', + }, + { + timestamp: '2025-08-19T10:49:57.398Z', + content: 'You currently have 61 open alerts in your environment. {reference(oQ5xL)}', + role: 'assistant', + metadata: { + contentReferences: { oQ5xL: { id: 'oQ5xL', type: 'SecurityAlertsPage' } }, + }, + traceData: { + traceId: 'f44d01b6095d35dce15aa8137df76e29', + transactionId: 'ee432e8be6ad3f9c', + }, + }, + ], + updatedAt: '2025-08-19T13:26:01.746Z', + replacements: {}, + namespace: 'default', + id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + }, + ]); + }); + }); + + describe('transformESToConversations', () => { + it('should correctly transform ES conversations', () => { + const esConversations = [getEsConversationMock()]; + const conversations = transformESToConversations(esConversations); + expect(conversations).toEqual([ + { + timestamp: '2025-08-19T10:49:52.884Z', + createdAt: '2025-08-19T10:49:52.884Z', + users: [userAsUser], + title: 'Viewing the Number of Open Alerts in Elastic Security', + category: 'assistant', + apiConfig: { actionTypeId: '.gen-ai', connectorId: 'gpt-4-1' }, + messages: [ + { + timestamp: '2025-08-19T10:49:53.799Z', + content: 'Hello there, how many opened alerts do I have?', + role: 'user', + }, + { + timestamp: '2025-08-19T10:49:57.398Z', + content: 'You currently have 61 open alerts in your environment. {reference(oQ5xL)}', + role: 'assistant', + metadata: { + contentReferences: { oQ5xL: { id: 'oQ5xL', type: 'SecurityAlertsPage' } }, + }, + traceData: { + traceId: 'f44d01b6095d35dce15aa8137df76e29', + transactionId: 'ee432e8be6ad3f9c', + }, + }, + ], + updatedAt: '2025-08-19T13:26:01.746Z', + replacements: {}, + namespace: 'default', + id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + }, + ]); + }); + }); + + describe('transformFieldNamesToSourceScheme', () => { + it('should correctly transform empty array', () => { + const sourceNames = transformFieldNamesToSourceScheme([]); + expect(sourceNames).toEqual([]); + }); + + it('should correctly transform field names', () => { + const fields = [ + 'timestamp', + 'apiConfig', + 'apiConfig.actionTypeId', + 'apiConfig.connectorId', + 'apiConfig.defaultSystemPromptId', + 'apiConfig.model', + 'apiConfig.provider', + ]; + const sourceNames = transformFieldNamesToSourceScheme(fields); + expect(sourceNames).toEqual([ + '@timestamp', + 'api_config', + 'api_config.action_type_id', + 'api_config.connector_id', + 'api_config.default_system_prompt_id', + 'api_config.model', + 'api_config.provider', + ]); + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts index d2b47e25476de..05eaf2da96388 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/transforms.ts @@ -6,13 +6,79 @@ */ import type { estypes } from '@elastic/elasticsearch'; -import { - ConversationResponse, - Replacements, - replaceOriginalValuesWithUuidValues, -} from '@kbn/elastic-assistant-common'; +import type { ConversationResponse, Replacements } from '@kbn/elastic-assistant-common'; +import { replaceOriginalValuesWithUuidValues } from '@kbn/elastic-assistant-common'; import _ from 'lodash'; -import { EsConversationSchema } from './types'; +import type { EsConversationSchema } from './types'; + +export const transformESToConversation = ( + conversationSchema: EsConversationSchema +): ConversationResponse => { + const replacements = conversationSchema.replacements?.reduce((acc: Record, r) => { + acc[r.uuid] = r.value; + return acc; + }, {}) as Replacements; + const conversation: ConversationResponse = { + timestamp: conversationSchema['@timestamp'], + createdAt: conversationSchema.created_at, + users: + conversationSchema.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + title: conversationSchema.title, + category: conversationSchema.category, + ...(conversationSchema.api_config + ? { + apiConfig: { + actionTypeId: conversationSchema.api_config.action_type_id, + connectorId: conversationSchema.api_config.connector_id, + defaultSystemPromptId: conversationSchema.api_config.default_system_prompt_id, + model: conversationSchema.api_config.model, + provider: conversationSchema.api_config.provider, + }, + } + : {}), + excludeFromLastConversationStorage: conversationSchema.exclude_from_last_conversation_storage, + messages: + // eslint-disable-next-line @typescript-eslint/no-explicit-any + conversationSchema.messages?.map((message: Record) => ({ + timestamp: message['@timestamp'], + // always return anonymized data from the client + content: replaceOriginalValuesWithUuidValues({ + messageContent: message.content, + replacements, + }), + ...(message.is_error ? { isError: message.is_error } : {}), + ...(message.reader ? { reader: message.reader } : {}), + ...(message.user ? { user: message.user } : {}), + role: message.role, + ...(message.metadata + ? { + metadata: { + ...(message.metadata.content_references + ? { contentReferences: message.metadata.content_references } + : {}), + }, + } + : {}), + ...(message.trace_data + ? { + traceData: { + traceId: message.trace_data?.trace_id, + transactionId: message.trace_data?.transaction_id, + }, + } + : {}), + })), + updatedAt: conversationSchema.updated_at, + replacements, + namespace: conversationSchema.namespace, + id: conversationSchema.id, + }; + + return conversation; +}; export const transformESSearchToConversations = ( response: estypes.SearchResponse @@ -22,145 +88,18 @@ export const transformESSearchToConversations = ( .map((hit) => { // eslint-disable-next-line @typescript-eslint/no-non-null-assertion const conversationSchema = hit._source!; - const conversation: ConversationResponse = { - timestamp: conversationSchema['@timestamp'], - createdAt: conversationSchema.created_at, - users: - conversationSchema.users?.map((user) => ({ - id: user.id, - name: user.name, - })) ?? [], - title: conversationSchema.title, - category: conversationSchema.category, - summary: conversationSchema.summary, - ...(conversationSchema.api_config - ? { - apiConfig: { - connectorId: conversationSchema.api_config.connector_id, - actionTypeId: conversationSchema.api_config.action_type_id, - defaultSystemPromptId: conversationSchema.api_config.default_system_prompt_id, - model: conversationSchema.api_config.model, - provider: conversationSchema.api_config.provider, - }, - } - : {}), - excludeFromLastConversationStorage: - conversationSchema.exclude_from_last_conversation_storage, - messages: - // eslint-disable-next-line @typescript-eslint/no-explicit-any - conversationSchema.messages?.map((message: Record) => ({ - timestamp: message['@timestamp'], - // always return anonymized data from the client - content: conversationSchema.replacements - ? replaceOriginalValuesWithUuidValues({ - messageContent: message.content, - replacements: conversationSchema.replacements?.reduce( - (acc: Record, r) => { - acc[r.uuid] = r.value; - return acc; - }, - {} - ), - }) - : message.content, - ...(message.is_error ? { isError: message.is_error } : {}), - ...(message.reader ? { reader: message.reader } : {}), - role: message.role, - ...(message.metadata - ? { - metadata: { - ...(message.metadata.content_references - ? { contentReferences: message.metadata.content_references } - : {}), - }, - } - : {}), - ...(message.trace_data - ? { - traceData: { - traceId: message.trace_data?.trace_id, - transactionId: message.trace_data?.transaction_id, - }, - } - : {}), - })) ?? [], - updatedAt: conversationSchema.updated_at, - replacements: conversationSchema.replacements?.reduce((acc: Record, r) => { - acc[r.uuid] = r.value; - return acc; - }, {}), - namespace: conversationSchema.namespace, + return transformESToConversation({ + ...conversationSchema, // eslint-disable-next-line @typescript-eslint/no-non-null-assertion id: hit._id!, - }; - - return conversation; + }); }); }; export const transformESToConversations = ( response: EsConversationSchema[] ): ConversationResponse[] => { - return response.map((conversationSchema) => { - const replacements = conversationSchema.replacements?.reduce( - (acc: Record, r) => { - acc[r.uuid] = r.value; - return acc; - }, - {} - ) as Replacements; - const conversation: ConversationResponse = { - timestamp: conversationSchema['@timestamp'], - createdAt: conversationSchema.created_at, - users: - conversationSchema.users?.map((user) => ({ - id: user.id, - name: user.name, - })) ?? [], - title: conversationSchema.title, - category: conversationSchema.category, - summary: conversationSchema.summary, - ...(conversationSchema.api_config - ? { - apiConfig: { - actionTypeId: conversationSchema.api_config.action_type_id, - connectorId: conversationSchema.api_config.connector_id, - defaultSystemPromptId: conversationSchema.api_config.default_system_prompt_id, - model: conversationSchema.api_config.model, - provider: conversationSchema.api_config.provider, - }, - } - : {}), - excludeFromLastConversationStorage: conversationSchema.exclude_from_last_conversation_storage, - messages: - // eslint-disable-next-line @typescript-eslint/no-explicit-any - conversationSchema.messages?.map((message: Record) => ({ - timestamp: message['@timestamp'], - // always return anonymized data from the client - content: replaceOriginalValuesWithUuidValues({ - messageContent: message.content, - replacements, - }), - ...(message.is_error ? { isError: message.is_error } : {}), - ...(message.reader ? { reader: message.reader } : {}), - role: message.role, - ...(message.trace_data - ? { - traceData: { - traceId: message.trace_data?.trace_id, - transactionId: message.trace_data?.transaction_id, - }, - } - : {}), - })), - updatedAt: conversationSchema.updated_at, - replacements, - namespace: conversationSchema.namespace, - id: conversationSchema.id, - }; - - return conversation; - }); + return response.map((conversationSchema) => transformESToConversation(conversationSchema)); }; export const transformFieldNamesToSourceScheme = (fields: string[]) => { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts index 81332bd9176ee..5bee0fcb2a5f4 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/types.ts @@ -40,6 +40,9 @@ export interface EsConversationSchema { transaction_id?: string; trace_id?: string; }; + metadata?: { + content_references?: unknown; + }; }>; api_config?: { connector_id: string; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts index 3bb84c4a5bffe..18fbf0fd60c2f 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts @@ -5,16 +5,14 @@ * 2.0. */ -import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; import { loggerMock } from '@kbn/logging-mocks'; -import { - UpdateConversationSchema, - transformToUpdateScheme, - updateConversation, -} from './update_conversation'; -import { getConversation } from './get_conversation'; +import type { ConversationUpdateProps } from '@kbn/elastic-assistant-common'; + +import type { UpdateConversationSchema } from './update_conversation'; +import { transformToUpdateScheme, updateConversation } from './update_conversation'; +import type { EsConversationSchema } from './types'; import { authenticatedUser } from '../../__mocks__/user'; -import { ConversationResponse, ConversationUpdateProps } from '@kbn/elastic-assistant-common'; +import type { DocumentsDataWriter } from '../../lib/data_stream/documents_data_writer'; export const getUpdateConversationOptionsMock = (): ConversationUpdateProps => ({ id: 'test', @@ -32,47 +30,70 @@ export const getUpdateConversationOptionsMock = (): ConversationUpdateProps => ( }); const mockUser1 = authenticatedUser; - -export const getConversationResponseMock = (): ConversationResponse => ({ - id: 'test', - title: 'test', - apiConfig: { - actionTypeId: '.gen-ai', - connectorId: '1', - defaultSystemPromptId: 'default-system-prompt', - model: 'test-model', - provider: 'OpenAI', - }, - category: 'assistant', - excludeFromLastConversationStorage: false, - messages: [ - { - content: 'Message 3', - role: 'user', - timestamp: '2024-02-14T22:29:43.862Z', +const userAsUser = { + id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', + name: 'elastic', +}; +const getEsConversationMock = (): EsConversationSchema => { + return { + '@timestamp': '2025-08-19T10:49:52.884Z', + updated_at: '2025-08-19T13:26:01.746Z', + api_config: { + action_type_id: '.gen-ai', + connector_id: 'gpt-4-1', }, - { - content: 'Message 4', - role: 'user', - timestamp: '2024-02-14T22:29:43.862Z', - }, - ], - replacements: {}, - createdAt: '2020-04-20T15:25:31.830Z', - namespace: 'default', - updatedAt: '2020-04-20T15:25:31.830Z', - timestamp: '2020-04-20T15:25:31.830Z', - users: [ - { - id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', - name: 'elastic', - }, - ], -}); + namespace: 'default', + created_at: '2025-08-19T10:49:52.884Z', + messages: [ + { + '@timestamp': '2025-08-19T10:49:53.799Z', + role: 'user', + content: 'Hello there, how many opened alerts do I have?', + }, + { + metadata: { + content_references: { + oQ5xL: { + id: 'oQ5xL', + type: 'SecurityAlertsPage', + }, + }, + }, + '@timestamp': '2025-08-19T10:49:57.398Z', + role: 'assistant', + is_error: false, + trace_data: { + transaction_id: 'ee432e8be6ad3f9c', + trace_id: 'f44d01b6095d35dce15aa8137df76e29', + }, + content: 'You currently have 61 open alerts in your environment. {reference(oQ5xL)}', + }, + ], + replacements: [], + title: 'Viewing the Number of Open Alerts in Elastic Security', + category: 'assistant', + users: [userAsUser], + id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + }; +}; + +const getNothingToUpdateErrorResponseMock = () => { + return { + errors: [ + { + status_code: 500, + conversations: [{ id: '', name: '' }], + message: + 'null_pointer_exception\n\tRoot causes:\n\t\tnull_pointer_exception: Cannot invoke "org.elasticsearch.xcontent.XContentType.xContent()" because "xContentType" is null', + }, + ], + docs_updated: [], + }; +}; -jest.mock('./get_conversation', () => ({ - getConversation: jest.fn(), -})); +const dataWriterMock = { + bulk: jest.fn(), +} as unknown as DocumentsDataWriter; describe('updateConversation', () => { beforeEach(() => { @@ -83,42 +104,117 @@ describe('updateConversation', () => { jest.clearAllMocks(); }); + test('it calls a `dataWriter.bulk` with the correct parameters', async () => { + const conversation: ConversationUpdateProps = getUpdateConversationOptionsMock(); + const updatedESConversation = getEsConversationMock(); + + (dataWriterMock.bulk as jest.Mock).mockResolvedValue({ + errors: [], + docs_updated: [updatedESConversation], + }); + + await updateConversation({ + conversationUpdateProps: conversation, + dataWriter: dataWriterMock, + logger: loggerMock.create(), + user: mockUser1, + }); + + expect(dataWriterMock.bulk).toHaveBeenCalledWith({ + documentsToUpdate: [ + { + api_config: { + action_type_id: '.gen-ai', + connector_id: '1', + default_system_prompt_id: 'default-system-prompt', + model: 'test-model', + provider: 'OpenAI', + }, + exclude_from_last_conversation_storage: false, + id: 'test', + messages: [], + replacements: [], + title: 'test', + updated_at: expect.anything(), + }, + ], + getUpdateScript: expect.anything(), + authenticatedUser: mockUser1, + }); + }); + test('it returns a conversation with serializer and deserializer', async () => { const conversation: ConversationUpdateProps = getUpdateConversationOptionsMock(); - const existingConversation = getConversationResponseMock(); - (getConversation as unknown as jest.Mock).mockResolvedValueOnce(existingConversation); + const updatedESConversation = getEsConversationMock(); - const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; - esClient.updateByQuery.mockResolvedValue({ updated: 1 }); + (dataWriterMock.bulk as jest.Mock).mockResolvedValue({ + errors: [], + docs_updated: [updatedESConversation], + }); const updatedList = await updateConversation({ - esClient, - logger: loggerMock.create(), - conversationIndex: 'index-1', conversationUpdateProps: conversation, + dataWriter: dataWriterMock, + logger: loggerMock.create(), user: mockUser1, }); - const expected: ConversationResponse = { - ...getConversationResponseMock(), - id: conversation.id, - title: 'test', - }; - expect(updatedList).toEqual(expected); + + expect(updatedList).toEqual({ + timestamp: '2025-08-19T10:49:52.884Z', + createdAt: '2025-08-19T10:49:52.884Z', + users: [userAsUser], + title: 'Viewing the Number of Open Alerts in Elastic Security', + category: 'assistant', + apiConfig: { + actionTypeId: '.gen-ai', + connectorId: 'gpt-4-1', + }, + messages: [ + { + timestamp: '2025-08-19T10:49:53.799Z', + content: 'Hello there, how many opened alerts do I have?', + role: 'user', + }, + { + timestamp: '2025-08-19T10:49:57.398Z', + content: 'You currently have 61 open alerts in your environment. {reference(oQ5xL)}', + role: 'assistant', + metadata: { + contentReferences: { + oQ5xL: { + id: 'oQ5xL', + type: 'SecurityAlertsPage', + }, + }, + }, + traceData: { + traceId: 'f44d01b6095d35dce15aa8137df76e29', + transactionId: 'ee432e8be6ad3f9c', + }, + }, + ], + updatedAt: '2025-08-19T13:26:01.746Z', + replacements: {}, + namespace: 'default', + id: 'a565baa8-5566-47b2-ab69-807248b2fc46', + }); }); test('it returns null when there is not a conversation to update', async () => { - (getConversation as unknown as jest.Mock).mockResolvedValueOnce(null); const conversation = getUpdateConversationOptionsMock(); + (dataWriterMock.bulk as jest.Mock).mockResolvedValue(getNothingToUpdateErrorResponseMock()); - const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + const mockedLogger = loggerMock.create(); const updatedList = await updateConversation({ - esClient, - logger: loggerMock.create(), - conversationIndex: 'index-1', conversationUpdateProps: conversation, + dataWriter: dataWriterMock, + logger: mockedLogger, user: mockUser1, }); expect(updatedList).toEqual(null); + expect(mockedLogger.warn).toHaveBeenCalledWith( + expect.stringContaining('Error updating conversation: null_pointer_exception') + ); }); }); @@ -129,8 +225,6 @@ describe('transformToUpdateScheme', () => { test('it returns a transformed conversation with converted string datetime to ISO from the client', async () => { const conversation: ConversationUpdateProps = getUpdateConversationOptionsMock(); - const existingConversation = getConversationResponseMock(); - (getConversation as unknown as jest.Mock).mockResolvedValueOnce(existingConversation); const updateAt = new Date().toISOString(); const transformed = transformToUpdateScheme(updateAt, { @@ -210,8 +304,6 @@ describe('transformToUpdateScheme', () => { }); test('it does not pass api_config if apiConfig is not updated', async () => { const conversation: ConversationUpdateProps = getUpdateConversationOptionsMock(); - const existingConversation = getConversationResponseMock(); - (getConversation as unknown as jest.Mock).mockResolvedValueOnce(existingConversation); const updateAt = new Date().toISOString(); const transformed = transformToUpdateScheme(updateAt, { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts index 2d1aa0acfc4a9..2dd40b8429b3a 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { AuthenticatedUser, ElasticsearchClient, Logger } from '@kbn/core/server'; +import { AuthenticatedUser, Logger } from '@kbn/core/server'; import { ConversationResponse, Reader, @@ -16,9 +16,10 @@ import { UUID, ContentReferences, } from '@kbn/elastic-assistant-common'; -import { getConversation } from './get_conversation'; import { getUpdateScript } from './helpers'; -import { EsReplacementSchema } from './types'; +import { EsConversationSchema, EsReplacementSchema } from './types'; +import type { DocumentsDataWriter } from '../../lib/data_stream/documents_data_writer'; +import { transformESToConversations } from './transforms'; export interface UpdateConversationSchema { id: UUID; @@ -52,84 +53,39 @@ export interface UpdateConversationSchema { } export interface UpdateConversationParams { - esClient: ElasticsearchClient; + conversationUpdateProps: ConversationUpdateProps; + dataWriter: DocumentsDataWriter; logger: Logger; user?: AuthenticatedUser; - conversationIndex: string; - conversationUpdateProps: ConversationUpdateProps; - isPatch?: boolean; } export const updateConversation = async ({ - esClient, - logger, - conversationIndex, conversationUpdateProps, - isPatch, + dataWriter, + logger, user, }: UpdateConversationParams): Promise => { const updatedAt = new Date().toISOString(); const params = transformToUpdateScheme(updatedAt, conversationUpdateProps); - const maxRetries = 3; - let attempt = 0; - let response; - while (attempt < maxRetries) { - try { - response = await esClient.updateByQuery({ - conflicts: 'proceed', - index: conversationIndex, - query: { - ids: { - values: [params.id], - }, - }, - refresh: true, - script: getUpdateScript({ conversation: params, isPatch }).script, - }); - if ( - (response?.updated && response?.updated > 0) || - (response?.failures && response?.failures.length > 0) - ) { - break; - } - if ( - response?.version_conflicts && - response?.version_conflicts > 0 && - response?.updated === 0 - ) { - attempt++; - if (attempt < maxRetries) { - logger.warn( - `Version conflict detected, retrying updateConversation (attempt ${ - attempt + 1 - }) for conversation ID: ${params.id}` - ); - await new Promise((resolve) => setTimeout(resolve, 100 * attempt)); - } - } else { - break; - } - } catch (err) { - logger.warn(`Error updating conversation: ${err} by ID: ${params.id}`); - throw err; - } - } + const { errors, docs_updated: docsUpdated } = await dataWriter.bulk({ + documentsToUpdate: [params], + getUpdateScript: (document: UpdateConversationSchema) => + getUpdateScript({ conversation: document }), + authenticatedUser: user, + }); - if (response && response?.failures && response?.failures.length > 0) { + if (errors && errors.length > 0) { logger.warn( - `Error updating conversation: ${response?.failures.map((f) => f.id)} by ID: ${params.id}` + `Error updating conversation: ${errors.map((err) => err.message)} by ID: ${params.id}` ); return null; } - const updatedConversation = await getConversation({ - esClient, - conversationIndex, - id: params.id, - logger, - user, - }); + const updatedConversation = transformESToConversations( + docsUpdated as EsConversationSchema[] + )?.[0]; + return updatedConversation; }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 2e9d466aa9088..88389928a29bb 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -200,13 +200,14 @@ export class DocumentsDataWriter implements DocumentsDataWriter { _id: document.id, _index: responseToUpdate?.hits.hits.find((c) => c._id === document.id)?._index, _source: true, + retry_on_conflict: 3, }, }, getUpdateScript(document, updatedAt), ]); }; - private getDeletedocumentsQuery = async ( + private getDeleteDocumentsQuery = async ( documentsToDelete: string[], authenticatedUser?: AuthenticatedUser ) => { @@ -259,7 +260,7 @@ export class DocumentsDataWriter implements DocumentsDataWriter { const documentDeletedBody = params.documentsToDelete && params.documentsToDelete.length > 0 - ? await this.getDeletedocumentsQuery(params.documentsToDelete, params.authenticatedUser) + ? await this.getDeleteDocumentsQuery(params.documentsToDelete, params.authenticatedUser) : []; const documentUpdatedBody = diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts index 5dc308840630e..185f9fb775d2b 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts @@ -14,9 +14,11 @@ import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; import { coreMock } from '@kbn/core/server/mocks'; import { INVOKE_ASSISTANT_ERROR_EVENT } from '../../lib/telemetry/event_based_telemetry'; import { PassThrough } from 'stream'; -import { getConversationResponseMock } from '../../ai_assistant_data_clients/conversations/update_conversation.test'; import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock'; -import { getFindAnonymizationFieldsResultWithSingleHit } from '../../__mocks__/response'; +import { + getConversationResponseMock, + getFindAnonymizationFieldsResultWithSingleHit, +} from '../../__mocks__/response'; import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { chatCompleteRoute } from './chat_complete_route'; import { licensingMock } from '@kbn/licensing-plugin/server/mocks'; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts index 2526d67e5ce20..5e604f4d6fbfd 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts @@ -15,9 +15,11 @@ import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; import { coreMock } from '@kbn/core/server/mocks'; import { INVOKE_ASSISTANT_ERROR_EVENT } from '../lib/telemetry/event_based_telemetry'; import { PassThrough } from 'stream'; -import { getConversationResponseMock } from '../ai_assistant_data_clients/conversations/update_conversation.test'; import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock'; -import { getFindAnonymizationFieldsResultWithSingleHit } from '../__mocks__/response'; +import { + getConversationResponseMock, + getFindAnonymizationFieldsResultWithSingleHit, +} from '../__mocks__/response'; import { defaultAssistantFeatures, ExecuteConnectorRequestBody, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts index b66ea1baf5d76..0d241306a1f3c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts @@ -209,7 +209,7 @@ export const bulkActionConversationsRoute = ( documentsToUpdate: body.update?.map((c) => transformToUpdateScheme(changedAt, c)), authenticatedUser, getUpdateScript: (document: UpdateConversationSchema) => - getUpdateScript({ conversation: document, isPatch: true }), + getUpdateScript({ conversation: document }), }); const created = docsCreated.length > 0