[Security Solution] [Attack discovery] Output chunking / refinement, LangGraph migration, and evaluation improvements

[andrew-goldstein]

[Security Solution] [Attack discovery] Output chunking / refinement, LangGraph migration, and evaluation improvements

Summary

This PR improves the Attack discovery user and developer experience with output chunking / refinement, migration to LangGraph, and improvements to evaluations.

The improvements were realized by transitioning from directly using lower-level LangChain apis to LangGraph in this PR, and a deeper integration with the evaluation features of LangSmith.

Output chunking

Output chunking increases the maximum and default number of alerts sent as context, working around the output token limitations of popular large language models (LLMs):

	Old	New
max alerts	100	500
default alerts	20	200

See Output chunking details below for more information.

Settings

A new settings modal makes it possible to configure the number of alerts sent as context directly from the Attack discovery page:

• Previously, users configured this value for Attack discovery via the security assistant Knowledge base settings, as documented here
• The new settings modal uses local storage (instead of the previously-shared assistant Knowledge base setting, which is stored in Elasticsearch)

Output refinement

Output refinement automatically combines related discoveries (that were previously represented as two or more discoveries):

• The refine step in the graph diagram above may (for example), combine three discoveries from the generate step into two discoveries when they are related

Hallucination detection

New hallucination detection displays an error in lieu of showing hallucinated output:

• A new tour step was added to the Attack discovery page to share the improvements:

Summary of improvements for developers

The following features improve the developer experience when running evaluations for Attack discovery:

Replay alerts in evaluations

This evaluation feature eliminates the need to populate a local environment with alerts to (re)run evaluations:

Alert replay skips the retrieve_anonymized_alerts step in the graph, because it uses the anonymizedAlerts and replacements provided as Input in a dataset example. See Replay alerts in evaluations details below for more information.

Override graph state

Override graph state via datatset examples to test prompt improvements and edge cases via evaluations:

To use this feature, add an overrides key to the Input of a dataset example. See Override graph state details below for more information.

New custom evaluator

Prior to this PR, an evaluator had to be manually added to each dataset in LangSmith to use an LLM as the judge for correctness.

This PR introduces a custom, programmatic evaluator that handles anonymization automatically, and eliminates the need to manually create evaluators in LangSmith. To use it, simply run evaluations from the Evaluation tab in settings.

New evaluation settings

This PR introduces new settings in the Evaluation tab:

New evaluation settings:

• Evaluator model (optional) - Judge the quality of predictions using a single model. (Default: use the same model as the connector)

This new setting is useful when you want to use the same model, e.g. GPT-4o to judge the quality of all the models evaluated in an experiment.

• Default max alerts - The default maximum number of alerts to send as context, which may be overridden by the example input

This new setting is useful when using the alerts in the local environment to run evaluations. Examples that use the Alerts replay feature will ignore this value, because the alerts in the example Input will be used instead.

Directory structure refactoring

• The server-side directory structure was refactored to consolidate the location of Attack discovery related files

Details

This section describes some of the improvements above in detail.

Output chunking details

The new output chunking feature increases the maximum and default number of alerts that may be sent as context. It achieves this improvement by working around output token limitations.

LLMs have different limits for the number of tokens accepted as input for requests, and the number of tokens available for output when generating responses.

Today, the output token limits of most popular models are significantly smaller than the input token limits.

For example, at the time of this writing, the Gemini 1.5 Pro model's limits are (source):

• Input token limit: 2,097,152
• Output token limit: 8,192

As a result of this relatively smaller output token limit, previous versions of Attack discovery would simply fail when an LLM ran out of output tokens when generating a response. This often happened "mid sentence", and resulted in errors or hallucinations being displayed to users.

The new output chunking feature detects incomplete responses from the LLM in the generate step of the Graph. When an incomplete response is detected, the generate step will run again with:

• The original prompt
• The Alerts provided as context
• The partially generated response
• Instructions to "continue where you left off"

The generate step in the graph will run until one of the following conditions is met:

• The incomplete response can be successfully parsed
• The maximum number of generation attempts (default: 10) is reached
• The maximum number of hallucinations detected (default: 5) is reached

Output refinement details

The new output refinement feature automatically combines related discoveries (that were previously represented as two or more discoveries).

The new refine step in the graph re-submits the discoveries from the generate step with a refinePrompt to combine related attack discoveries.

The refine step is subject to the model's output token limits, just like the generate step. That means a response to the refine prompt from the LLM may be cut off "mid" sentence. To that end:

• The refine step will re-run until the (same, shared) maxGenerationAttempts and maxHallucinationFailures limits as the generate step are reached
• The maximum number of attempts (default: 10) is shared with the generate step. For example, if it took 7 tries (generationAttempts) to complete the generate step, the refine step will only run up to 3 times.

The refine step will return unrefined results from the generate step when:

• The generate step uses all 10 generation attempts. When this happens, the refine step will be skipped, and the unrefined output of the generate step will be returned to the user
• If the refine step uses all remaining attempts, but fails to produce a refined response, due to output token limitations, or hallucinations in the refined response

Hallucination detection details

Before this PR, Attack discovery directly used lower level LangChain APIs to parse responses from the LLM. After this PR, Attack discovery uses LangGraph.

In the previous implementation, when Attack discovery received an incomplete response because the output token limits of a model were hit, the LangChain APIs automatically re-submitted the incomplete response in an attempt to "repair" it. However, the re-submitted results didn't include all of the original context (i.e. alerts that generated them). The repair process often resulted in hallucinated results being presented to users, especially with some models i.e. Claude 3.5 Haiku.

