[Security Solution] [Elastic AI Assistant] Retrieval Augmented Generation (RAG) for Alerts

[andrew-goldstein]

[Security Solution] [Elastic AI Assistant] Retrieval Augmented Generation (RAG) for Alerts

This PR implements Retrieval Augmented Generation (RAG) for Alerts in the Security Solution. This feature enables users to ask the assistant questions about the latest and riskiest open alerts in their environment using natural language, for example:

• How many alerts are currently open?
• Which alerts should I look at first?
• Did we have any alerts with suspicious activity on Windows machines?

More context

Previously, the assistant relied solely on the knowledge of the configured LLM and singular alerts or events passed by the client to the LLM as prompt context. This new feature:

• Enables multiple alerts to be passed by the server as context to the LLM, via LangChain tools
• Applies the user's anonymization settings to those alerts
  • Only fields allowed by the user will be sent as context to the LLM
  • Users may enable or disable anonymization for specific fields (via settings)
  • Click the conversation's Show anonymized toggle to see the anonymized values sent to / received from the LLM:
    

Settings

This feature is enabled and configured via the Knowledge Base > Alerts settings in the screenshot below:

• The Alerts toggle enables or disables the feature
• The slider has a range of 10 - 100 alerts (default: 20)

When the setting above is enabled, up to n alerts (as determined by the slider) that meet the following criteria will be returned:

• the kibana.alert.workflow_status must be open
• the alert must have been generated in the last 24 hours
• the alert must NOT be a kibana.alert.building_block_type alert
• the n alerts are ordered by kibana.alert.risk_score, to prioritize the riskiest alerts

Feature flag

To use this feature:

1. Add the assistantRagOnAlerts feature flag to the xpack.securitySolution.enableExperimental setting in config/kibana.yml (or config/kibana.dev.yml in local development environments), per the example below:

xpack.securitySolution.enableExperimental: ['assistantRagOnAlerts']

2. Enable the Alerts toggle in the Assistant's Knowledge Base settings, per the screenshot below:

How it works

• When the Alerts settings toggle is enabled, http POST requests to the /internal/elastic_assistant/actions/connector/{id}/_execute route include the following new (optional) parameters:
  • alertsIndexPattern, the alerts index for the current Kibana Space, e.g. .alerts-security.alerts-default
  • allow, the user's Allowed fields in the Anonymization settings, e.g. ["@timestamp", "cloud.availability_zone", "file.name", "user.name", ...]
  • allowReplacement, the user's Anonymized fields in the Anonymization settings, e.g. ["cloud.availability_zone", "host.name", "user.name", ...]
  • replacements, a Record<string, string> of replacements (generated on the server) that starts empty for a new conversation, and accumulates anonymized values until the conversation is cleared, e.g.

"replacements": {
    "e4f935c0-5a80-47b2-ac7f-816610790364": "Host-itk8qh4tjm",
    "cf61f946-d643-4b15-899f-6ffe3fd36097": "rpwmjvuuia",
    "7f80b092-fb1a-48a2-a634-3abc61b32157": "6astve9g6s",
    "f979c0d5-db1b-4506-b425-500821d00813": "Host-odqbow6tmc",
    // ...
},

• size, the numeric value set by the slider in the user's Knowledge Base > Alerts setting, e.g. 20

• The postActionsConnectorExecuteRoute function in x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts was updated to accept the new optional parameters, and to return an updated replacements with every response. (Every new request that is processed on the server may add additional anonymized values to the replacements returned in the response.)

• The callAgentExecutor function in x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts previously used a hard-coded array of LangChain tools that had just one entry, for the ESQLKnowledgeBaseTool tool. That hard-coded array was replaced in this PR with a call to the (new) getApplicableTools function:

  const tools: Tool[] = getApplicableTools({
    allow,
    allowReplacement,
    alertsIndexPattern,
    assistantLangChain,
    chain,
    esClient,
    modelExists,
    onNewReplacements,
    replacements,
    request,
    size,
  });

• The getApplicableTools function in x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts examines the parameters in the KibanaRequest and only returns a filtered set of LangChain tools. If the request doesn't contain all the parameters required by a tool, it will NOT be returned by getApplicableTools. For example, if the required anonymization parameters are not included in the request, the open-alerts tool will not be returned.

• The new alert-counts LangChain tool returned by the getAlertCountsTool function in x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts provides the LLM the results of an aggregation on the last 24 hours of alerts (in the current Kibana Space), grouped by kibana.alert.severity. See the getAlertsCountQuery function in x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_query.ts for details

• The new open-alerts LangChain tool returned by the getOpenAlertsTool function in x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.ts provides the LLM up to size non-building-block alerts generated in the last 24 hours (in the current Kibana Space) with an open workflow status, ordered by kibana.alert.risk_score to prioritize the riskiest alerts. See the getOpenAlertsQuery function in x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_query.ts for details.

• On the client, a conversation continues to accumulate additional replacements (and send them in subsequent requests) until the conversation is cleared

• Anonymization functions that were only invoked by the browser were moved from the (browser) kbn-elastic-assistant package in x-pack/packages/kbn-elastic-assistant/ to a new common package: x-pack/packages/kbn-elastic-assistant-common

  • The new kbn-elastic-assistant-common package is also consumed by the elastic_assistant (server) plugin: x-pack/plugins/elastic_assistant

[elasticmachine]

Pinging @elastic/security-solution (Team: SecuritySolution)

[andrew-goldstein]

@elasticmachine merge upstream

[stephmilovic]

