[8.19] [Security Solution][AI Assistant] Handle conflicts during the bulk update (#233288)

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: [
     {

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,
   };
 };

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 { BulkRequest } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 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',
+              },
+            },
+          ],
+        },
       },
-    });
+    ]);
   });
 });

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.ts

@@ -35,7 +35,7 @@ export class AIAssistantConversationsDataClient extends AIAssistantDataClient {
   }
 
   /**
-   * Updates a conversation with the new messages.
+   * Gets a conversation by its id.
    * @param options
    * @param options.id The existing conversation id.
    * @param options.authenticatedUser Current authenticated user.
@@ -127,19 +127,15 @@ export class AIAssistantConversationsDataClient extends AIAssistantDataClient {
   public updateConversation = async ({
     conversationUpdateProps,
     authenticatedUser,
-    isPatch,
   }: {
     conversationUpdateProps: ConversationUpdateProps;
     authenticatedUser?: AuthenticatedUser;
-    isPatch?: boolean;
   }): Promise<ConversationResponse | null> => {
-    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,
     });
   };