[Security AI Assistant] Persist prompts

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

@@ -52,3 +52,6 @@ export * from './knowledge_base/bulk_crud_knowledge_base_route.gen';
 export * from './knowledge_base/common_attributes.gen';
 export * from './knowledge_base/crud_knowledge_base_route.gen';
 export * from './knowledge_base/find_knowledge_base_entries_route.gen';
+
+export * from './prompts/find_prompts_route.gen';
+export { PromptResponse, PromptTypeEnum } from './prompts/bulk_crud_prompts_route.gen';

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts

@@ -34,6 +34,14 @@ export const PromptDetailsInError = z.object({
   name: z.string().optional(),
 });
 
+/**
+ * Prompt type
+ */
+export type PromptType = z.infer<typeof PromptType>;
+export const PromptType = z.enum(['system', 'quick']);
+export type PromptTypeEnum = typeof PromptType.enum;
+export const PromptTypeEnum = PromptType.enum;
+
 export type NormalizedPromptError = z.infer<typeof NormalizedPromptError>;
 export const NormalizedPromptError = z.object({
   message: z.string(),
@@ -47,11 +55,13 @@ export const PromptResponse = z.object({
   id: NonEmptyString,
   timestamp: NonEmptyString.optional(),
   name: z.string(),
-  promptType: z.string(),
+  promptType: PromptType,
   content: z.string(),
+  categories: z.array(z.string()).optional(),
+  color: z.string().optional(),
   isNewConversationDefault: z.boolean().optional(),
   isDefault: z.boolean().optional(),
-  isShared: z.boolean().optional(),
+  consumer: z.string().optional(),
   updatedAt: z.string().optional(),
   updatedBy: z.string().optional(),
   createdAt: z.string().optional(),
@@ -107,20 +117,24 @@ export const BulkActionBase = z.object({
 export type PromptCreateProps = z.infer<typeof PromptCreateProps>;
 export const PromptCreateProps = z.object({
   name: z.string(),
-  promptType: z.string(),
+  promptType: PromptType,
   content: z.string(),
+  color: z.string().optional(),
+  categories: z.array(z.string()).optional(),
   isNewConversationDefault: z.boolean().optional(),
   isDefault: z.boolean().optional(),
-  isShared: z.boolean().optional(),
+  consumer: z.string().optional(),
 });
 
 export type PromptUpdateProps = z.infer<typeof PromptUpdateProps>;
 export const PromptUpdateProps = z.object({
   id: z.string(),
   content: z.string().optional(),
+  color: z.string().optional(),
+  categories: z.array(z.string()).optional(),
   isNewConversationDefault: z.boolean().optional(),
   isDefault: z.boolean().optional(),
-  isShared: z.boolean().optional(),
+  consumer: z.string().optional(),
 });
 
 export type PerformBulkActionRequestBody = z.infer<typeof PerformBulkActionRequestBody>;

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml

@@ -78,6 +78,13 @@ components:
       required:
         - id
 
+    PromptType:
+      type: string
+      description: Prompt type
+      enum:
+        - system
+        - quick
+
     NormalizedPromptError:
       type: object
       properties:
@@ -111,15 +118,21 @@ components:
         name:
           type: string
         promptType:
-          type: string
+          $ref: '#/components/schemas/PromptType'
         content:
           type: string
+        categories:
+          type: array
+          items:
+            type: string
+        color:
+          type: string
         isNewConversationDefault:
           type: boolean
         isDefault:
           type: boolean
-        isShared:
-          type: boolean
+        consumer:
+          type: string
         updatedAt:
           type: string
         updatedBy:
@@ -231,15 +244,21 @@ components:
         name:
           type: string
         promptType:
-          type: string
+          $ref: '#/components/schemas/PromptType'
         content:
           type: string
+        color:
+          type: string
+        categories:
+          type: array
+          items:
+            type: string
         isNewConversationDefault:
           type: boolean
         isDefault:
           type: boolean
-        isShared:
-          type: boolean
+        consumer:
+          type: string
 
     PromptUpdateProps:
       type: object
@@ -250,9 +269,15 @@ components:
           type: string
         content:
           type: string
+        color:
+          type: string
+        categories:
+          type: array
+          items:
+            type: string
         isNewConversationDefault:
           type: boolean
         isDefault:
           type: boolean
-        isShared:
-          type: boolean
+        consumer:
+          type: string

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx

@@ -11,6 +11,7 @@ import { API_ERROR } from '../translations';
 import { getOptionalRequestParams } from '../helpers';
 import { TraceOptions } from '../types';
 export * from './conversations';
+export * from './prompts';
 
 export interface FetchConnectorExecuteAction {
   conversationId: string;

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/prompts/bulk_update_prompts.test.ts

@@ -0,0 +1,137 @@
+/*
+ * 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 {
+  API_VERSIONS,
+  ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION,
+} from '@kbn/elastic-assistant-common';
+import { httpServiceMock } from '@kbn/core-http-browser-mocks';
+import { IToasts } from '@kbn/core-notifications-browser';
+import { bulkUpdatePrompts } from './bulk_update_prompts';
+import { PromptTypeEnum } from '@kbn/elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen';
+
+const prompt1 = {
+  id: 'field1',
+  content: 'Prompt 1',
+  name: 'test',
+  promptType: PromptTypeEnum.system,
+};
+const prompt2 = {
+  ...prompt1,
+  id: 'field2',
+  content: 'Prompt 2',
+  name: 'test2',
+  promptType: PromptTypeEnum.system,
+};
+const toasts = {
+  addError: jest.fn(),
+};
+describe('bulkUpdatePrompts', () => {
+  let httpMock: ReturnType<typeof httpServiceMock.createSetupContract>;
+
+  beforeEach(() => {
+    httpMock = httpServiceMock.createSetupContract();
+
+    jest.clearAllMocks();
+  });
+  it('should send a POST request with the correct parameters and receive a successful response', async () => {
+    const promptsActions = {
+      create: [],
+      update: [],
+      delete: { ids: [] },
+    };
+
+    await bulkUpdatePrompts(httpMock, promptsActions);
+
+    expect(httpMock.fetch).toHaveBeenCalledWith(ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION, {
+      method: 'POST',
+      version: API_VERSIONS.internal.v1,
+      body: JSON.stringify({
+        create: [],
+        update: [],
+        delete: { ids: [] },
+      }),
+    });
+  });
+
+  it('should transform the prompts dictionary to an array of fields to create', async () => {
+    const promptsActions = {
+      create: [prompt1, prompt2],
+      update: [],
+      delete: { ids: [] },
+    };
+
+    await bulkUpdatePrompts(httpMock, promptsActions);
+
+    expect(httpMock.fetch).toHaveBeenCalledWith(ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION, {
+      method: 'POST',
+      version: API_VERSIONS.internal.v1,
+      body: JSON.stringify({
+        create: [prompt1, prompt2],
+        update: [],
+        delete: { ids: [] },
+      }),
+    });
+  });
+
+  it('should transform the prompts dictionary to an array of fields to update', async () => {
+    const promptsActions = {
+      update: [prompt1, prompt2],
+      delete: { ids: [] },
+    };
+
+    await bulkUpdatePrompts(httpMock, promptsActions);
+
+    expect(httpMock.fetch).toHaveBeenCalledWith(ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION, {
+      method: 'POST',
+      version: API_VERSIONS.internal.v1,
+      body: JSON.stringify({
+        update: [prompt1, prompt2],
+        delete: { ids: [] },
+      }),
+    });
+  });
+
+  it('should throw an error with the correct message when receiving an unsuccessful response', async () => {
+    httpMock.fetch.mockResolvedValue({
+      success: false,
+      attributes: {
+        errors: [
+          {
+            statusCode: 400,
+            message: 'Error updating prompt',
+            prompts: [{ id: prompt1.id, name: prompt1.content }],
+          },
+        ],
+      },
+    });
+    const promptsActions = {
+      create: [],
+      update: [prompt1],
+      delete: { ids: [] },
+    };
+    await bulkUpdatePrompts(httpMock, promptsActions, toasts as unknown as IToasts);
+    expect(toasts.addError.mock.calls[0][0]).toEqual(
+      new Error('Error message: Error updating prompt for prompt Prompt 1')
+    );
+  });
+
+  it('should handle cases where result.attributes.errors is undefined', async () => {
+    httpMock.fetch.mockResolvedValue({
+      success: false,
+      attributes: {},
+    });
+    const promptsActions = {
+      create: [],
+      update: [],
+      delete: { ids: [] },
+    };
+
+    await bulkUpdatePrompts(httpMock, promptsActions, toasts as unknown as IToasts);
+    expect(toasts.addError.mock.calls[0][0]).toEqual(new Error(''));
+  });
+});

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/prompts/bulk_update_prompts.ts

@@ -0,0 +1,54 @@
+/*
+ * 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 { i18n } from '@kbn/i18n';
+import { HttpSetup, IToasts } from '@kbn/core/public';
+import {
+  ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION,
+  API_VERSIONS,
+} from '@kbn/elastic-assistant-common';
+import {
+  PerformBulkActionRequestBody,
+  PerformBulkActionResponse,
+} from '@kbn/elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen';
+
+export const bulkUpdatePrompts = async (
+  http: HttpSetup,
+  prompts: PerformBulkActionRequestBody,
+  toasts?: IToasts
+) => {
+  try {
+    const result = await http.fetch<PerformBulkActionResponse>(
+      ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION,
+      {
+        method: 'POST',
+        version: API_VERSIONS.internal.v1,
+        body: JSON.stringify(prompts),
+      }
+    );
+
+    if (!result.success) {
+      const serverError = result.attributes.errors
+        ?.map(
+          (e) =>
+            `${e.status_code ? `Error code: ${e.status_code}. ` : ''}Error message: ${
+              e.message
+            } for prompt ${e.prompts.map((c) => c.name).join(',')}`
+        )
+        .join(',\n');
+      throw new Error(serverError);
+    }
+    return result;
+  } catch (error) {
+    toasts?.addError(error.body && error.body.message ? new Error(error.body.message) : error, {
+      title: i18n.translate('xpack.elasticAssistant.prompts.bulkActionspromptsError', {
+        defaultMessage: 'Error updating prompts {error}',
+        values: { error },
+      }),
+    });
+  }
+};

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/prompts/index.tsx

@@ -0,0 +1,9 @@
+/*
+ * 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 * from './bulk_update_prompts';
+export * from './use_fetch_prompts';