Skip to content

Commit

Permalink
[Security Solution] [Elastic AI Assistant] Include acknowledged alert…
Browse files Browse the repository at this point in the history
…s in the context sent to the LLM (Retrieval Augmented Generation (RAG) for Alerts) (#173121)

## [Security Solution] [Elastic AI Assistant] Include `acknowledged` alerts in the context sent to the LLM (Retrieval Augmented Generation (RAG) for Alerts)

This PR updates the query used by [[Security Solution] [Elastic AI Assistant] Retrieval Augmented Generation (RAG) for Alerts #172542](#172542) to include alerts with a `kibana.alert.workflow_status` value of `acknowledged`.

The query previously only returned alerts with a status of `open`. This change ensures both `open` and `acknowledged` alerts are provided as context to the LLM.

### Updated Anonymization defaults

Three fields, detailed below, were added as anonymization defaults because they improve the quality of responses from the LLM when it answers questions about alerts.

For example, the LLM can refer to specific alerts by ID when the `_id` field is provided.

This PR makes the following additive changes to the Assistant's `Anonymization` defaults:

| Field                          | Allow by default | Anonymize by default | Value add                                                                                                                                                                                                                                                                                                                         |
|--------------------------------|------------------|----------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `_id`                          | ✅                | ✅                    | An anonymized `_id` field enables responses from the LLM to refer to specific documents (but doesn't provide it the actual document IDs).                                                                                                                                                                                         |
| `kibana.alert.risk_score`      | ✅                | ❌                    | The `getOpenAndAcknowledgedAlertsQuery` query sorts alerts by `kibana.alert.risk_score` to return the `n` riskiest alerts.  Allowing this field (by default) enables the LLM to include actual alert risk scores in responses.                                                                                                    |
| `kibana.alert.workflow_status` | ✅                | ❌                    | The `getOpenAndAcknowledgedAlertsQuery` query filters alerts by `kibana.alert.workflow_status` to ensure only `open` and `acknowledged` alerts are provided as context to the LLM.  Allowing this field (by default) enables the LLM answer questions about workflow status, and echo the workflow status of alerts in responses. |

- Clicking the `Reset` button shown in the screenshot below will reset the user's `Anonymization` defaults, such that they include the additive changes in the table above:

![01__id_allowed_and_anonymized_by_default](https://github.com/elastic/kibana/assets/4459398/7b8bcfa0-deb3-478b-b32f-c7fe3cf039b0)

### Updated settings text

The text in the settings below was also updated:

![updated_settings](https://github.com/elastic/kibana/assets/4459398/3b2d9f93-f89f-410a-b943-dd4c98e37a29)

### Desk testing

To desk test this change:

- Enable the `assistantRagOnAlerts` feature flag described in [#172542](#172542) must be enabled, per the following example:

```
xpack.securitySolution.enableExperimental: ['assistantRagOnAlerts']
```

- The `Alerts` feature must be enabled in the assistant settings, per the screenshot below:

 ![alerts_enabled](https://github.com/elastic/kibana/assets/4459398/edd273c8-b9e7-4ecf-93bb-81fbf0f3e6b0)

1) Navigate to Security > Alerts

2) Click the `AI Assistant` button to open the assistant

3) Click the `Settings` gear to open the assistant settings

4) Click the `Anonymization` category

5) Click the `Reset` button, shown in the screenshot below

![01__id_allowed_and_anonymized_by_default](https://github.com/elastic/kibana/assets/4459398/7b8bcfa0-deb3-478b-b32f-c7fe3cf039b0)

**Expected results**

- `65` fields are allowed by default, per the screenshot above
- `12` fields are anonymized by default, per the screenshot above
- The `_id` field is allowed by default, per the screenshot above
- The `_id` field is anonymized by default, per the screenshot above

6) Type `kibana.alert.risk` in the search box

**Expected result**

- The `kibana.alert.risk_score` field is allowed by default

7) Type `kibana.alert.workflow` in the search box

**Expected result**

- The `kibana.alert.workflow_status` field is allowed by default

8) Click `Save`

9) Click the `X` button to clear the conversation

10) Close the assistant

11) Add the following two fields as columns to the Alerts page table:

- `kibana.alert.workflow_status`
- `_id`

12) Sort the table, first by `kibana.alert.risk_score` from high to low, and then by `@timestamp` from new to old, per the screenshot below:

![fields_sorted](https://github.com/elastic/kibana/assets/4459398/e84f06d4-790d-4227-afbf-a233d4848178)

13) Filter the alerts page to only show `open` and `acknowledged` alerts

**Expected result**

- The alerts page has custom columns, sorting, and filtering, per the screenshot below:

![alerts_page_custom_filter_sort_and_columns](https://github.com/elastic/kibana/assets/4459398/ba191f83-ee4c-4c90-be4f-197eff6be68e)

14) Click the `AI Assistant` button to open the assistant

