[Security Solution] Security AI Assistant persistent storage.

x-pack/packages/kbn-elastic-assistant-common/constants.ts

@@ -0,0 +1,26 @@
+/*
+ * 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.
+ */
+
+export const ELASTIC_AI_ASSISTANT_API_CURRENT_VERSION = '2023-10-31';
+export const ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION = '1';
+
+export const ELASTIC_AI_ASSISTANT_URL = '/api/elastic_assistant';
+
+export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL = `${ELASTIC_AI_ASSISTANT_URL}/current_user/conversations`;
+export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID = `${ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL}/{id}`;
+export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID_MESSAGES = `${ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID}/messages`;
+
+export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL}/_bulk_action`;
+export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_FIND = `${ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL}/_find`;
+
+export const ELASTIC_AI_ASSISTANT_PROMPTS_URL = `${ELASTIC_AI_ASSISTANT_URL}/prompts`;
+export const ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_PROMPTS_URL}/_bulk_action`;
+export const ELASTIC_AI_ASSISTANT_PROMPTS_URL_FIND = `${ELASTIC_AI_ASSISTANT_PROMPTS_URL}/_find`;
+
+export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL = `${ELASTIC_AI_ASSISTANT_URL}/anonymization_fields`;
+export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL}/_bulk_action`;
+export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_FIND = `${ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL}/_find`;

x-pack/packages/kbn-elastic-assistant-common/impl/data_anonymization/helpers/index.ts

@@ -5,6 +5,8 @@
  * 2.0.
  */
 
+import { Replacement } from '../../schemas';
+
 export const getIsDataAnonymizable = (rawData: string | Record<string, string[]>): boolean =>
   typeof rawData !== 'string';
 
@@ -21,3 +23,31 @@ export const isAnonymized = ({
   allowReplacementSet: Set<string>;
   field: string;
 }): boolean => allowReplacementSet.has(field);
+
+export const replaceAnonymizedValuesWithOriginalValues = ({
+  messageContent,
+  replacements,
+}: {
+  messageContent: string;
+  replacements: Replacement[];
+}): string =>
+  replacements != null
+    ? replacements.reduce((acc, replacement) => {
+        const value = replacement.value;
+        return replacement.uuid && value ? acc.replaceAll(replacement.uuid, value) : acc;
+      }, messageContent)
+    : messageContent;
+
+export const replaceOriginalValuesWithUuidValues = ({
+  messageContent,
+  replacements,
+}: {
+  messageContent: string;
+  replacements: Replacement[];
+}): string =>
+  replacements != null
+    ? replacements.reduce((acc, replacement) => {
+        const value = replacement.value;
+        return replacement.uuid && value ? acc.replaceAll(value, replacement.uuid) : acc;
+      }, messageContent)
+    : messageContent;

x-pack/packages/kbn-elastic-assistant-common/impl/data_anonymization/transform_raw_data/index.test.tsx

@@ -20,7 +20,7 @@ describe('transformRawData', () => {
     const result = transformRawData({
       allow: inputRawData.allow,
       allowReplacement: inputRawData.allowReplacement,
-      currentReplacements: {},
+      currentReplacements: [],
       getAnonymizedValue: mockGetAnonymizedValue,
       onNewReplacements: () => {},
       rawData: inputRawData.rawData,
@@ -42,13 +42,13 @@ describe('transformRawData', () => {
     transformRawData({
       allow: inputRawData.allow,
       allowReplacement: inputRawData.allowReplacement,
-      currentReplacements: {},
+      currentReplacements: [],
       getAnonymizedValue: mockGetAnonymizedValue,
       onNewReplacements,
       rawData: inputRawData.rawData,
     });
 
-    expect(onNewReplacements).toHaveBeenCalledWith({ '1eulav': 'value1' });
+    expect(onNewReplacements).toHaveBeenCalledWith([{ uuid: '1eulav', value: 'value1' }]);
   });
 
   it('returns the expected mix of anonymized and non-anonymized data as a CSV string', () => {
@@ -62,7 +62,7 @@ describe('transformRawData', () => {
     const result = transformRawData({
       allow: inputRawData.allow,
       allowReplacement: inputRawData.allowReplacement,
-      currentReplacements: {},
+      currentReplacements: [],
       getAnonymizedValue: mockGetAnonymizedValue,
       onNewReplacements: () => {},
       rawData: inputRawData.rawData,
@@ -86,7 +86,7 @@ describe('transformRawData', () => {
     const result = transformRawData({
       allow: inputRawData.allow,
       allowReplacement: inputRawData.allowReplacement,
-      currentReplacements: {},
+      currentReplacements: [],
       getAnonymizedValue: mockGetAnonymizedValue,
       onNewReplacements: () => {},
       rawData: inputRawData.rawData,

x-pack/packages/kbn-elastic-assistant-common/impl/data_anonymization/transform_raw_data/index.tsx

@@ -5,6 +5,7 @@
  * 2.0.
  */
 
+import { Replacement } from '../../schemas';
 import { getAnonymizedData } from '../get_anonymized_data';
 import { getAnonymizedValues } from '../get_anonymized_values';
 import { getCsvFromData } from '../get_csv_from_data';
@@ -19,15 +20,15 @@ export const transformRawData = ({
 }: {
   allow: string[];
   allowReplacement: string[];
-  currentReplacements: Record<string, string> | undefined;
+  currentReplacements: Replacement[] | undefined;
   getAnonymizedValue: ({
     currentReplacements,
     rawValue,
   }: {
     currentReplacements: Record<string, string> | undefined;
     rawValue: string;
   }) => string;
-  onNewReplacements?: (replacements: Record<string, string>) => void;
+  onNewReplacements?: (replacements: Replacement[]) => void;
   rawData: string | Record<string, unknown[]>;
 }): string => {
   if (typeof rawData === 'string') {
@@ -37,14 +38,22 @@ export const transformRawData = ({
   const anonymizedData = getAnonymizedData({
     allow,
     allowReplacement,
-    currentReplacements,
+    currentReplacements: currentReplacements?.reduce((acc: Record<string, string>, r) => {
+      acc[r.uuid] = r.value;
+      return acc;
+    }, {}),
     rawData,
     getAnonymizedValue,
     getAnonymizedValues,
   });
 
   if (onNewReplacements != null) {
-    onNewReplacements(anonymizedData.replacements);
+    onNewReplacements(
+      Object.keys(anonymizedData.replacements).map((key) => ({
+        uuid: key,
+        value: anonymizedData.replacements[key],
+      }))
+    );
   }
 
   return getCsvFromData(anonymizedData.anonymizedData);

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts

@@ -0,0 +1,68 @@
+/*
+ * 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 { z } from 'zod';
+
+/*
+ * NOTICE: Do not edit this file manually.
+ * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator.
+ *
+ * info:
+ *   title: Execute Connector API endpoint
+ *   version: 1
+ */
+
+import { UUID, Replacement } from '../conversations/common_attributes.gen';
+
+export type ExecuteConnectorRequestParams = z.infer<typeof ExecuteConnectorRequestParams>;
+export const ExecuteConnectorRequestParams = z.object({
+  /**
+   * The connector's `id` value.
+   */
+  connectorId: z.string(),
+});
+export type ExecuteConnectorRequestParamsInput = z.input<typeof ExecuteConnectorRequestParams>;
+
+export type ExecuteConnectorRequestBody = z.infer<typeof ExecuteConnectorRequestBody>;
+export const ExecuteConnectorRequestBody = z.object({
+  conversationId: UUID.optional(),
+  message: z.string().optional(),
+  model: z.string().optional(),
+  subAction: z.enum(['invokeAI', 'invokeStream']),
+  alertsIndexPattern: z.string().optional(),
+  allow: z.array(z.string()).optional(),
+  allowReplacement: z.array(z.string()).optional(),
+  isEnabledKnowledgeBase: z.boolean().optional(),
+  isEnabledRAGAlerts: z.boolean().optional(),
+  replacements: z.array(Replacement),
+  size: z.number().optional(),
+  llmType: z.enum(['bedrock', 'openai']),
+});
+export type ExecuteConnectorRequestBodyInput = z.input<typeof ExecuteConnectorRequestBody>;
+
+export type ExecuteConnectorResponse = z.infer<typeof ExecuteConnectorResponse>;
+export const ExecuteConnectorResponse = z.object({
+  data: z.string().optional(),
+  connector_id: z.string().optional(),
+  replacements: z.array(Replacement).optional(),
+  status: z.string().optional(),
+  /**
+   * Trace Data
+   */
+  trace_data: z
+    .object({
+      /**
+       * Could be any string, not necessarily a UUID
+       */
+      transactionId: z.string().optional(),
+      /**
+       * Could be any string, not necessarily a UUID
+       */
+      traceId: z.string().optional(),
+    })
+    .optional(),
+});

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml

@@ -0,0 +1,110 @@
+openapi: 3.0.0
+info:
+  title: Execute Connector API endpoint
+  version: '1'
+paths:
+  /internal/elastic_assistant/actions/connector/{connectorId}/_execute:
+    post:
+      operationId: ExecuteConnector
+      x-codegen-enabled: true
+      description: Execute Elastic Assistant connector by id
+      summary: Execute Elastic Assistant connector
+      tags:
+        - Connector API
+      parameters:
+        - name: connectorId
+          in: path
+          required: true
+          description: The connector's `id` value.
+          schema:
+            type: string
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              required:
+                - params
+                - llmType
+                - replacements
+                - subAction
+              properties:
+                conversationId:
+                  $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/UUID'
+                message:
+                  type: string
+                model:
+                  type: string
+                subAction:
+                  type: string
+                  enum:
+                    - invokeAI
+                    - invokeStream
+                alertsIndexPattern:
+                  type: string
+                allow:
+                  type: array
+                  items:
+                    type: string
+                allowReplacement:
+                  type: array
+                  items:
+                    type: string
+                isEnabledKnowledgeBase:
+                  type: boolean
+                isEnabledRAGAlerts:
+                  type: boolean
+                replacements:
+                  type: array
+                  items:
+                    $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacement'
+                size:
+                  type: number
+                llmType:
+                  type: string
+                  enum:
+                    - bedrock
+                    - openai
+      responses:
+        '200':
+          description: Successful response
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  data:
+                    type: string
+                  connector_id:
+                    type: string
+                  replacements:
+                    type: array
+                    items:
+                      $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacement'
+                  status:
+                    type: string
+                  trace_data:
+                    type: object
+                    description: Trace Data
+                    properties:
+                      transactionId:
+                        type: string
+                        description: Could be any string, not necessarily a UUID
+                      traceId:
+                        type: string
+                        description: Could be any string, not necessarily a UUID
+        '400':
+          description: Generic Error
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  statusCode:
+                    type: number
+                  error:
+                    type: string
+                  message:
+                    type: string
+