Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andrew-goldstein merged 4 commits into from
Dec 13, 2023
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
85e32af
## [Security Solution] [Elastic AI Assistant] Include `acknowledged` …
andrew-goldstein Dec 11, 2023
36150e6
- pr feedback
andrew-goldstein Dec 12, 2023
bdf02b8
- updated the description in settings to include acknowledged alerts
andrew-goldstein Dec 13, 2023
0982303
- added "Select fewer alerts if the model's maximum context length is…
andrew-goldstein Dec 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.tsx
View file
Open in desktop
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
16 changes: 15 additions & 1 deletion x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/translations.ts
View file
Open in desktop
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
4 changes: 2 additions & 2 deletions x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts
View file
Open in desktop
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
22 changes: 17 additions & 5 deletions ...open_alerts/get_open_alerts_query.test.ts → ...pen_and_acknowledged_alerts_query.test.ts
View file
Open in desktop
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
18 changes: 15 additions & 3 deletions ...ools/open_alerts/get_open_alerts_query.ts → ...get_open_and_acknowledged_alerts_query.ts
View file
Open in desktop
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
32 changes: 22 additions & 10 deletions .../open_alerts/get_open_alerts_tool.test.ts → ...open_and_acknowledged_alerts_tool.test.ts
View file
Open in desktop
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
10 changes: 5 additions & 5 deletions ...tools/open_alerts/get_open_alerts_tool.ts → .../get_open_and_acknowledged_alerts_tool.ts
View file
Open in desktop
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';
Copy link
Copy Markdown
Member

@spong spong Dec 13, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: open an acknowledged

Found in managing the merge w/ #172234. I got this 👍


/**
* 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
0 ...ngchain/tools/open_alerts/helpers.test.ts → ...n_and_acknowledged_alerts/helpers.test.ts
View file
Open in desktop
File renamed without changes.
0 ...ib/langchain/tools/open_alerts/helpers.ts → ...s/open_and_acknowledged_alerts/helpers.ts
View file
Open in desktop
File renamed without changes.
4 changes: 4 additions & 0 deletions x-pack/plugins/security_solution/public/assistant/content/anonymization/index.ts
View file
Open in desktop
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',
Copy link
Copy Markdown
Contributor Author

@andrew-goldstein andrew-goldstein Dec 13, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

'@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