15) Ask the assistant:

```
What is the workflow status of my alerts?
```

**Expected result**

- The assistant will report on the workflow status of alerts, per the example response below:

```
The workflow status for your alerts is currently 'open'. This status was observed on alerts related to processes started by Mimikatz, a known tool used in many cyberattacks, and sequences of processes that are often indicative of malicious activity. The severity of most of these alerts is 'high' or 'critical'. You may want to investigate these alerts further to ensure there's no ongoing threat to your system.
```

![assistant_open_alerts](https://github.com/elastic/kibana/assets/4459398/93730bfe-9d6a-44f6-8270-14b6dd141514)

16) Close the assistant

17) Change the workflow status of an alert in the Alerts table from `open` to `acknowledged`

**Expected result**

- The alerts table shows the updated alert, per the screenshot below:

![acknowledged_alert](https://github.com/elastic/kibana/assets/4459398/df685381-ffd5-4253-897a-4058f1641747)

18) Once again, open the assistant

19) Once again, ask the (same) question:

```
What is the workflow status of my alerts?
```

**Expected result**

- The response from the assistant makes reference to the alert who's workflow status was changed from `open` to `acknowledged`, per the example response below:

```
Based on the latest information, your alerts mainly show 'open' status, indicating that they have not been resolved yet. Some alerts have been acknowledged. Most of these unaddressed alerts have a critical severity rating and are primarily triggered by a Mimikatz process start and an EQL process sequence. You may want to prioritize these if the severity of the threat they pose is truly high or critical. It's also noteworthy that some alerts have a high severity rating. You should review all of these alerts as soon as possible to ensure your systems are secure.
```

![assistant_some_alerts_acknowledged](https://github.com/elastic/kibana/assets/4459398/53bef5b6-1661-480e-bbf1-bdb14202480e)

20) Ask the assistant for details about the acknowledged alerts:

```
What are the details of the acknowledged alerts?
```

**Expected result**

The assistant for details about the acknowledged alert that, for example, includes the `kibana.alert.risk_score`, per the example response below:

```
In response to your previous question, here are the details of the acknowledged alerts:

1. There is a 'mimikatz process started' alert, which is of 'critical' severity and 'acknowledged' status. It has a high risk score of 99. Its threat tactic is 'Command and Control'. The process involved was 'mimikatz.exe' running with arguments '--fo1'.

2. A 'Threshold rule' alert of 'critical' severity and 'open' status has also been detected with a risk score of 99 and threat tactic 'Collection'.

3. Lastly, there are several 'EQL process sequence' alerts of 'high' severity with 'open' status. These alerts involve execution of various processes including 'mimikatz.exe', 'lsass.exe', and 'notepad.exe'. Risk score for these alerts is 73 and the threat tactic involved is 'Execution'.

Please, take appropriate action to address these alerts.
```

![details_of_acknowledged_alert](https://github.com/elastic/kibana/assets/4459398/191400e3-f123-4a0d-8ebc-9b5c88f242ed)

21) Ask the assistant for the `_id` of the acknowledged alert:

```
What is the id of the acknowledged alert?
```

**Expected results**

- The response from the assistant contains the `_id` of the `acknowledged` alert, per the example response below:

```
The id of the acknowledged alert is 'db9e3dbaf40a37e3b7b95d8015e99c5721b416731e04b9140536675f6e4fd170'. This alert was for a 'mimikatz process started' event with a severity rating of 'critical' and a risk score of 99. The host name associated with this alert is 'Host-terkvbzvtj'.
```

