[Security solution] AI Assistant Fix Up (#161784)

This commit is contained in:
Steph Milovic 2023-07-13 10:53:48 -06:00 committed by GitHub
parent 8f57280853
commit 11cf01ea43
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 210 additions and 152 deletions

View file

@ -0,0 +1,115 @@
/*
* 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 { getWelcomeConversation } from './helpers';
import { enterpriseMessaging } from './use_conversation/sample_conversations';
describe('getWelcomeConversation', () => {
describe('isAssistantEnabled = false', () => {
const isAssistantEnabled = false;
it('When no conversation history, return only enterprise messaging', () => {
const conversation = {
id: 'conversation_id',
theme: {},
messages: [],
apiConfig: {},
};
const result = getWelcomeConversation(conversation, isAssistantEnabled);
expect(result.messages).toEqual(enterpriseMessaging);
expect(result.messages.length).toEqual(1);
});
it('When conversation history and the last message is not enterprise messaging, appends enterprise messaging to conversation', () => {
const conversation = {
id: 'conversation_id',
theme: {},
messages: [
{
role: 'user' as const,
content: 'Hello',
timestamp: '',
presentation: {
delay: 0,
stream: false,
},
},
],
apiConfig: {},
};
const result = getWelcomeConversation(conversation, isAssistantEnabled);
expect(result.messages.length).toEqual(2);
});
it('returns the conversation without changes when the last message is enterprise messaging', () => {
const conversation = {
id: 'conversation_id',
theme: {},
messages: enterpriseMessaging,
apiConfig: {},
};
const result = getWelcomeConversation(conversation, isAssistantEnabled);
expect(result.messages.length).toEqual(1);
expect(result.messages).toEqual(enterpriseMessaging);
});
it('returns the conversation with new enterprise message when conversation has enterprise messaging, but not as the last message', () => {
const conversation = {
id: 'conversation_id',
theme: {},
messages: [
...enterpriseMessaging,
{
role: 'user' as const,
content: 'Hello',
timestamp: '',
presentation: {
delay: 0,
stream: false,
},
},
],
apiConfig: {},
};
const result = getWelcomeConversation(conversation, isAssistantEnabled);
expect(result.messages.length).toEqual(3);
});
});
describe('isAssistantEnabled = true', () => {
const isAssistantEnabled = true;
it('when no conversation history, returns the welcome conversation', () => {
const conversation = {
id: 'conversation_id',
theme: {},
messages: [],
apiConfig: {},
};
const result = getWelcomeConversation(conversation, isAssistantEnabled);
expect(result.messages.length).toEqual(3);
});
it('returns a conversation history with the welcome conversation appended', () => {
const conversation = {
id: 'conversation_id',
theme: {},
messages: [
{
role: 'user' as const,
content: 'Hello',
timestamp: '',
presentation: {
delay: 0,
stream: false,
},
},
],
apiConfig: {},
};
const result = getWelcomeConversation(conversation, isAssistantEnabled);
expect(result.messages.length).toEqual(4);
});
});
});

View file

@ -5,10 +5,9 @@
* 2.0.
*/
import { BASE_CONVERSATIONS, Conversation } from '../..';
import { Conversation } from '../..';
import type { Message } from '../assistant_context/types';
import { WELCOME_CONVERSATION_TITLE } from './use_conversation/translations';
import { enterpriseMessaging } from './use_conversation/sample_conversations';
import { enterpriseMessaging, WELCOME_CONVERSATION } from './use_conversation/sample_conversations';
export const getMessageFromRawResponse = (rawResponse: string): Message => {
const dateTimeString = new Date().toLocaleString(); // TODO: Pull from response
@ -27,13 +26,13 @@ export const getMessageFromRawResponse = (rawResponse: string): Message => {
}
};
export const getWelcomeConversation = (isAssistantEnabled: boolean): Conversation => {
const conversation = BASE_CONVERSATIONS[WELCOME_CONVERSATION_TITLE];
const doesConversationHaveMessages = conversation.messages.length > 0;
export const getWelcomeConversation = (
conversation: Conversation,
isAssistantEnabled: boolean
): Conversation => {
if (!isAssistantEnabled) {
if (
!doesConversationHaveMessages ||
conversation.messages.length === 0 ||
conversation.messages[conversation.messages.length - 1].content !==
enterpriseMessaging[0].content
) {
@ -47,6 +46,6 @@ export const getWelcomeConversation = (isAssistantEnabled: boolean): Conversatio
return {
...conversation,
messages: BASE_CONVERSATIONS[WELCOME_CONVERSATION_TITLE].messages,
messages: [...conversation.messages, ...WELCOME_CONVERSATION.messages],
};
};

View file

@ -26,6 +26,7 @@ import { css } from '@emotion/react';
import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/gen_ai/constants';
import { ActionConnectorProps } from '@kbn/triggers-actions-ui-plugin/public/types';
import { WELCOME_CONVERSATION_TITLE } from './use_conversation/translations';
import { AssistantTitle } from './assistant_title';
import { UpgradeButtons } from '../upgrade/upgrade_buttons';
import { getMessageFromRawResponse, getWelcomeConversation } from './helpers';
@ -46,7 +47,6 @@ import * as i18n from './translations';
import { QuickPrompts } from './quick_prompts/quick_prompts';
import { useLoadConnectors } from '../connectorland/use_load_connectors';
import { useConnectorSetup } from '../connectorland/connector_setup';
import { WELCOME_CONVERSATION_TITLE } from './use_conversation/translations';
import { AssistantSettingsButton } from './settings/assistant_settings_button';
import { ConnectorMissingCallout } from '../connectorland/connector_missing_callout';
@ -63,7 +63,7 @@ export interface Props {
* quick prompts for common actions, settings, and prompt context providers.
*/
const AssistantComponent: React.FC<Props> = ({
conversationId = WELCOME_CONVERSATION_TITLE,
conversationId,
isAssistantEnabled,
promptContextId = '',
shouldRefocusPrompt = false,
@ -80,9 +80,11 @@ const AssistantComponent: React.FC<Props> = ({
http,
promptContexts,
setLastConversationId,
localStorageLastConversationId,
title,
allSystemPrompts,
} = useAssistantContext();
const [selectedPromptContexts, setSelectedPromptContexts] = useState<
Record<string, SelectedPromptContext>
>({});
@ -109,23 +111,38 @@ const AssistantComponent: React.FC<Props> = ({
[connectors]
);
const [selectedConversationId, setSelectedConversationId] = useState<string>(
isAssistantEnabled
? // if a conversationId has been provided, use that
// if not, check local storage
// last resort, go to welcome conversation
conversationId ?? localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE
: WELCOME_CONVERSATION_TITLE
);
const currentConversation = useMemo(
() =>
conversations[selectedConversationId] ??
createConversation({ conversationId: selectedConversationId }),
[conversations, createConversation, selectedConversationId]
);
// Welcome setup state
const isWelcomeSetup = useMemo(() => (connectors?.length ?? 0) === 0, [connectors?.length]);
const isWelcomeSetup = useMemo(() => {
// if any conversation has a connector id, we're not in welcome set up
return Object.keys(conversations).some(
(conversation) => conversations[conversation].apiConfig.connectorId != null
)
? false
: (connectors?.length ?? 0) === 0;
}, [connectors?.length, conversations]);
const isDisabled = isWelcomeSetup || !isAssistantEnabled;
// Welcome conversation is a special 'setup' case when no connector exists, mostly extracted to `ConnectorSetup` component,
// but currently a bit of state is littered throughout the assistant component. TODO: clean up/isolate this state
const welcomeConversation = useMemo(
() => getWelcomeConversation(isAssistantEnabled),
[isAssistantEnabled]
);
const [selectedConversationId, setSelectedConversationId] = useState<string>(
isWelcomeSetup ? welcomeConversation.id : conversationId
);
const currentConversation = useMemo(
() => conversations[selectedConversationId] ?? createConversation({ conversationId }),
[conversationId, conversations, createConversation, selectedConversationId]
() => getWelcomeConversation(currentConversation, isAssistantEnabled),
[currentConversation, isAssistantEnabled]
);
// Remember last selection for reuse after keyboard shortcut is pressed.
@ -519,7 +536,7 @@ const AssistantComponent: React.FC<Props> = ({
<EuiModalBody>
{comments}
{!isWelcomeSetup && showMissingConnectorCallout && (
{!isDisabled && showMissingConnectorCallout && (
<>
<EuiSpacer />
<EuiFlexGroup justifyContent="spaceAround">
@ -569,8 +586,8 @@ const AssistantComponent: React.FC<Props> = ({
onPromptSubmit={handleSendMessage}
ref={promptTextAreaRef}
handlePromptChange={setPromptTextPreview}
value={isDisabled ? '' : suggestedUserPrompt ?? ''}
isDisabled={isDisabled}
value={isSendingDisabled ? '' : suggestedUserPrompt ?? ''}
isDisabled={isSendingDisabled}
/>
</EuiFlexItem>
@ -594,7 +611,7 @@ const AssistantComponent: React.FC<Props> = ({
<EuiButtonIcon
display="base"
iconType="cross"
isDisabled={isDisabled}
isDisabled={isSendingDisabled}
aria-label={i18n.CLEAR_CHAT}
color="danger"
onClick={() => {

View file

@ -11,15 +11,16 @@ import userEvent from '@testing-library/user-event';
import { mockSystemPrompt } from '../../../mock/system_prompt';
import { SystemPrompt } from '.';
import { BASE_CONVERSATIONS, Conversation } from '../../../..';
import { Conversation } from '../../../..';
import { DEFAULT_CONVERSATION_TITLE } from '../../use_conversation/translations';
import { Prompt } from '../../types';
import { TestProviders } from '../../../mock/test_providers/test_providers';
import { TEST_IDS } from '../../constants';
import { useAssistantContext } from '../../../assistant_context';
import { WELCOME_CONVERSATION } from '../../use_conversation/sample_conversations';
const BASE_CONVERSATION: Conversation = {
...BASE_CONVERSATIONS[DEFAULT_CONVERSATION_TITLE],
...WELCOME_CONVERSATION,
apiConfig: {
defaultSystemPromptId: mockSystemPrompt.id,
},

View file

@ -8,127 +8,55 @@
import { Conversation, Message } from '../../assistant_context/types';
import * as i18n from '../../content/prompts/welcome/translations';
import {
DEFAULT_CONVERSATION_TITLE,
ELASTIC_AI_ASSISTANT,
ELASTIC_AI_ASSISTANT_TITLE,
WELCOME_CONVERSATION_TITLE,
} from './translations';
export const BASE_CONVERSATIONS: Record<string, Conversation> = {
[DEFAULT_CONVERSATION_TITLE]: {
id: DEFAULT_CONVERSATION_TITLE,
messages: [
{
role: 'user',
content: 'Can you provide a sample KQL query please? (Kibana)',
timestamp: '5/12/2023, 12:04:07 AM',
},
{
role: 'assistant',
content:
'Certainly! Here\'s a sample KQL (Kibana Query Language) query for a fictional web server log data:\n\n```\nresponse.keyword: "200" AND extension.keyword: "html" AND NOT (geo.src: "US" OR geo.src: "CA")\n```\n\nThis query will return all documents where:\n\n1. The `response` field has a value of "200" (HTTP status code for a successful request).\n2. The `extension` field has a value of "html" (indicating an HTML file was requested).\n3. The `geo.src` field is not "US" or "CA" (meaning the request did not come from the United States or Canada).\n\nYou can use this query in Kibana\'s Discover or Dashboard panels to filter the data based on these conditions.',
timestamp: '5/12/2023, 12:04:21 AM',
},
{
role: 'user',
content:
'What about an EQL sequence query for 4 failed logons followed by a successful login?',
timestamp: '5/12/2023, 12:08:31 AM',
},
{
role: 'assistant',
content:
'Here\'s an EQL (Event Query Language) sequence query to detect 4 failed logons followed by a successful login for a fictional logon event data:\n\n```eql\nsequence by user.name with maxspan=5m\n [authentication where result == "failure" and event.action == "logon_attempt"]{4}\n [authentication where result == "success" and event.action == "logon_attempt"]\n```\n\nThis query will return sequences where:\n\n1. The events are grouped by the `user.name` field.\n2. The maximum time span between the first and last event in the sequence is 5 minutes (`maxspan=5m`).\n3. There are 4 consecutive `authentication` events with a `result` field value of "failure" and an `event.action` field value of "logon_attempt".\n4. After the 4 failed logon attempts, there\'s an `authentication` event with a `result` field value of "success" and an `event.action` field value of "logon_attempt".\n\nYou can use this EQL query in Elasticsearch or Kibana to detect potential brute force attacks or suspicious logon behavior.',
timestamp: '5/12/2023, 12:08:46 AM',
},
{
role: 'user',
content: "Can you fix the markdown please, it's not rendering",
timestamp: '5/12/2023, 12:09:00 AM',
},
{
role: 'assistant',
content:
'Apologies for the formatting issue. Here\'s the corrected EQL sequence query:\n\n```\nsequence by user.name with maxspan=5m\n [authentication where result == "failure" and event.action == "logon_attempt"]{4}\n [authentication where result == "success" and event.action == "logon_attempt"]\n```\n\nThis query will return sequences where:\n\n1. The events are grouped by the `user.name` field.\n2. The maximum time span between the first and last event in the sequence is 5 minutes (`maxspan=5m`).\n3. There are 4 consecutive `authentication` events with a `result` field value of "failure" and an `event.action` field value of "logon_attempt".\n4. After the 4 failed logon attempts, there\'s an `authentication` event with a `result` field value of "success" and an `event.action` field value of "logon_attempt".\n\nYou can use this EQL query in Elasticsearch or Kibana to detect potential brute force attacks or suspicious logon behavior.',
timestamp: '5/12/2023, 12:09:16 AM',
},
{
role: 'user',
content:
'Can you response with a message that contains both the first KQL query you sent me, and also the EQL sequence query you sent?',
timestamp: '5/12/2023, 12:11:49 AM',
},
{
role: 'assistant',
content:
'Certainly! Here are both the KQL and EQL queries in a single response:\n\n**KQL Query (Kibana Query Language):**\n\n```\nresponse.keyword: "200" AND extension.keyword: "html" AND NOT (geo.src: "US" OR geo.src: "CA")\n```\n\nThis KQL query returns all documents where:\n\n1. The `response` field has a value of "200" (HTTP status code for a successful request).\n2. The `extension` field has a value of "html" (indicating an HTML file was requested).\n3. The `geo.src` field is not "US" or "CA" (meaning the request did not come from the United States or Canada).\n\nYou can use this query in Kibana\'s Discover or Dashboard panels to filter the data based on these conditions.\n\n**EQL Sequence Query (Event Query Language):**\n\n```\nsequence by user.name with maxspan=5m\n [authentication where result == "failure" and event.action == "logon_attempt"]{4}\n [authentication where result == "success" and event.action == "logon_attempt"]\n```\n\nThis EQL query returns sequences where:\n\n1. The events are grouped by the `user.name` field.\n2. The maximum time span between the first and last event in the sequence is 5 minutes (`maxspan=5m`).\n3. There are 4 consecutive `authentication` events with a `result` field value of "failure" and an `event.action` field value of "logon_attempt".\n4. After the 4 failed logon attempts, there\'s an `authentication` event with a `result` field value of "success" and an `event.action` field value of "logon_attempt".\n\nYou can use this EQL query in Elasticsearch or Kibana to detect potential brute force attacks or suspicious logon behavior.',
timestamp: '5/12/2023, 12:12:09 AM',
},
{
role: 'user',
content:
'Last but not least, could you write an Elasticsearch QueryDSL for filtering events where `event.action` is `logon_attempt`',
timestamp: '5/12/2023, 12:30:38 AM',
},
{
role: 'assistant',
content:
'Certainly! Here\'s an Elasticsearch QueryDSL for filtering events where `file.name` is `fake_behavior.exe`:\n\n```json\n{\n "match_phrase": {\n "file.name": {\n "query": "fake_behavior.exe"\n }\n }\n}\n```\n\nThis QueryDSL is searching for documents where the file.name field matches the exact phrase "fake_behavior.exe". It is useful for finding documents that have a specific file name in an Elasticsearch index.',
timestamp: '5/12/2023, 12:30:49 AM',
},
],
apiConfig: {},
},
timeline: {
id: 'timeline',
messages: [],
apiConfig: {},
},
[WELCOME_CONVERSATION_TITLE]: {
id: WELCOME_CONVERSATION_TITLE,
theme: {
title: ELASTIC_AI_ASSISTANT_TITLE,
titleIcon: 'logoSecurity',
assistant: {
name: ELASTIC_AI_ASSISTANT,
icon: 'logoSecurity',
},
system: {
icon: 'logoElastic',
},
user: {},
export const WELCOME_CONVERSATION: Conversation = {
id: WELCOME_CONVERSATION_TITLE,
theme: {
title: ELASTIC_AI_ASSISTANT_TITLE,
titleIcon: 'logoSecurity',
assistant: {
name: ELASTIC_AI_ASSISTANT,
icon: 'logoSecurity',
},
messages: [
{
role: 'assistant',
content: i18n.WELCOME_GENERAL,
timestamp: '',
presentation: {
delay: 2 * 1000,
stream: true,
},
},
{
role: 'assistant',
content: i18n.WELCOME_GENERAL_2,
timestamp: '',
presentation: {
delay: 1000,
stream: true,
},
},
{
role: 'assistant',
content: i18n.WELCOME_GENERAL_3,
timestamp: '',
presentation: {
delay: 1000,
stream: true,
},
},
],
apiConfig: {},
system: {
icon: 'logoElastic',
},
user: {},
},
messages: [
{
role: 'assistant',
content: i18n.WELCOME_GENERAL,
timestamp: '',
presentation: {
delay: 2 * 1000,
stream: true,
},
},
{
role: 'assistant',
content: i18n.WELCOME_GENERAL_2,
timestamp: '',
presentation: {
delay: 1000,
stream: true,
},
},
{
role: 'assistant',
content: i18n.WELCOME_GENERAL_3,
timestamp: '',
presentation: {
delay: 1000,
stream: true,
},
},
],
apiConfig: {},
};
export const enterpriseMessaging: Message[] = [

View file

@ -20,7 +20,8 @@ import {
OpenAiProviderType,
} from '@kbn/stack-connectors-plugin/public/common';
import { ActionConnectorProps } from '@kbn/triggers-actions-ui-plugin/public/types';
import { BASE_CONVERSATIONS, Conversation, Message } from '../../..';
import { WELCOME_CONVERSATION } from '../../assistant/use_conversation/sample_conversations';
import { Conversation, Message } from '../../..';
import { useLoadActionTypes } from '../use_load_action_types';
import { StreamingText } from '../../assistant/streaming_text';
import { ConnectorButton } from '../connector_button';
@ -28,7 +29,6 @@ import { useConversation } from '../../assistant/use_conversation';
import { clearPresentationData, conversationHasNoPresentationData } from './helpers';
import * as i18n from '../translations';
import { useAssistantContext } from '../../assistant_context';
import { WELCOME_CONVERSATION_TITLE } from '../../assistant/use_conversation/translations';
const ConnectorButtonWrapper = styled.div`
margin-bottom: 10px;
@ -53,7 +53,7 @@ export interface ConnectorSetupProps {
export const useConnectorSetup = ({
actionTypeRegistry,
conversation = BASE_CONVERSATIONS[WELCOME_CONVERSATION_TITLE],
conversation = WELCOME_CONVERSATION,
http,
isConnectorConfigured = false,
onSetupComplete,

View file

@ -71,9 +71,10 @@ export { useAssistantOverlay } from './impl/assistant/use_assistant_overlay';
/** a helper that enriches content returned from a query with action buttons */
export { analyzeMarkdown } from './impl/assistant/use_conversation/helpers';
// Sample content is exported with the following:
/** sample content */
export { BASE_CONVERSATIONS } from './impl/assistant/use_conversation/sample_conversations';
export {
ELASTIC_AI_ASSISTANT_TITLE,
WELCOME_CONVERSATION_TITLE,
} from './impl/assistant/use_conversation/translations';
/** i18n translations of system prompts */
export * as SYSTEM_PROMPTS from './impl/content/prompts/system/translations';

View file

@ -5,10 +5,7 @@
* 2.0.
*/
import {
ELASTIC_AI_ASSISTANT_TITLE,
WELCOME_CONVERSATION_TITLE,
} from '@kbn/elastic-assistant/impl/assistant/use_conversation/translations';
import { ELASTIC_AI_ASSISTANT_TITLE, WELCOME_CONVERSATION_TITLE } from '@kbn/elastic-assistant';
import type { Conversation } from '@kbn/elastic-assistant';
import { DATA_QUALITY_DASHBOARD_CONVERSATION_ID } from '@kbn/ecs-data-quality-dashboard/impl/data_quality/data_quality_panel/tabs/summary_tab/callout_summary/translations';
import { DETECTION_RULES_CONVERSATION_ID } from '../../../detections/pages/detection_engine/rules/translations';