Skip to content

Commit 2e384fb

Browse files
stephmilovicstephmilovic
authored
[Security Assistant] Context prompts (#224956)
1 parent bf0003f commit 2e384fb

39 files changed

Lines changed: 1262 additions & 115 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,12 @@ export const ELASTIC_AI_ASSISTANT_ALERT_SUMMARY_URL_BULK_ACTION =
6767
export const ELASTIC_AI_ASSISTANT_ALERT_SUMMARY_URL_FIND =
6868
`${ELASTIC_AI_ASSISTANT_ALERT_SUMMARY_URL}/_find` as const;
6969

70+
// Security AI Prompts (prompt integration)
71+
export const ELASTIC_AI_ASSISTANT_SECURITY_AI_PROMPTS_URL =
72+
`${ELASTIC_AI_ASSISTANT_URL}/security_ai_prompts` as const;
73+
export const ELASTIC_AI_ASSISTANT_SECURITY_AI_PROMPTS_URL_FIND =
74+
`${ELASTIC_AI_ASSISTANT_SECURITY_AI_PROMPTS_URL}/_find` as const;
75+
7076
// Defend insights
7177
export const DEFEND_INSIGHTS_ID = 'defend-insights';
7278
export const DEFEND_INSIGHTS = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/defend_insights`;

‎x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/common_attributes.gen.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ export const BulkActionBase = z.object({
102102
});
103103

104104
/**
105-
* User screen context.
105+
* IDs for a specific prompt within a group of prompts.
106106
*/
107107
export type PromptIds = z.infer<typeof PromptIds>;
108108
export const PromptIds = z.object({

‎x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ components:
101101
- "5678"
102102

103103
PromptIds:
104-
description: User screen context.
104+
description: IDs for a specific prompt within a group of prompts.
105105
type: object
106106
required:
107107
- promptId

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,4 +75,6 @@ export * from './prompts/find_prompts_route.gen';
7575
export * from './alert_summary/bulk_crud_alert_summary_route.gen';
7676
export * from './alert_summary/find_alert_summary_route.gen';
7777

78+
export * from './security_ai_prompts/find_prompts_route.gen';
79+
7880
export { PromptResponse, PromptTypeEnum } from './prompts/bulk_crud_prompts_route.gen';

‎x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/security_ai_prompts/common_attributes.gen.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
/*
9+
* NOTICE: Do not edit this file manually.
10+
* This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator.
11+
*
12+
* info:
13+
* title: Common Security AI Prompts Attributes
14+
* version: not applicable
15+
*/
16+
17+
import { z } from '@kbn/zod';
18+
19+
export type PromptItem = z.infer<typeof PromptItem>;
20+
export const PromptItem = z.object({
21+
promptId: z.string(),
22+
prompt: z.string(),
23+
});
24+
25+
/**
26+
* Prompt array by prompt group id and prompt id.
27+
*/
28+
export type PromptItemArray = z.infer<typeof PromptItemArray>;
29+
export const PromptItemArray = z.array(PromptItem);

‎x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/security_ai_prompts/common_attributes.schema.yaml‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
openapi: 3.0.0
2+
info:
3+
title: Common Security AI Prompts Attributes
4+
version: 'not applicable'
5+
paths: {}
6+
components:
7+
x-codegen-enabled: true
8+
schemas:
9+
PromptItem:
10+
type: object
11+
properties:
12+
promptId:
13+
type: string
14+
example: systemPrompt
15+
prompt:
16+
type: string
17+
example: This is the prompt
18+
required:
19+
- promptId
20+
- prompt
21+
22+
PromptItemArray:
23+
type: array
24+
description: Prompt array by prompt group id and prompt id.
25+
items:
26+
$ref: '#/components/schemas/PromptItem'

‎x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/security_ai_prompts/find_prompts_route.gen.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
/*
9+
* NOTICE: Do not edit this file manually.
10+
* This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator.
11+
*
12+
* info:
13+
* title: Get Security AI Prompt endpoint
14+
* version: 2023-10-31
15+
*/
16+
17+
import { z } from '@kbn/zod';
18+
import { ArrayFromString } from '@kbn/zod-helpers';
19+
20+
import { PromptItemArray } from './common_attributes.gen';
21+
22+
export type FindSecurityAIPromptsRequestQuery = z.infer<typeof FindSecurityAIPromptsRequestQuery>;
23+
export const FindSecurityAIPromptsRequestQuery = z.object({
24+
/**
25+
* Connector id used for prompt lookup
26+
*/
27+
connector_id: z.string().optional(),
28+
/**
29+
* The unique identifier for the prompt group
30+
*/
31+
prompt_group_id: z.string(),
32+
/**
33+
* Comma-separated list of prompt IDs to retrieve
34+
*/
35+
prompt_ids: ArrayFromString(z.string()),
36+
});
37+
export type FindSecurityAIPromptsRequestQueryInput = z.input<
38+
typeof FindSecurityAIPromptsRequestQuery
39+
>;
40+
41+
export type FindSecurityAIPromptsResponse = z.infer<typeof FindSecurityAIPromptsResponse>;
42+
export const FindSecurityAIPromptsResponse = z.object({
43+
prompts: PromptItemArray,
44+
});

‎x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/security_ai_prompts/find_prompts_route.schema.yaml‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
openapi: 3.0.0
2+
info:
3+
title: Get Security AI Prompt endpoint
4+
version: '2023-10-31'
5+
paths:
6+
/api/security_ai_assistant/security_ai_prompts/_find:
7+
get:
8+
x-codegen-enabled: true
9+
x-internal: true
10+
x-labels: [ess, serverless]
11+
operationId: FindSecurityAIPrompts
12+
description: Gets Security AI prompts
13+
summary: Gets Security AI prompts
14+
tags:
15+
- Security AI Prompts API
16+
parameters:
17+
- name: connector_id
18+
in: query
19+
description: Connector id used for prompt lookup
20+
required: false
21+
schema:
22+
type: string
23+
- name: prompt_group_id
24+
in: query
25+
description: The unique identifier for the prompt group
26+
required: true
27+
schema:
28+
type: string
29+
- name: prompt_ids
30+
in: query
31+
description: Comma-separated list of prompt IDs to retrieve
32+
required: true
33+
style: form
34+
explode: true
35+
schema:
36+
type: array
37+
items:
38+
type: string
39+
responses:
40+
'200':
41+
description: Indicates a successful call.
42+
content:
43+
application/json:
44+
schema:
45+
type: object
46+
properties:
47+
prompts:
48+
$ref: './common_attributes.schema.yaml#/components/schemas/PromptItemArray'
49+
required:
50+
- prompts
51+
'400':
52+
description: Generic Error
53+
content:
54+
application/json:
55+
schema:
56+
type: object
57+
properties:
58+
status_code:
59+
type: number
60+
error:
61+
type: string
62+
message:
63+
type: string

‎x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/security_ai_prompts/use_find_prompts.test.tsx‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import { waitFor, renderHook } from '@testing-library/react';
9+
import { useFindPrompts, UseFindPromptsParams } from './use_find_prompts';
10+
import {
11+
API_VERSIONS,
12+
ELASTIC_AI_ASSISTANT_SECURITY_AI_PROMPTS_URL_FIND,
13+
} from '@kbn/elastic-assistant-common';
14+
import { TestProviders } from '../../../mock/test_providers/test_providers';
15+
import { IToasts } from '@kbn/core-notifications-browser';
16+
17+
const mockHttpFetch = jest.fn();
18+
const mockToasts = {
19+
addSuccess: jest.fn(),
20+
addError: jest.fn(),
21+
} as unknown as IToasts;
22+
23+
describe('useFindPrompts', () => {
24+
const params: UseFindPromptsParams = {
25+
context: {
26+
isAssistantEnabled: true,
27+
httpFetch: mockHttpFetch,
28+
toasts: mockToasts,
29+
},
30+
params: {
31+
connector_id: 'connector-1',
32+
prompt_ids: ['prompt-1'],
33+
prompt_group_id: 'group-1',
34+
},
35+
signal: undefined,
36+
};
37+
38+
beforeEach(() => {
39+
jest.clearAllMocks();
40+
});
41+
42+
it('returns initial and placeholder data when loading', async () => {
43+
mockHttpFetch.mockResolvedValueOnce({ prompts: [{ id: 'prompt-1', name: 'Prompt 1' }] });
44+
const { result } = renderHook(() => useFindPrompts(params), {
45+
wrapper: TestProviders,
46+
});
47+
expect(result.current.data).toEqual({ prompts: [] });
48+
await waitFor(() => result.current.isSuccess);
49+
50+
await waitFor(() =>
51+
expect(result.current.data).toEqual({ prompts: [{ id: 'prompt-1', name: 'Prompt 1' }] })
52+
);
53+
});
54+
55+
it('disables the query if isAssistantEnabled is false', async () => {
56+
const disabledParams = { ...params, context: { ...params.context, isAssistantEnabled: false } };
57+
const { result } = renderHook(() => useFindPrompts(disabledParams), { wrapper: TestProviders });
58+
expect(result.current.isLoading).toBe(false);
59+
expect(result.current.data).toEqual({ prompts: [] });
60+
expect(mockHttpFetch).not.toHaveBeenCalled();
61+
});
62+
63+
it('calls httpFetch with correct arguments', async () => {
64+
mockHttpFetch.mockResolvedValueOnce({ prompts: [{ id: 'prompt-1', name: 'Prompt 1' }] });
65+
const { result } = renderHook(() => useFindPrompts(params), {
66+
wrapper: TestProviders,
67+
});
68+
await waitFor(() => result.current.isSuccess);
69+
expect(mockHttpFetch).toHaveBeenCalledWith(
70+
ELASTIC_AI_ASSISTANT_SECURITY_AI_PROMPTS_URL_FIND,
71+
expect.objectContaining({
72+
method: 'GET',
73+
version: API_VERSIONS.public.v1,
74+
query: {
75+
connector_id: 'connector-1',
76+
prompt_ids: ['prompt-1'],
77+
prompt_group_id: 'group-1',
78+
},
79+
})
80+
);
81+
});
82+
83+
it('shows a toast on error', async () => {
84+
const error = { body: { message: 'Something went wrong' } };
85+
mockHttpFetch.mockRejectedValueOnce(error);
86+
const { result } = renderHook(() => useFindPrompts(params), {
87+
wrapper: TestProviders,
88+
});
89+
await waitFor(() => result.current.isError);
90+
expect(mockToasts.addError).toHaveBeenCalledWith(
91+
expect.any(Error),
92+
expect.objectContaining({
93+
title: expect.any(String),
94+
})
95+
);
96+
});
97+
});

0 commit comments

Comments
 (0)