![id_of_acknowledged_alert](https://github.com/elastic/kibana/assets/4459398/b360b957-b198-4c9f-a101-581589b279bd)

- The `_id` shown in the assistant is the same `_id` of the acknowledged alert on the alerts page, per the screeenshot below:

![id_of_alert](https://github.com/elastic/kibana/assets/4459398/b9d4ef55-83ee-46cb-ab32-35bc467f4150)

22) Click the `Show anonymized` toggle in the assistant

**Expected result**

- The `_id` shown in the latest result is replaced with the actual anonymized value that was sent to the LLM, per the example screenshot below:

![show_anonymized_alert_id](https://github.com/elastic/kibana/assets/4459398/8ce47dd6-4bc6-4478-87c9-97d24b6eae82)
  • Loading branch information
andrew-goldstein authored Dec 13, 2023
1 parent 2f69f1b commit 0d9c261
Show file tree
Hide file tree
Showing 10 changed files with 92 additions and 26 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,18 @@ const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSetting
<span>{i18n.LATEST_AND_RISKIEST_OPEN_ALERTS}</span>
</EuiText>
</EuiFlexItem>

<EuiFlexItem grow={false}>
<EuiText color="subdued" size="xs">
<span>{i18n.YOUR_ANONYMIZATION_SETTINGS}</span>
</EuiText>
</EuiFlexItem>

<EuiFlexItem grow={false}>
<EuiText color="subdued" size="xs">
<span>{i18n.SELECT_FEWER_ALERTS}</span>
</EuiText>
</EuiFlexItem>
</EuiFlexGroup>
</>
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,23 @@ export const ASK_QUESTIONS_ABOUT = i18n.translate(

export const LATEST_AND_RISKIEST_OPEN_ALERTS = i18n.translate(
'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.latestAndRiskiestOpenAlertsLabel',
{
defaultMessage: 'latest and riskiest open and acknowledged alerts in your environment.',
}
);

export const YOUR_ANONYMIZATION_SETTINGS = i18n.translate(
'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.yourAnonymizationSettingsLabel',
{
defaultMessage: 'Your Anonymization settings will be applied to the alerts.',
}
);

export const SELECT_FEWER_ALERTS = i18n.translate(
'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.selectFewerAlertsLabel',
{
defaultMessage:
'latest and riskiest open alerts in your environment. Your Anonymization settings will be applied to the alerts',
"Select fewer alerts if the model's maximum context length is frequently exceeded.",
}
);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import { Tool } from 'langchain/tools';

import { getAlertCountsTool } from './alert_counts/get_alert_counts_tool';
import { getEsqlLanguageKnowledgeBaseTool } from './esql_language_knowledge_base/get_esql_language_knowledge_base_tool';
import { getOpenAlertsTool } from './open_alerts/get_open_alerts_tool';
import { getOpenAndAcknowledgedAlertsTool } from './open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool';
import type { RequestBody } from '../types';

export interface GetApplicableTools {
Expand Down Expand Up @@ -50,7 +50,7 @@ export const getApplicableTools = ({
replacements,
request,
}) ?? [],
getOpenAlertsTool({
getOpenAndAcknowledgedAlertsTool({
alertsIndexPattern,
allow,
allowReplacement,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,15 @@
* 2.0.
*/

import { getOpenAlertsQuery } from './get_open_alerts_query';
import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query';

describe('getOpenAlertsQuery', () => {
describe('getOpenAndAcknowledgedAlertsQuery', () => {
it('returns the expected query', () => {
const alertsIndexPattern = 'alerts-*';
const allow = ['field1', 'field2'];
const size = 10;

const query = getOpenAlertsQuery({ alertsIndexPattern, allow, size });
const query = getOpenAndAcknowledgedAlertsQuery({ alertsIndexPattern, allow, size });

expect(query).toEqual({
allow_no_indices: true,
Expand All @@ -30,8 +30,20 @@ describe('getOpenAlertsQuery', () => {
must: [],
filter: [
{
match_phrase: {
'kibana.alert.workflow_status': 'open',
bool: {
should: [
{
match_phrase: {
'kibana.alert.workflow_status': 'open',
},
},
{
match_phrase: {
'kibana.alert.workflow_status': 'acknowledged',
},
},
],
minimum_should_match: 1,
},
},
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
* 2.0.
*/

export const getOpenAlertsQuery = ({
export const getOpenAndAcknowledgedAlertsQuery = ({
alertsIndexPattern,
allow,
size,
Expand All @@ -28,8 +28,20 @@ export const getOpenAlertsQuery = ({
must: [],
filter: [
{
match_phrase: {
'kibana.alert.workflow_status': 'open',
bool: {
should: [
{
match_phrase: {
'kibana.alert.workflow_status': 'open',
},
},
{
match_phrase: {
'kibana.alert.workflow_status': 'acknowledged',
},
},
],
minimum_should_match: 1,
},
},
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,12 @@ import type { KibanaRequest } from '@kbn/core-http-server';
import { DynamicTool } from 'langchain/tools';
import { omit } from 'lodash/fp';

import { getOpenAlertsTool } from './get_open_alerts_tool';
import { getOpenAndAcknowledgedAlertsTool } from './get_open_and_acknowledged_alerts_tool';
import { mockAlertsFieldsApi } from '../../../../__mocks__/alerts';
import type { RequestBody } from '../../types';
import { MAX_SIZE } from './helpers';

describe('getOpenAlertsTool', () => {
describe('getOpenAndAcknowledgedAlertsTool', () => {
const alertsIndexPattern = 'alerts-index';
const esClient = {
search: jest.fn().mockResolvedValue(mockAlertsFieldsApi),
Expand All @@ -37,7 +37,7 @@ describe('getOpenAlertsTool', () => {
});

it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => {
const tool: DynamicTool = getOpenAlertsTool({
const tool: DynamicTool = getOpenAndAcknowledgedAlertsTool({
alertsIndexPattern,
allow: request.body.allow,
allowReplacement: request.body.allowReplacement,
Expand Down Expand Up @@ -75,8 +75,20 @@ describe('getOpenAlertsTool', () => {
bool: {
filter: [
{
match_phrase: {
'kibana.alert.workflow_status': 'open',
bool: {
should: [
{
match_phrase: {
'kibana.alert.workflow_status': 'open',
},
},
{
match_phrase: {
'kibana.alert.workflow_status': 'acknowledged',
},
},
],
minimum_should_match: 1,
},
},
{
Expand Down Expand Up @@ -130,7 +142,7 @@ describe('getOpenAlertsTool', () => {
RequestBody
>;

const tool = getOpenAlertsTool({
const tool = getOpenAndAcknowledgedAlertsTool({
alertsIndexPattern,
allow: requestWithMissingParams.body.allow,
allowReplacement: requestWithMissingParams.body.allowReplacement,
Expand All @@ -145,7 +157,7 @@ describe('getOpenAlertsTool', () => {
});

it('returns null when alertsIndexPattern is undefined', () => {
const tool = getOpenAlertsTool({
const tool = getOpenAndAcknowledgedAlertsTool({
// alertsIndexPattern is undefined
allow: request.body.allow,
allowReplacement: request.body.allowReplacement,
Expand All @@ -160,7 +172,7 @@ describe('getOpenAlertsTool', () => {
});

it('returns null when size is undefined', () => {
const tool = getOpenAlertsTool({
const tool = getOpenAndAcknowledgedAlertsTool({
alertsIndexPattern,
allow: request.body.allow,
allowReplacement: request.body.allowReplacement,
Expand All @@ -175,7 +187,7 @@ describe('getOpenAlertsTool', () => {
});

it('returns null when size out of range', () => {
const tool = getOpenAlertsTool({
const tool = getOpenAndAcknowledgedAlertsTool({
alertsIndexPattern,
allow: request.body.allow,
allowReplacement: request.body.allowReplacement,
Expand All @@ -190,7 +202,7 @@ describe('getOpenAlertsTool', () => {
});

it('returns a tool instance with the expected tags', () => {
const tool = getOpenAlertsTool({
const tool = getOpenAndAcknowledgedAlertsTool({
alertsIndexPattern,
allow: request.body.allow,
allowReplacement: request.body.allowReplacement,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,17 +13,17 @@ import { DynamicTool, Tool } from 'langchain/tools';
import { requestHasRequiredAnonymizationParams } from '../../helpers';
import { RequestBody } from '../../types';

import { getOpenAlertsQuery } from './get_open_alerts_query';
import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query';
import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers';

export const OPEN_ALERTS_TOOL_DESCRIPTION =
'Call this for knowledge about the latest n open alerts (sorted by `kibana.alert.risk_score`) in the environment, or when answering questions about open alerts';

/**
* Returns a tool for querying open alerts, or null if the request
* doesn't have all the required parameters.
* Returns a tool for querying open and acknowledged alerts, or null if the
* request doesn't have all the required parameters.
*/
export const getOpenAlertsTool = ({
export const getOpenAndAcknowledgedAlertsTool = ({
alertsIndexPattern,
allow,
allowReplacement,
Expand Down Expand Up @@ -55,7 +55,7 @@ export const getOpenAlertsTool = ({
name: 'open-alerts',
description: OPEN_ALERTS_TOOL_DESCRIPTION,
func: async () => {
const query = getOpenAlertsQuery({
const query = getOpenAndAcknowledgedAlertsQuery({
alertsIndexPattern,
allow: allow ?? [],
size,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

/** By default, these fields are allowed to be sent to the assistant */
export const DEFAULT_ALLOW = [
'_id',
'@timestamp',
'cloud.availability_zone',
'cloud.provider',
Expand All @@ -28,6 +29,7 @@ export const DEFAULT_ALLOW = [
'host.risk.calculated_level',
'host.risk.calculated_score_norm',
'kibana.alert.last_detected',
'kibana.alert.risk_score',
'kibana.alert.rule.description',
'kibana.alert.rule.name',
'kibana.alert.rule.references',
Expand All @@ -42,6 +44,7 @@ export const DEFAULT_ALLOW = [
'kibana.alert.rule.threat.technique.subtechnique.name',
'kibana.alert.rule.threat.technique.subtechnique.reference',
'kibana.alert.severity',
'kibana.alert.workflow_status',
'process.args',
'process.command_line',
'process.executable',
Expand Down Expand Up @@ -73,6 +76,7 @@ export const DEFAULT_ALLOW = [

/** By default, these fields will be anonymized */
export const DEFAULT_ALLOW_REPLACEMENT = [
'_id', // the document's _id is replaced with an anonymized value
'cloud.availability_zone',
'cloud.provider',
'cloud.region',
Expand Down

0 comments on commit 0d9c261

Please sign in to comment.