In this PR, the generate and refine steps detect (some) hallucinations. When hallucinations are detected:

• The current accumulated generations or refinements are (respectively) discarded, effectively restarting the generate or refine process
• The generate and refine steps will be retried until the maximum generation attempts (default: 10) or hallucinations detected (default: 5) limits are reached

Hitting the hallucination limit during the generate step will result in an error being displayed to the user.

Hitting the hallucination limit during the refine step will result in the unrefined discoveries being displayed to the user.

Replay alerts in evaluations details

Alerts replay makes it possible to re-run evaluations, even when your local deployment has zero alerts.

This feature eliminates the chore of populating your local instance with specific alerts for each example.

Every example in a dataset may (optionally) specify a different set of alerts.

Alert replay skips the retrieve_anonymized_alerts step in the graph, because it uses the anonymizedAlerts and replacements provided as Input in a dataset example.

The following instructions document the process of creating a new LangSmith dataset example that uses the Alerts replay feature:

1. In Kibana, navigate to Security > Attack discovery

2. Click Generate to generate Attack discoveries

3. In LangSmith, navigate to Projects > Your project

4. In the Runs tab of the LangSmith project, click on the latest Attack discovery entry to open the trace

5. IMPORTANT: In the trace, select the LAST ChannelWriteChannelWrite<attackDiscoveries,attackDisc... entry. The last entry will appear inside the LAST refine step in the trace, as illustrated by the screenshot below:

6. With the last ChannelWriteChannelWrite<attackDiscoveries,attackDisc... entry selected, click Add to > Add to Dataset

7. Copy-paste the Input to the Output, because evaluation Experiments always compare the current run with the Output in an example.

• This step is always required to create a dataset.
• If you don't want to use the Alert replay feature, replace Input with an empty object:

{}

8. Choose an existing dataset, or create a new one

9. Click the Submit button to add the example to the dataset.

After completing the steps above, the dataset is ready to be run in evaluations.

Override graph state details

When a dataset is run in an evaluation (to create Experiments):

• The (optional) anonymizedAlerts and replacements provided as Input in the example will be replayed, bypassing the retrieve_anonymized_alerts step in the graph
• The rest of the properties in Input will not be used as inputs to the graph
• In contrast, an empty object {} in Input means the latest and riskiest alerts in the last 24 hours in the local environment will be queried

In addition to the above, you may add an optional overrides key in the Input of a dataset example to test changes or edge cases. This is useful for evaluating changes without updating the code directly.

The overrides set the initial state of the graph before it's run in an evaluation.

The example Input below overrides the prompts used in the generate and refine steps:

{
  "overrides": {
    "refinePrompt": "This overrides the refine prompt",
    "attackDiscoveryPrompt": "This overrides the attack discovery prompt"
  }
}

To use the overrides feature in evaluations to set the initial state of the graph:

1. Create a dataset example, as documented in the Replay alerts in evaluations details section above

2. In LangSmith, navigate to Datasets & Testing > Your Dataset

3. In the dataset, click the Examples tab

4. Click an example to open it in the flyout

5. Click the Edit button to edit the example

6. Add the overrides key shown below to the Input e.g.:

{
  "overrides": {
    "refinePrompt": "This overrides the refine prompt",
    "attackDiscoveryPrompt": "This overrides the attack discovery prompt"
  }
}

7. Edit the overrides in the example Input above to add (or remove) entries that will determine the initial state of the graph.

All of the overides shown in step 6 are optional. The refinePrompt and attackDiscoveryPrompt could be removed from the overrides example above, and replaced with maxGenerationAttempts to test a higher limit.

All valid graph state may be specified in overrides.

[elasticmachine]

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

[stephmilovic]

nit: Evals is always developer only, right? Never user facing? We probably don't need to be using i18n resources

[andrew-goldstein]

Yes, the audience of evals is developers, who must enable this UI via the assistantModelEvaluation feature flag. The original author of this file also used i18n, so the new entries are consistent with that style.

[stephmilovic]

It's a nit so no big deal, but it might be nice to remove this pattern to avoid unnecessary resource use

[kibanamachine]

Starting backport for target branches: 8.x

https://github.com/elastic/kibana/actions/runs/11348436218

[kibanamachine]

💚 All backports created successfully

Status	Branch	Result
✅	8.x

Note: Successful backport PRs will be merged automatically after passing CI.

Questions ?

Please refer to the Backport tool documentation

[jbudz]

9.0 was reverted with dbe6d82.
https://buildkite.com/elastic/kibana-on-merge/builds/52349

I'm monitoring 8.x, will only revert if required.

[andrew-goldstein]

Before merging, I asked the ops team in this thread if this PR was safe to merge, because the build state was out of sync with the build comments.

The build comments being out of sync is a known build issue, so I got an explicit OK / safe ✅ to merge in that thread, but another PR that added a new lint rule merged first:
#195456

The new lint rule requires all imports from react-use to change from:

import { useLocalStorage } from 'react-use';

to:

import useLocalStorage from 'react-use/lib/useLocalStorage';

As a result, this PR started failing CI, but only after it merged to main.

[jbudz]

The revert has been re-applied in #196440 / ad2ac71. Dropping the revert label.

[andrew-goldstein]

Thank you @jbudz!