Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Security Solution] [Attack discovery] Output chunking / refinement, …
…LangGraph migration, and evaluation improvements (elastic#195669) ## [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](https://www.elastic.co/guide/en/security/8.15/attack-discovery.html#attack-discovery-generate-discoveries) - 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](https://ai.google.dev/gemini-api/docs/models/gemini)): - 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: ```json {} ``` 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: ```json { "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.: ```json { "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`.
- Loading branch information
1 parent
204f9d3
commit 2c21adb