Should we instead consider moving the package code into plugins/elastic_assistant/public and this directory would be plugins/elastic_assistant/common? Are there arguments to be made still that it should be a package?

[andrew-goldstein]

Should we instead consider moving the package code into plugins/elastic_assistant/public and this directory would be plugins/elastic_assistant/common? Are there arguments to be made still that it should be a package?

It feels like moving the common code into plugins/elastic_assistant/public would create a circular dependency, because the elastic_assistant's x-pack/plugins/elastic_assistant/tsconfig.json already has a reference to the package @kbn/elastic-assistant:

  "kbn_references": [
     // ...
    "@kbn/elastic-assistant", // <--existing reference
    "@kbn/elastic-assistant-common", // <-- new reference
     // ...
  ],

The anonymization functions moving to the new elastic-assistant-common are pure functions that are truly common / usable on both the client and server. (They are used on both in this PR.) Some of the helper functions may also be applicable in other contexts where anonymization is a useful feature.

It therefore feels like moving these anonymization functions to the new elastic-assistant-common is package prudent, because:

• It does not create a circular dependency
• The code is truly shared (by client and server) common functions, as described by the @kbn/elastic-assistant-common package's x-pack/packages/kbn-elastic-assistant-common/kibana.jsonc:

{
  "id": "@kbn/elastic-assistant-common",
  "owner": "@elastic/security-solution",
  "type": "shared-common" // <-- truly shared by client and server
}

• The common functions may be applicable in other contexts that require anonymization. Given that Kibana Packages are a lighter weight abstraction for sharing code compared to Plugins (because Plugins have, for example, lifecycle), it feels prudent to make these potentially reusable functions available via the lighter weight (Package) abstraction, given they are pure functions that don't require the capabilities of a Plugin.

Most of the types imported today from @kbn/elastic-assistant by the elastic_assistant plugin may be moved to the new @kbn/elastic-assistant-common package (in a future PR). If that happened, it should be possible to remove the @kbn/elastic-assistant dependency from x-pack/plugins/elastic_assistant/tsconfig.json.

[spong]

Thanks for including testing instructions!

[spong]

Thank you for updating and making these types accurate 👍

[spong]

nit: consider more human readable name for understanding role here and disambiguating ragOnAlerts: (the FF itself). canQueryAlerts, alertQueryingEnabled?

[spong]

note: this is the fix for the existing bug on main where errors were not bubbling up. Ty for the fix!

[spong]

Ahh! Here's the details on the naming from above -- perfect! 🙂

[spong]

Thanks for the further abstraction here 👍

[spong]

note: Not written to individual Message in localstorage, but rather the entire conversation as intended

[spong]

nite: lotta alert data, if uniqueness not super necessary maybe templatize one and modify severity with accessor?

[spong]

note: replacement test to come in a follow-up since control flow logic has changed here.

[andrew-goldstein]

the following replacement test:

  it('returns null if modelExists is false (the ELSER model is not installed)', () => {
    const tool = getEsqlLanguageKnowledgeBaseTool({
      assistantLangChain: true,
      chain,
      modelExists: false, // <-- ELSER model is not installed
    });

    expect(tool).toBeNull();
  });

is now added to x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts.

The replacement test asserts the new expected behavior:

• when the ELSER model is not installed (modelExists is false) the getEsqlLanguageKnowledgeBaseTool returns null
• an error will NOT thrown when the ELSER model is not installed

[spong]

nit: maybe bundle anonymization fields into an object for readability (incl: onNewReplacements, replacements)

[spong]

Yahs! Tool abstraction is here, thank you! 🎉

[spong]

nit: as discussed when pairing, at some point we'll scrape the dependencies off the request and use those instead (perhaps I'll do this when wrapping up the tool registration work :).

[spong]

nit: kinda weird this has to get injected since we're already in 'langchain land', but makes sense as required by getApplicableTools. This should get cleaner as we rework tool registration, feature flags, and remove the non-langchain path.

[kibana-ci]

💛 Build succeeded, but was flaky

• Buildkite Build
• Commit: 79918b1

Failed CI Steps

• Defend Workflows Cypress Tests on Serverless #6

Metrics [docs]

Module Count

Fewer modules leads to a faster build time

id	before	after	diff
securitySolution	4812	4816	+4

Public APIs missing comments

Total count of every public API that lacks a comment. Target amount is 0. Run node scripts/build_api_docs --plugin [yourplugin] --stats comments for more detailed information.

id	before	after	diff
@kbn/elastic-assistant	77	78	+1
@kbn/elastic-assistant-common	-	30	+30
total			+31

Async chunks

Total size of all lazy-loaded chunks that will be downloaded as the user navigates the app

id	before	after	diff
securitySolution	15.3MB	15.3MB	+20.0KB

Page load bundle

Size of the bundles that are downloaded on every page load. Target size is below 100kb

id	before	after	diff
securitySolution	66.7KB	66.7KB	+24.0B

Unknown metric groups

API count

id	before	after	diff
@kbn/elastic-assistant	97	98	+1
@kbn/elastic-assistant-common	-	30	+30
total			+31

ESLint disabled line counts

id	before	after	diff
@kbn/elastic-assistant-common	-	1	+1

Total ESLint disabled count

id	before	after	diff
@kbn/elastic-assistant-common	-	1	+1

History

• 💛 Build #182026 was flaky b7871169f57b669abfefe4024212a4ea07b2beef
• 💔 Build #181666 failed b8f0c4039f51e8cb80c07f1987c11ea97414b483

To update your PR or re-run it, just comment with:
@elasticmachine merge upstream

cc @andrew-goldstein