[8.x] [Security Solution] [Attack discovery] Output chunking / refinement, LangGraph migration, and evaluation improvements (#195669) (#196334)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[Security Solution] [Attack discovery] Output chunking / refinement,
LangGraph migration, and evaluation improvements
(#195669)](https://github.com/elastic/kibana/pull/195669)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Andrew
Macri","email":"andrew.macri@elastic.co"},"sourceCommit":{"committedDate":"2024-10-15T14:39:48Z","message":"[Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation improvements (#195669)\n\n## [Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation improvements\r\n\r\n### Summary\r\n\r\nThis PR
improves the Attack discovery user and developer experience with output
chunking / refinement, migration to LangGraph, and improvements to
evaluations.\r\n\r\nThe 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.\r\n\r\n#### Output chunking\r\n\r\n_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):\r\n\r\n| | Old | New
|\r\n|----------------|-------|-------|\r\n| max alerts | `100` | `500`
|\r\n| default alerts | `20` | `200` |\r\n\r\nSee _Output chunking
details_ below for more information.\r\n\r\n#### Settings\r\n\r\nA new
settings modal makes it possible to configure the number of alerts sent
as context directly from the Attack discovery
page:\r\n\r\n![settings](https://github.com/user-attachments/assets/3f5ab4e9-5eae-4f99-8490-e392c758fa6e)\r\n\r\n-
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)\r\n-
The new settings modal uses local storage (instead of the
previously-shared assistant Knowledge base setting, which is stored in
Elasticsearch)\r\n\r\n#### Output refinement\r\n\r\n_Output refinement_
automatically combines related discoveries (that were previously
represented as two or more discoveries):\r\n\r\n
![default_attack_discovery_graph](https://github.com/user-attachments/assets/c092bb42-a41e-4fba-85c2-a4b2c1ef3053)\r\n\r\n-
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\r\n\r\n### Hallucination detection\r\n\r\nNew
_hallucination detection_ displays an error in lieu of showing
hallucinated
output:\r\n\r\n![hallucination_detection](https://github.com/user-attachments/assets/1d849908-3f10-4fe8-8741-c0cf418b1524)\r\n\r\n-
A new tour step was added to the Attack discovery page to share the
improvements:\r\n\r\n![tour_step](https://github.com/user-attachments/assets/0cedf770-baba-41b1-8ec6-b12b14c0c57a)\r\n\r\n###
Summary of improvements for developers\r\n\r\nThe following features
improve the developer experience when running evaluations for Attack
discovery:\r\n\r\n#### Replay alerts in evaluations\r\n\r\nThis
evaluation feature eliminates the need to populate a local environment
with alerts to (re)run evaluations:\r\n\r\n
![alerts_as_input](https://github.com/user-attachments/assets/b29dc847-3d53-4b17-8757-ed59852c1623)\r\n\r\nAlert
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.\r\n\r\n#### Override graph state\r\n\r\nOverride graph
state via datatset examples to test prompt improvements and edge cases
via evaluations:\r\n\r\n
![override_graph_input](https://github.com/user-attachments/assets/a685177b-1e07-4f49-9b8d-c0b652975237)\r\n\r\nTo
use this feature, add an `overrides` key to the `Input` of a dataset
example. See _Override graph state details_ below for more
information.\r\n\r\n#### New custom evaluator\r\n\r\nPrior to this PR,
an evaluator had to be manually added to each dataset in LangSmith to
use an LLM as the judge for correctness.\r\n\r\nThis 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.\r\n\r\n#### New evaluation settings\r\n\r\nThis PR introduces
new settings in the `Evaluation`
tab:\r\n\r\n![new_evaluation_settings](https://github.com/user-attachments/assets/ca72aa2a-b0dc-4bec-9409-386d77d6a2f4)\r\n\r\nNew
evaluation settings:\r\n\r\n- `Evaluator model (optional)` - Judge the
quality of predictions using a single model. (Default: use the same
model as the connector)\r\n\r\nThis 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.\r\n\r\n- `Default max alerts` - The
default maximum number of alerts to send as context, which may be
overridden by the example input\r\n\r\nThis 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.\r\n\r\n####
Directory structure refactoring\r\n\r\n- The server-side directory
structure was refactored to consolidate the location of Attack discovery
related files\r\n\r\n### Details\r\n\r\nThis section describes some of
the improvements above in detail.\r\n\r\n#### Output chunking
details\r\n\r\nThe 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.\r\n\r\nLLMs have
different limits for the number of tokens accepted as _input_ for
requests, and the number of tokens available for _output_ when
generating responses.\r\n\r\nToday, the output token limits of most
popular models are significantly smaller than the input token
limits.\r\n\r\nFor 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)):\r\n\r\n-
Input token limit: `2,097,152`\r\n- Output token limit:
`8,192`\r\n\r\nAs 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.\r\n\r\nThe 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:\r\n\r\n- The original prompt\r\n- The Alerts provided as
context\r\n- The partially generated response\r\n- Instructions to
\"continue where you left off\"\r\n\r\nThe `generate` step in the graph
will run until one of the following conditions is met:\r\n\r\n- The
incomplete response can be successfully parsed\r\n- The maximum number
of generation attempts (default: `10`) is reached\r\n- The maximum
number of hallucinations detected (default: `5`) is reached\r\n\r\n####
Output refinement details\r\n\r\nThe new output refinement feature
automatically combines related discoveries (that were previously
represented as two or more discoveries).\r\n\r\nThe new `refine` step in
the graph re-submits the discoveries from the `generate` step with a
`refinePrompt` to combine related attack discoveries.\r\n\r\nThe
`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:\r\n\r\n- The refine
step will re-run until the (same, shared) `maxGenerationAttempts` and
`maxHallucinationFailures` limits as the `generate` step are
reached\r\n- 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.\r\n\r\nThe `refine` step will
return _unrefined_ results from the `generate` step when:\r\n\r\n- 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\r\n- 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\r\n\r\n#### Hallucination detection details\r\n\r\nBefore this
PR, Attack discovery directly used lower level LangChain APIs to parse
responses from the LLM. After this PR, Attack discovery uses
LangGraph.\r\n\r\nIn 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`.\r\n\r\nIn this PR, the `generate` and
`refine` steps detect (some) hallucinations. When hallucinations are
detected:\r\n\r\n- The current accumulated `generations` or
`refinements` are (respectively) discarded, effectively restarting the
`generate` or `refine` process\r\n- The `generate` and `refine` steps
will be retried until the maximum generation attempts (default: `10`) or
hallucinations detected (default: `5`) limits are reached\r\n\r\nHitting
the hallucination limit during the `generate` step will result in an
error being displayed to the user.\r\n\r\nHitting the hallucination
limit during the `refine` step will result in the unrefined discoveries
being displayed to the user.\r\n\r\n#### Replay alerts in evaluations
details\r\n\r\nAlerts replay makes it possible to re-run evaluations,
even when your local deployment has zero alerts.\r\n\r\nThis feature
eliminates the chore of populating your local instance with specific
alerts for each example.\r\n\r\nEvery example in a dataset may
(optionally) specify a different set of alerts.\r\n\r\nAlert replay
skips the `retrieve_anonymized_alerts` step in the graph, because it
uses the `anonymizedAlerts` and `replacements` provided as `Input` in a
dataset example.\r\n\r\nThe following instructions document the process
of creating a new LangSmith dataset example that uses the Alerts replay
feature:\r\n\r\n1) In Kibana, navigate to Security > Attack
discovery\r\n\r\n2) Click `Generate` to generate Attack
discoveries\r\n\r\n3) In LangSmith, navigate to Projects > _Your
project_\r\n\r\n4) In the `Runs` tab of the LangSmith project, click on
the latest `Attack discovery` entry to open the trace\r\n\r\n5)
**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:\r\n\r\n![last_channel_write](https://github.com/user-attachments/assets/c57fc803-3bbb-4603-b99f-d2b130428201)\r\n\r\n6)
With the last `ChannelWriteChannelWrite<attackDiscoveries,attackDisc...`
entry selected, click `Add to` > `Add to Dataset`\r\n\r\n7) Copy-paste
the `Input` to the `Output`, because evaluation Experiments always
compare the current run with the `Output` in an example.\r\n\r\n- This
step is _always_ required to create a dataset.\r\n- If you don't want to
use the Alert replay feature, replace `Input` with an empty
object:\r\n\r\n```json\r\n{}\r\n```\r\n\r\n8) Choose an existing
dataset, or create a new one\r\n\r\n9) Click the `Submit` button to add
the example to the dataset.\r\n\r\nAfter completing the steps above, the
dataset is ready to be run in evaluations.\r\n\r\n#### Override graph
state details\r\n\r\nWhen a dataset is run in an evaluation (to create
Experiments):\r\n\r\n- The (optional) `anonymizedAlerts` and
`replacements` provided as `Input` in the example will be replayed,
bypassing the `retrieve_anonymized_alerts` step in the graph\r\n- The
rest of the properties in `Input` will not be used as inputs to the
graph\r\n- 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\r\n\r\nIn 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.\r\n\r\nThe `overrides` set the initial state of the graph
before it's run in an evaluation.\r\n\r\nThe example `Input` below
overrides the prompts used in the `generate` and `refine`
steps:\r\n\r\n```json\r\n{\r\n \"overrides\": {\r\n \"refinePrompt\":
\"This overrides the refine prompt\",\r\n \"attackDiscoveryPrompt\":
\"This overrides the attack discovery prompt\"\r\n
}\r\n}\r\n```\r\n\r\nTo use the `overrides` feature in evaluations to
set the initial state of the graph:\r\n\r\n1) Create a dataset example,
as documented in the _Replay alerts in evaluations details_ section
above\r\n\r\n2) In LangSmith, navigate to Datasets & Testing > _Your
Dataset_\r\n\r\n3) In the dataset, click the Examples tab\r\n\r\n4)
Click an example to open it in the flyout\r\n\r\n5) Click the `Edit`
button to edit the example\r\n\r\n6) Add the `overrides` key shown below
to the `Input` e.g.:\r\n\r\n```json\r\n{\r\n \"overrides\": {\r\n
\"refinePrompt\": \"This overrides the refine prompt\",\r\n
\"attackDiscoveryPrompt\": \"This overrides the attack discovery
prompt\"\r\n }\r\n}\r\n```\r\n\r\n7) Edit the `overrides` in the example
`Input` above to add (or remove) entries that will determine the initial
state of the graph.\r\n\r\nAll 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.\r\n\r\nAll valid graph
state may be specified in
`overrides`.","sha":"2c21adb8faafc0016ad7a6591837118f6bdf0907","branchLabelMapping":{"^v9.0.0$":"main","^v8.16.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:enhancement","v9.0.0","Team:
SecuritySolution","ci:cloud-deploy","ci:cloud-persist-deployment","Team:Security
Generative AI","v8.16.0","backport:version"],"title":"[Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation
improvements","number":195669,"url":"https://github.com/elastic/kibana/pull/195669","mergeCommit":{"message":"[Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation improvements (#195669)\n\n## [Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation improvements\r\n\r\n### Summary\r\n\r\nThis PR
improves the Attack discovery user and developer experience with output
chunking / refinement, migration to LangGraph, and improvements to
evaluations.\r\n\r\nThe 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.\r\n\r\n#### Output chunking\r\n\r\n_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):\r\n\r\n| | Old | New
|\r\n|----------------|-------|-------|\r\n| max alerts | `100` | `500`
|\r\n| default alerts | `20` | `200` |\r\n\r\nSee _Output chunking
details_ below for more information.\r\n\r\n#### Settings\r\n\r\nA new
settings modal makes it possible to configure the number of alerts sent
as context directly from the Attack discovery
page:\r\n\r\n![settings](https://github.com/user-attachments/assets/3f5ab4e9-5eae-4f99-8490-e392c758fa6e)\r\n\r\n-
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)\r\n-
The new settings modal uses local storage (instead of the
previously-shared assistant Knowledge base setting, which is stored in
Elasticsearch)\r\n\r\n#### Output refinement\r\n\r\n_Output refinement_
automatically combines related discoveries (that were previously
represented as two or more discoveries):\r\n\r\n
![default_attack_discovery_graph](https://github.com/user-attachments/assets/c092bb42-a41e-4fba-85c2-a4b2c1ef3053)\r\n\r\n-
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\r\n\r\n### Hallucination detection\r\n\r\nNew
_hallucination detection_ displays an error in lieu of showing
hallucinated
output:\r\n\r\n![hallucination_detection](https://github.com/user-attachments/assets/1d849908-3f10-4fe8-8741-c0cf418b1524)\r\n\r\n-
A new tour step was added to the Attack discovery page to share the
improvements:\r\n\r\n![tour_step](https://github.com/user-attachments/assets/0cedf770-baba-41b1-8ec6-b12b14c0c57a)\r\n\r\n###
Summary of improvements for developers\r\n\r\nThe following features
improve the developer experience when running evaluations for Attack
discovery:\r\n\r\n#### Replay alerts in evaluations\r\n\r\nThis
evaluation feature eliminates the need to populate a local environment
with alerts to (re)run evaluations:\r\n\r\n
![alerts_as_input](https://github.com/user-attachments/assets/b29dc847-3d53-4b17-8757-ed59852c1623)\r\n\r\nAlert
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.\r\n\r\n#### Override graph state\r\n\r\nOverride graph
state via datatset examples to test prompt improvements and edge cases
via evaluations:\r\n\r\n
![override_graph_input](https://github.com/user-attachments/assets/a685177b-1e07-4f49-9b8d-c0b652975237)\r\n\r\nTo
use this feature, add an `overrides` key to the `Input` of a dataset
example. See _Override graph state details_ below for more
information.\r\n\r\n#### New custom evaluator\r\n\r\nPrior to this PR,
an evaluator had to be manually added to each dataset in LangSmith to
use an LLM as the judge for correctness.\r\n\r\nThis 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.\r\n\r\n#### New evaluation settings\r\n\r\nThis PR introduces
new settings in the `Evaluation`
tab:\r\n\r\n![new_evaluation_settings](https://github.com/user-attachments/assets/ca72aa2a-b0dc-4bec-9409-386d77d6a2f4)\r\n\r\nNew
evaluation settings:\r\n\r\n- `Evaluator model (optional)` - Judge the
quality of predictions using a single model. (Default: use the same
model as the connector)\r\n\r\nThis 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.\r\n\r\n- `Default max alerts` - The
default maximum number of alerts to send as context, which may be
overridden by the example input\r\n\r\nThis 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.\r\n\r\n####
Directory structure refactoring\r\n\r\n- The server-side directory
structure was refactored to consolidate the location of Attack discovery
related files\r\n\r\n### Details\r\n\r\nThis section describes some of
the improvements above in detail.\r\n\r\n#### Output chunking
details\r\n\r\nThe 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.\r\n\r\nLLMs have
different limits for the number of tokens accepted as _input_ for
requests, and the number of tokens available for _output_ when
generating responses.\r\n\r\nToday, the output token limits of most
popular models are significantly smaller than the input token
limits.\r\n\r\nFor 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)):\r\n\r\n-
Input token limit: `2,097,152`\r\n- Output token limit:
`8,192`\r\n\r\nAs 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.\r\n\r\nThe 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:\r\n\r\n- The original prompt\r\n- The Alerts provided as
context\r\n- The partially generated response\r\n- Instructions to
\"continue where you left off\"\r\n\r\nThe `generate` step in the graph
will run until one of the following conditions is met:\r\n\r\n- The
incomplete response can be successfully parsed\r\n- The maximum number
of generation attempts (default: `10`) is reached\r\n- The maximum
number of hallucinations detected (default: `5`) is reached\r\n\r\n####
Output refinement details\r\n\r\nThe new output refinement feature
automatically combines related discoveries (that were previously
represented as two or more discoveries).\r\n\r\nThe new `refine` step in
the graph re-submits the discoveries from the `generate` step with a
`refinePrompt` to combine related attack discoveries.\r\n\r\nThe
`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:\r\n\r\n- The refine
step will re-run until the (same, shared) `maxGenerationAttempts` and
`maxHallucinationFailures` limits as the `generate` step are
reached\r\n- 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.\r\n\r\nThe `refine` step will
return _unrefined_ results from the `generate` step when:\r\n\r\n- 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\r\n- 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\r\n\r\n#### Hallucination detection details\r\n\r\nBefore this
PR, Attack discovery directly used lower level LangChain APIs to parse
responses from the LLM. After this PR, Attack discovery uses
LangGraph.\r\n\r\nIn 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`.\r\n\r\nIn this PR, the `generate` and
`refine` steps detect (some) hallucinations. When hallucinations are
detected:\r\n\r\n- The current accumulated `generations` or
`refinements` are (respectively) discarded, effectively restarting the
`generate` or `refine` process\r\n- The `generate` and `refine` steps
will be retried until the maximum generation attempts (default: `10`) or
hallucinations detected (default: `5`) limits are reached\r\n\r\nHitting
the hallucination limit during the `generate` step will result in an
error being displayed to the user.\r\n\r\nHitting the hallucination
limit during the `refine` step will result in the unrefined discoveries
being displayed to the user.\r\n\r\n#### Replay alerts in evaluations
details\r\n\r\nAlerts replay makes it possible to re-run evaluations,
even when your local deployment has zero alerts.\r\n\r\nThis feature
eliminates the chore of populating your local instance with specific
alerts for each example.\r\n\r\nEvery example in a dataset may
(optionally) specify a different set of alerts.\r\n\r\nAlert replay
skips the `retrieve_anonymized_alerts` step in the graph, because it
uses the `anonymizedAlerts` and `replacements` provided as `Input` in a
dataset example.\r\n\r\nThe following instructions document the process
of creating a new LangSmith dataset example that uses the Alerts replay
feature:\r\n\r\n1) In Kibana, navigate to Security > Attack
discovery\r\n\r\n2) Click `Generate` to generate Attack
discoveries\r\n\r\n3) In LangSmith, navigate to Projects > _Your
project_\r\n\r\n4) In the `Runs` tab of the LangSmith project, click on
the latest `Attack discovery` entry to open the trace\r\n\r\n5)
**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:\r\n\r\n![last_channel_write](https://github.com/user-attachments/assets/c57fc803-3bbb-4603-b99f-d2b130428201)\r\n\r\n6)
With the last `ChannelWriteChannelWrite<attackDiscoveries,attackDisc...`
entry selected, click `Add to` > `Add to Dataset`\r\n\r\n7) Copy-paste
the `Input` to the `Output`, because evaluation Experiments always
compare the current run with the `Output` in an example.\r\n\r\n- This
step is _always_ required to create a dataset.\r\n- If you don't want to
use the Alert replay feature, replace `Input` with an empty
object:\r\n\r\n```json\r\n{}\r\n```\r\n\r\n8) Choose an existing
dataset, or create a new one\r\n\r\n9) Click the `Submit` button to add
the example to the dataset.\r\n\r\nAfter completing the steps above, the
dataset is ready to be run in evaluations.\r\n\r\n#### Override graph
state details\r\n\r\nWhen a dataset is run in an evaluation (to create
Experiments):\r\n\r\n- The (optional) `anonymizedAlerts` and
`replacements` provided as `Input` in the example will be replayed,
bypassing the `retrieve_anonymized_alerts` step in the graph\r\n- The
rest of the properties in `Input` will not be used as inputs to the
graph\r\n- 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\r\n\r\nIn 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.\r\n\r\nThe `overrides` set the initial state of the graph
before it's run in an evaluation.\r\n\r\nThe example `Input` below
overrides the prompts used in the `generate` and `refine`
steps:\r\n\r\n```json\r\n{\r\n \"overrides\": {\r\n \"refinePrompt\":
\"This overrides the refine prompt\",\r\n \"attackDiscoveryPrompt\":
\"This overrides the attack discovery prompt\"\r\n
}\r\n}\r\n```\r\n\r\nTo use the `overrides` feature in evaluations to
set the initial state of the graph:\r\n\r\n1) Create a dataset example,
as documented in the _Replay alerts in evaluations details_ section
above\r\n\r\n2) In LangSmith, navigate to Datasets & Testing > _Your
Dataset_\r\n\r\n3) In the dataset, click the Examples tab\r\n\r\n4)
Click an example to open it in the flyout\r\n\r\n5) Click the `Edit`
button to edit the example\r\n\r\n6) Add the `overrides` key shown below
to the `Input` e.g.:\r\n\r\n```json\r\n{\r\n \"overrides\": {\r\n
\"refinePrompt\": \"This overrides the refine prompt\",\r\n
\"attackDiscoveryPrompt\": \"This overrides the attack discovery
prompt\"\r\n }\r\n}\r\n```\r\n\r\n7) Edit the `overrides` in the example
`Input` above to add (or remove) entries that will determine the initial
state of the graph.\r\n\r\nAll 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.\r\n\r\nAll valid graph
state may be specified in
`overrides`.","sha":"2c21adb8faafc0016ad7a6591837118f6bdf0907"}},"sourceBranch":"main","suggestedTargetBranches":["8.x"],"targetPullRequestStates":[{"branch":"main","label":"v9.0.0","branchLabelMappingKey":"^v9.0.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/195669","number":195669,"mergeCommit":{"message":"[Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation improvements (#195669)\n\n## [Security
Solution] [Attack discovery] Output chunking / refinement, LangGraph
migration, and evaluation improvements\r\n\r\n### Summary\r\n\r\nThis PR
improves the Attack discovery user and developer experience with output
chunking / refinement, migration to LangGraph, and improvements to
evaluations.\r\n\r\nThe 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.\r\n\r\n#### Output chunking\r\n\r\n_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):\r\n\r\n| | Old | New
|\r\n|----------------|-------|-------|\r\n| max alerts | `100` | `500`
|\r\n| default alerts | `20` | `200` |\r\n\r\nSee _Output chunking
details_ below for more information.\r\n\r\n#### Settings\r\n\r\nA new
settings modal makes it possible to configure the number of alerts sent
as context directly from the Attack discovery
page:\r\n\r\n![settings](https://github.com/user-attachments/assets/3f5ab4e9-5eae-4f99-8490-e392c758fa6e)\r\n\r\n-
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)\r\n-
The new settings modal uses local storage (instead of the
previously-shared assistant Knowledge base setting, which is stored in
Elasticsearch)\r\n\r\n#### Output refinement\r\n\r\n_Output refinement_
automatically combines related discoveries (that were previously
represented as two or more discoveries):\r\n\r\n
![default_attack_discovery_graph](https://github.com/user-attachments/assets/c092bb42-a41e-4fba-85c2-a4b2c1ef3053)\r\n\r\n-
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\r\n\r\n### Hallucination detection\r\n\r\nNew
_hallucination detection_ displays an error in lieu of showing
hallucinated
output:\r\n\r\n![hallucination_detection](https://github.com/user-attachments/assets/1d849908-3f10-4fe8-8741-c0cf418b1524)\r\n\r\n-
A new tour step was added to the Attack discovery page to share the
improvements:\r\n\r\n![tour_step](https://github.com/user-attachments/assets/0cedf770-baba-41b1-8ec6-b12b14c0c57a)\r\n\r\n###
Summary of improvements for developers\r\n\r\nThe following features
improve the developer experience when running evaluations for Attack
discovery:\r\n\r\n#### Replay alerts in evaluations\r\n\r\nThis
evaluation feature eliminates the need to populate a local environment
with alerts to (re)run evaluations:\r\n\r\n
![alerts_as_input](https://github.com/user-attachments/assets/b29dc847-3d53-4b17-8757-ed59852c1623)\r\n\r\nAlert
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.\r\n\r\n#### Override graph state\r\n\r\nOverride graph
state via datatset examples to test prompt improvements and edge cases
via evaluations:\r\n\r\n
![override_graph_input](https://github.com/user-attachments/assets/a685177b-1e07-4f49-9b8d-c0b652975237)\r\n\r\nTo
use this feature, add an `overrides` key to the `Input` of a dataset
example. See _Override graph state details_ below for more
information.\r\n\r\n#### New custom evaluator\r\n\r\nPrior to this PR,
an evaluator had to be manually added to each dataset in LangSmith to
use an LLM as the judge for correctness.\r\n\r\nThis 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.\r\n\r\n#### New evaluation settings\r\n\r\nThis PR introduces
new settings in the `Evaluation`
tab:\r\n\r\n![new_evaluation_settings](https://github.com/user-attachments/assets/ca72aa2a-b0dc-4bec-9409-386d77d6a2f4)\r\n\r\nNew
evaluation settings:\r\n\r\n- `Evaluator model (optional)` - Judge the
quality of predictions using a single model. (Default: use the same
model as the connector)\r\n\r\nThis 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.\r\n\r\n- `Default max alerts` - The
default maximum number of alerts to send as context, which may be
overridden by the example input\r\n\r\nThis 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.\r\n\r\n####
Directory structure refactoring\r\n\r\n- The server-side directory
structure was refactored to consolidate the location of Attack discovery
related files\r\n\r\n### Details\r\n\r\nThis section describes some of
the improvements above in detail.\r\n\r\n#### Output chunking
details\r\n\r\nThe 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.\r\n\r\nLLMs have
different limits for the number of tokens accepted as _input_ for
requests, and the number of tokens available for _output_ when
generating responses.\r\n\r\nToday, the output token limits of most
popular models are significantly smaller than the input token
limits.\r\n\r\nFor 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)):\r\n\r\n-
Input token limit: `2,097,152`\r\n- Output token limit:
`8,192`\r\n\r\nAs 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.\r\n\r\nThe 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:\r\n\r\n- The original prompt\r\n- The Alerts provided as
context\r\n- The partially generated response\r\n- Instructions to
\"continue where you left off\"\r\n\r\nThe `generate` step in the graph
will run until one of the following conditions is met:\r\n\r\n- The
incomplete response can be successfully parsed\r\n- The maximum number
of generation attempts (default: `10`) is reached\r\n- The maximum
number of hallucinations detected (default: `5`) is reached\r\n\r\n####
Output refinement details\r\n\r\nThe new output refinement feature
automatically combines related discoveries (that were previously
represented as two or more discoveries).\r\n\r\nThe new `refine` step in
the graph re-submits the discoveries from the `generate` step with a
`refinePrompt` to combine related attack discoveries.\r\n\r\nThe
`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:\r\n\r\n- The refine
step will re-run until the (same, shared) `maxGenerationAttempts` and
`maxHallucinationFailures` limits as the `generate` step are
reached\r\n- 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.\r\n\r\nThe `refine` step will
return _unrefined_ results from the `generate` step when:\r\n\r\n- 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\r\n- 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\r\n\r\n#### Hallucination detection details\r\n\r\nBefore this
PR, Attack discovery directly used lower level LangChain APIs to parse
responses from the LLM. After this PR, Attack discovery uses
LangGraph.\r\n\r\nIn 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`.\r\n\r\nIn this PR, the `generate` and
`refine` steps detect (some) hallucinations. When hallucinations are
detected:\r\n\r\n- The current accumulated `generations` or
`refinements` are (respectively) discarded, effectively restarting the
`generate` or `refine` process\r\n- The `generate` and `refine` steps
will be retried until the maximum generation attempts (default: `10`) or
hallucinations detected (default: `5`) limits are reached\r\n\r\nHitting
the hallucination limit during the `generate` step will result in an
error being displayed to the user.\r\n\r\nHitting the hallucination
limit during the `refine` step will result in the unrefined discoveries
being displayed to the user.\r\n\r\n#### Replay alerts in evaluations
details\r\n\r\nAlerts replay makes it possible to re-run evaluations,
even when your local deployment has zero alerts.\r\n\r\nThis feature
eliminates the chore of populating your local instance with specific
alerts for each example.\r\n\r\nEvery example in a dataset may
(optionally) specify a different set of alerts.\r\n\r\nAlert replay
skips the `retrieve_anonymized_alerts` step in the graph, because it
uses the `anonymizedAlerts` and `replacements` provided as `Input` in a
dataset example.\r\n\r\nThe following instructions document the process
of creating a new LangSmith dataset example that uses the Alerts replay
feature:\r\n\r\n1) In Kibana, navigate to Security > Attack
discovery\r\n\r\n2) Click `Generate` to generate Attack
discoveries\r\n\r\n3) In LangSmith, navigate to Projects > _Your
project_\r\n\r\n4) In the `Runs` tab of the LangSmith project, click on
the latest `Attack discovery` entry to open the trace\r\n\r\n5)
**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:\r\n\r\n![last_channel_write](https://github.com/user-attachments/assets/c57fc803-3bbb-4603-b99f-d2b130428201)\r\n\r\n6)
With the last `ChannelWriteChannelWrite<attackDiscoveries,attackDisc...`
entry selected, click `Add to` > `Add to Dataset`\r\n\r\n7) Copy-paste
the `Input` to the `Output`, because evaluation Experiments always
compare the current run with the `Output` in an example.\r\n\r\n- This
step is _always_ required to create a dataset.\r\n- If you don't want to
use the Alert replay feature, replace `Input` with an empty
object:\r\n\r\n```json\r\n{}\r\n```\r\n\r\n8) Choose an existing
dataset, or create a new one\r\n\r\n9) Click the `Submit` button to add
the example to the dataset.\r\n\r\nAfter completing the steps above, the
dataset is ready to be run in evaluations.\r\n\r\n#### Override graph
state details\r\n\r\nWhen a dataset is run in an evaluation (to create
Experiments):\r\n\r\n- The (optional) `anonymizedAlerts` and
`replacements` provided as `Input` in the example will be replayed,
bypassing the `retrieve_anonymized_alerts` step in the graph\r\n- The
rest of the properties in `Input` will not be used as inputs to the
graph\r\n- 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\r\n\r\nIn 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.\r\n\r\nThe `overrides` set the initial state of the graph
before it's run in an evaluation.\r\n\r\nThe example `Input` below
overrides the prompts used in the `generate` and `refine`
steps:\r\n\r\n```json\r\n{\r\n \"overrides\": {\r\n \"refinePrompt\":
\"This overrides the refine prompt\",\r\n \"attackDiscoveryPrompt\":
\"This overrides the attack discovery prompt\"\r\n
}\r\n}\r\n```\r\n\r\nTo use the `overrides` feature in evaluations to
set the initial state of the graph:\r\n\r\n1) Create a dataset example,
as documented in the _Replay alerts in evaluations details_ section
above\r\n\r\n2) In LangSmith, navigate to Datasets & Testing > _Your
Dataset_\r\n\r\n3) In the dataset, click the Examples tab\r\n\r\n4)
Click an example to open it in the flyout\r\n\r\n5) Click the `Edit`
button to edit the example\r\n\r\n6) Add the `overrides` key shown below
to the `Input` e.g.:\r\n\r\n```json\r\n{\r\n \"overrides\": {\r\n
\"refinePrompt\": \"This overrides the refine prompt\",\r\n
\"attackDiscoveryPrompt\": \"This overrides the attack discovery
prompt\"\r\n }\r\n}\r\n```\r\n\r\n7) Edit the `overrides` in the example
`Input` above to add (or remove) entries that will determine the initial
state of the graph.\r\n\r\nAll 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.\r\n\r\nAll valid graph
state may be specified in
`overrides`.","sha":"2c21adb8faafc0016ad7a6591837118f6bdf0907"}},{"branch":"8.x","label":"v8.16.0","branchLabelMappingKey":"^v8.16.0$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Andrew Macri <andrew.macri@elastic.co>

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/get_open_and_acknowledged_alerts_query/index.test.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query';
+import { getOpenAndAcknowledgedAlertsQuery } from '.';
 
 describe('getOpenAndAcknowledgedAlertsQuery', () => {
   it('returns the expected query', () => {

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/get_open_and_acknowledged_alerts_query/index.ts

@@ -5,8 +5,13 @@
  * 2.0.
  */
 
-import type { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
+import type { AnonymizationFieldResponse } from '../../schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
 
+/**
+ * This query returns open and acknowledged (non-building block) alerts in the last 24 hours.
+ *
+ * The alerts are ordered by risk score, and then from the most recent to the oldest.
+ */
 export const getOpenAndAcknowledgedAlertsQuery = ({
   alertsIndexPattern,
   anonymizationFields,

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.test.ts

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getRawDataOrDefault } from '.';
+
+describe('getRawDataOrDefault', () => {
+  it('returns the raw data when it is valid', () => {
+    const rawData = {
+      field1: [1, 2, 3],
+      field2: ['a', 'b', 'c'],
+    };
+
+    expect(getRawDataOrDefault(rawData)).toEqual(rawData);
+  });
+
+  it('returns an empty object when the raw data is invalid', () => {
+    const rawData = {
+      field1: [1, 2, 3],
+      field2: 'invalid',
+    };
+
+    expect(getRawDataOrDefault(rawData)).toEqual({});
+  });
+});

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { isRawDataValid } from '../is_raw_data_valid';
+import type { MaybeRawData } from '../types';
+
+/** Returns the raw data if it valid, or a default if it's not */
+export const getRawDataOrDefault = (rawData: MaybeRawData): Record<string, unknown[]> =>
+  isRawDataValid(rawData) ? rawData : {};

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/is_raw_data_valid/index.test.ts

@@ -0,0 +1,51 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { isRawDataValid } from '.';
+
+describe('isRawDataValid', () => {
+  it('returns true for valid raw data', () => {
+    const rawData = {
+      field1: [1, 2, 3], // the Fields API may return a number array
+      field2: ['a', 'b', 'c'], // the Fields API may return a string array
+    };
+
+    expect(isRawDataValid(rawData)).toBe(true);
+  });
+
+  it('returns true when a field array is empty', () => {
+    const rawData = {
+      field1: [1, 2, 3], // the Fields API may return a number array
+      field2: ['a', 'b', 'c'], // the Fields API may return a string array
+      field3: [], // the Fields API may return an empty array
+    };
+
+    expect(isRawDataValid(rawData)).toBe(true);
+  });
+
+  it('returns false when a field does not have an array of values', () => {
+    const rawData = {
+      field1: [1, 2, 3],
+      field2: 'invalid',
+    };
+
+    expect(isRawDataValid(rawData)).toBe(false);
+  });
+
+  it('returns true for empty raw data', () => {
+    const rawData = {};
+
+    expect(isRawDataValid(rawData)).toBe(true);
+  });
+
+  it('returns false when raw data is an unexpected type', () => {
+    const rawData = 1234;
+
+    // @ts-expect-error
+    expect(isRawDataValid(rawData)).toBe(false);
+  });
+});

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/is_raw_data_valid/index.ts

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MaybeRawData } from '../types';
+
+export const isRawDataValid = (rawData: MaybeRawData): rawData is Record<string, unknown[]> =>
+  typeof rawData === 'object' && Object.keys(rawData).every((x) => Array.isArray(rawData[x]));

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/size_is_out_of_range/index.test.ts

@@ -0,0 +1,47 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { sizeIsOutOfRange } from '.';
+import { MAX_SIZE, MIN_SIZE } from '../types';
+
+describe('sizeIsOutOfRange', () => {
+  it('returns true when size is undefined', () => {
+    const size = undefined;
+
+    expect(sizeIsOutOfRange(size)).toBe(true);
+  });
+
+  it('returns true when size is less than MIN_SIZE', () => {
+    const size = MIN_SIZE - 1;
+
+    expect(sizeIsOutOfRange(size)).toBe(true);
+  });
+
+  it('returns true when size is greater than MAX_SIZE', () => {
+    const size = MAX_SIZE + 1;
+
+    expect(sizeIsOutOfRange(size)).toBe(true);
+  });
+
+  it('returns false when size is exactly MIN_SIZE', () => {
+    const size = MIN_SIZE;
+
+    expect(sizeIsOutOfRange(size)).toBe(false);
+  });
+
+  it('returns false when size is exactly MAX_SIZE', () => {
+    const size = MAX_SIZE;
+
+    expect(sizeIsOutOfRange(size)).toBe(false);
+  });
+
+  it('returns false when size is within the valid range', () => {
+    const size = MIN_SIZE + 1;
+
+    expect(sizeIsOutOfRange(size)).toBe(false);
+  });
+});

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/size_is_out_of_range/index.ts

@@ -0,0 +1,12 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MAX_SIZE, MIN_SIZE } from '../types';
+
+/** Return true if the provided size is out of range */
+export const sizeIsOutOfRange = (size?: number): boolean =>
+  size == null || size < MIN_SIZE || size > MAX_SIZE;

x-pack/packages/kbn-elastic-assistant-common/impl/alerts/helpers/types.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types';
+
+export const MIN_SIZE = 10;
+export const MAX_SIZE = 10000;
+
+/** currently the same shape as "fields" property in the ES response */
+export type MaybeRawData = SearchResponse['fields'] | undefined;

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/attack_discovery/common_attributes.gen.ts

@@ -39,7 +39,7 @@ export const AttackDiscovery = z.object({
   /**
    * A short (no more than a sentence) summary of the attack discovery featuring only the host.name and user.name fields (when they are applicable), using the same syntax
    */
-  entitySummaryMarkdown: z.string(),
+  entitySummaryMarkdown: z.string().optional(),
   /**
    * An array of MITRE ATT&CK tactic for the attack discovery
    */
@@ -55,7 +55,7 @@ export const AttackDiscovery = z.object({
   /**
    * The time the attack discovery was generated
    */
-  timestamp: NonEmptyString,
+  timestamp: NonEmptyString.optional(),
 });
 
 /**

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/attack_discovery/common_attributes.schema.yaml

@@ -12,9 +12,7 @@ components:
       required:
         - 'alertIds'
         - 'detailsMarkdown'
-        - 'entitySummaryMarkdown'
         - 'summaryMarkdown'
-        - 'timestamp'
         - 'title'
       properties:
         alertIds:

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/evaluation/post_evaluate_route.gen.ts

@@ -22,10 +22,12 @@ export type PostEvaluateBody = z.infer<typeof PostEvaluateBody>;
 export const PostEvaluateBody = z.object({
   graphs: z.array(z.string()),
   datasetName: z.string(),
+  evaluatorConnectorId: z.string().optional(),
   connectorIds: z.array(z.string()),
   runName: z.string().optional(),
   alertsIndexPattern: z.string().optional().default('.alerts-security.alerts-default'),
   langSmithApiKey: z.string().optional(),
+  langSmithProject: z.string().optional(),
   replacements: Replacements.optional().default({}),
   size: z.number().optional().default(20),
 });

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/evaluation/post_evaluate_route.schema.yaml

@@ -61,6 +61,8 @@ components:
             type: string
         datasetName:
           type: string
+        evaluatorConnectorId:
+          type: string
         connectorIds:
           type: array
           items:
@@ -72,6 +74,8 @@ components:
           default: ".alerts-security.alerts-default"
         langSmithApiKey:
           type: string
+        langSmithProject:
+          type: string
         replacements:
           $ref: "../conversations/common_attributes.schema.yaml#/components/schemas/Replacements"
           default: {}

x-pack/packages/kbn-elastic-assistant-common/index.ts

@@ -25,3 +25,19 @@ export {
 export { transformRawData } from './impl/data_anonymization/transform_raw_data';
 export { parseBedrockBuffer, handleBedrockChunk } from './impl/utils/bedrock';
 export * from './constants';
+
+/** currently the same shape as "fields" property in the ES response */
+export { type MaybeRawData } from './impl/alerts/helpers/types';
+
+/**
+ * This query returns open and acknowledged (non-building block) alerts in the last 24 hours.
+ *
+ * The alerts are ordered by risk score, and then from the most recent to the oldest.
+ */
+export { getOpenAndAcknowledgedAlertsQuery } from './impl/alerts/get_open_and_acknowledged_alerts_query';
+
+/** Returns the raw data if it valid, or a default if it's not */
+export { getRawDataOrDefault } from './impl/alerts/helpers/get_raw_data_or_default';
+
+/** Return true if the provided size is out of range */
+export { sizeIsOutOfRange } from './impl/alerts/helpers/size_is_out_of_range';

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.tsx

@@ -16,7 +16,7 @@ import * as i18n from '../../../knowledge_base/translations';
 export const MIN_LATEST_ALERTS = 10;
 export const MAX_LATEST_ALERTS = 100;
 export const TICK_INTERVAL = 10;
-export const RANGE_CONTAINER_WIDTH = 300; // px
+export const RANGE_CONTAINER_WIDTH = 600; // px
 const LABEL_WRAPPER_MIN_WIDTH = 95; // px
 
 interface Props {
@@ -52,6 +52,7 @@ const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSetting
           <AlertsRange
             knowledgeBase={knowledgeBase}
             setUpdatedKnowledgeBaseSettings={setUpdatedKnowledgeBaseSettings}
+            value={knowledgeBase.latestAlerts}
           />
           <EuiSpacer size="s" />
         </EuiFlexItem>

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings_management.tsx

@@ -40,6 +40,7 @@ export const AlertsSettingsManagement: React.FC<Props> = React.memo(
           knowledgeBase={knowledgeBase}
           setUpdatedKnowledgeBaseSettings={setUpdatedKnowledgeBaseSettings}
           compressed={false}
+          value={knowledgeBase.latestAlerts}
         />
       </EuiPanel>
     );

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/evaluation_settings/evaluation_settings.tsx

@@ -17,28 +17,34 @@ import {
   EuiComboBox,
   EuiButton,
   EuiComboBoxOptionOption,
+  EuiComboBoxSingleSelectionShape,
   EuiTextColor,
   EuiFieldText,
+  EuiFieldNumber,
   EuiFlexItem,
   EuiFlexGroup,
   EuiLink,
   EuiPanel,
 } from '@elastic/eui';
-
 import { css } from '@emotion/react';
 import { FormattedMessage } from '@kbn/i18n-react';
 import type {
   GetEvaluateResponse,
   PostEvaluateRequestBodyInput,
 } from '@kbn/elastic-assistant-common';
+import { isEmpty } from 'lodash/fp';
+
 import * as i18n from './translations';
 import { useAssistantContext } from '../../../assistant_context';
+import { DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS } from '../../../assistant_context/constants';
 import { useLoadConnectors } from '../../../connectorland/use_load_connectors';
 import { getActionTypeTitle, getGenAiConfig } from '../../../connectorland/helpers';
 import { PRECONFIGURED_CONNECTOR } from '../../../connectorland/translations';
 import { usePerformEvaluation } from '../../api/evaluate/use_perform_evaluation';
 import { useEvaluationData } from '../../api/evaluate/use_evaluation_data';
 
+const AS_PLAIN_TEXT: EuiComboBoxSingleSelectionShape = { asPlainText: true };
+
 /**
  * Evaluation Settings -- development-only feature for evaluating models
  */
@@ -121,6 +127,18 @@ export const EvaluationSettings: React.FC = React.memo(() => {
     },
     [setSelectedModelOptions]
   );
+
+  const [selectedEvaluatorModel, setSelectedEvaluatorModel] = useState<
+    Array<EuiComboBoxOptionOption<string>>
+  >([]);
+
+  const onSelectedEvaluatorModelChange = useCallback(
+    (selected: Array<EuiComboBoxOptionOption<string>>) => setSelectedEvaluatorModel(selected),
+    []
+  );
+
+  const [size, setSize] = useState<string>(`${DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS}`);
+
   const visColorsBehindText = euiPaletteComplementary(connectors?.length ?? 0);
   const modelOptions = useMemo(() => {
     return (
@@ -170,19 +188,40 @@ export const EvaluationSettings: React.FC = React.memo(() => {
 
   // Perform Evaluation Button
   const handlePerformEvaluation = useCallback(async () => {
+    const evaluatorConnectorId =
+      selectedEvaluatorModel[0]?.key != null
+        ? { evaluatorConnectorId: selectedEvaluatorModel[0].key }
+        : {};
+
+    const langSmithApiKey = isEmpty(traceOptions.langSmithApiKey)
+      ? undefined
+      : traceOptions.langSmithApiKey;
+
+    const langSmithProject = isEmpty(traceOptions.langSmithProject)
+      ? undefined
+      : traceOptions.langSmithProject;
+
     const evalParams: PostEvaluateRequestBodyInput = {
       connectorIds: selectedModelOptions.flatMap((option) => option.key ?? []).sort(),
       graphs: selectedGraphOptions.map((option) => option.label).sort(),
       datasetName: selectedDatasetOptions[0]?.label,
+      ...evaluatorConnectorId,
+      langSmithApiKey,
+      langSmithProject,
       runName,
+      size: Number(size),
     };
     performEvaluation(evalParams);
   }, [
     performEvaluation,
     runName,
     selectedDatasetOptions,
+    selectedEvaluatorModel,
     selectedGraphOptions,
     selectedModelOptions,
+    size,
+    traceOptions.langSmithApiKey,
+    traceOptions.langSmithProject,
   ]);
 
   const getSection = (title: string, description: string) => (
@@ -355,6 +394,29 @@ export const EvaluationSettings: React.FC = React.memo(() => {
             onChange={onGraphOptionsChange}
           />
         </EuiFormRow>
+
+        <EuiFormRow
+          display="rowCompressed"
+          helpText={i18n.EVALUATOR_MODEL_DESCRIPTION}
+          label={i18n.EVALUATOR_MODEL}
+        >
+          <EuiComboBox
+            aria-label={i18n.EVALUATOR_MODEL}
+            compressed
+            onChange={onSelectedEvaluatorModelChange}
+            options={modelOptions}
+            selectedOptions={selectedEvaluatorModel}
+            singleSelection={AS_PLAIN_TEXT}
+          />
+        </EuiFormRow>
+
+        <EuiFormRow
+          display="rowCompressed"
+          helpText={i18n.DEFAULT_MAX_ALERTS_DESCRIPTION}
+          label={i18n.DEFAULT_MAX_ALERTS}
+        >
+          <EuiFieldNumber onChange={(e) => setSize(e.target.value)} value={size} />
+        </EuiFormRow>
       </EuiAccordion>
       <EuiHorizontalRule margin={'s'} />
       <EuiFlexGroup alignItems="center">

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/evaluation_settings/translations.ts

@@ -78,6 +78,36 @@ export const CONNECTORS_LABEL = i18n.translate(
   }
 );
 
+export const EVALUATOR_MODEL = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.evaluatorModelLabel',
+  {
+    defaultMessage: 'Evaluator model (optional)',
+  }
+);
+
+export const DEFAULT_MAX_ALERTS = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.defaultMaxAlertsLabel',
+  {
+    defaultMessage: 'Default max alerts',
+  }
+);
+
+export const EVALUATOR_MODEL_DESCRIPTION = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.evaluatorModelDescription',
+  {
+    defaultMessage:
+      'Judge the quality of all predictions using a single model. (Default: use the same model as the connector)',
+  }
+);
+
+export const DEFAULT_MAX_ALERTS_DESCRIPTION = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.defaultMaxAlertsDescription',
+  {
+    defaultMessage:
+      'The default maximum number of alerts to send as context, which may be overridden by the Example input',
+  }
+);
+
 export const CONNECTORS_DESCRIPTION = i18n.translate(
   'xpack.elasticAssistant.assistant.settings.evaluationSettings.connectorsDescription',
   {

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx

@@ -10,7 +10,9 @@ import { KnowledgeBaseConfig } from '../assistant/types';
 export const ATTACK_DISCOVERY_STORAGE_KEY = 'attackDiscovery';
 export const DEFAULT_ASSISTANT_NAMESPACE = 'elasticAssistantDefault';
 export const LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY = 'lastConversationId';
+export const MAX_ALERTS_LOCAL_STORAGE_KEY = 'maxAlerts';
 export const KNOWLEDGE_BASE_LOCAL_STORAGE_KEY = 'knowledgeBase';
+export const SHOW_SETTINGS_TOUR_LOCAL_STORAGE_KEY = 'showSettingsTour';
 export const STREAMING_LOCAL_STORAGE_KEY = 'streaming';
 export const TRACE_OPTIONS_SESSION_STORAGE_KEY = 'traceOptions';
 export const CONVERSATION_TABLE_SESSION_STORAGE_KEY = 'conversationTable';
@@ -21,6 +23,9 @@ export const ANONYMIZATION_TABLE_SESSION_STORAGE_KEY = 'anonymizationTable';
 /** The default `n` latest alerts, ordered by risk score, sent as context to the assistant */
 export const DEFAULT_LATEST_ALERTS = 20;
 
+/** The default maximum number of alerts to be sent as context when generating Attack discoveries */
+export const DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS = 200;
+
 export const DEFAULT_KNOWLEDGE_BASE_SETTINGS: KnowledgeBaseConfig = {
   latestAlerts: DEFAULT_LATEST_ALERTS,
 };

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx

@@ -261,7 +261,10 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       docLinks,
       getComments,
       http,
-      knowledgeBase: { ...DEFAULT_KNOWLEDGE_BASE_SETTINGS, ...localStorageKnowledgeBase },
+      knowledgeBase: {
+        ...DEFAULT_KNOWLEDGE_BASE_SETTINGS,
+        ...localStorageKnowledgeBase,
+      },
       promptContexts,
       navigateToApp,
       nameSpace,

x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.tsx

@@ -7,7 +7,7 @@
 
 import { EuiRange, useGeneratedHtmlId } from '@elastic/eui';
 import { css } from '@emotion/react';
-import React from 'react';
+import React, { useCallback } from 'react';
 import {
   MAX_LATEST_ALERTS,
   MIN_LATEST_ALERTS,
@@ -16,35 +16,57 @@ import {
 import { KnowledgeBaseConfig } from '../assistant/types';
 import { ALERTS_RANGE } from './translations';
 
+export type SingleRangeChangeEvent =
+  | React.ChangeEvent<HTMLInputElement>
+  | React.KeyboardEvent<HTMLInputElement>
+  | React.MouseEvent<HTMLButtonElement>;
+
 interface Props {
-  knowledgeBase: KnowledgeBaseConfig;
-  setUpdatedKnowledgeBaseSettings: React.Dispatch<React.SetStateAction<KnowledgeBaseConfig>>;
   compressed?: boolean;
+  maxAlerts?: number;
+  minAlerts?: number;
+  onChange?: (e: SingleRangeChangeEvent) => void;
+  knowledgeBase?: KnowledgeBaseConfig;
+  setUpdatedKnowledgeBaseSettings?: React.Dispatch<React.SetStateAction<KnowledgeBaseConfig>>;
+  step?: number;
+  value: string | number;
 }
 
 const MAX_ALERTS_RANGE_WIDTH = 649; // px
 
 export const AlertsRange: React.FC<Props> = React.memo(
-  ({ knowledgeBase, setUpdatedKnowledgeBaseSettings, compressed = true }) => {
+  ({
+    compressed = true,
+    knowledgeBase,
+    maxAlerts = MAX_LATEST_ALERTS,
+    minAlerts = MIN_LATEST_ALERTS,
+    onChange,
+    setUpdatedKnowledgeBaseSettings,
+    step = TICK_INTERVAL,
+    value,
+  }) => {
     const inputRangeSliderId = useGeneratedHtmlId({ prefix: 'inputRangeSlider' });
 
+    const handleOnChange = useCallback(
+      (e: SingleRangeChangeEvent) => {
+        if (knowledgeBase != null && setUpdatedKnowledgeBaseSettings != null) {
+          setUpdatedKnowledgeBaseSettings({
+            ...knowledgeBase,
+            latestAlerts: Number(e.currentTarget.value),
+          });
+        }
+
+        if (onChange != null) {
+          onChange(e);
+        }
+      },
+      [knowledgeBase, onChange, setUpdatedKnowledgeBaseSettings]
+    );
+
     return (
       <EuiRange
         aria-label={ALERTS_RANGE}
         compressed={compressed}
-        data-test-subj="alertsRange"
-        id={inputRangeSliderId}
-        max={MAX_LATEST_ALERTS}
-        min={MIN_LATEST_ALERTS}
-        onChange={(e) =>
-          setUpdatedKnowledgeBaseSettings({
-            ...knowledgeBase,
-            latestAlerts: Number(e.currentTarget.value),
-          })
-        }
-        showTicks
-        step={TICK_INTERVAL}
-        value={knowledgeBase.latestAlerts}
         css={css`
           max-inline-size: ${MAX_ALERTS_RANGE_WIDTH}px;
           & .euiRangeTrack {
@@ -52,6 +74,14 @@ export const AlertsRange: React.FC<Props> = React.memo(
             margin-inline-end: 0;
           }
         `}
+        data-test-subj="alertsRange"
+        id={inputRangeSliderId}
+        max={maxAlerts}
+        min={minAlerts}
+        onChange={handleOnChange}
+        showTicks
+        step={step}
+        value={value}
       />
     );
   }

x-pack/packages/kbn-elastic-assistant/index.ts

@@ -77,10 +77,17 @@ export { AssistantAvatar } from './impl/assistant/assistant_avatar/assistant_ava
 export { ConnectorSelectorInline } from './impl/connectorland/connector_selector_inline/connector_selector_inline';
 
 export {
+  /** The Attack discovery local storage key */
   ATTACK_DISCOVERY_STORAGE_KEY,
   DEFAULT_ASSISTANT_NAMESPACE,
+  /** The default maximum number of alerts to be sent as context when generating Attack discoveries */
+  DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS,
   DEFAULT_LATEST_ALERTS,
   KNOWLEDGE_BASE_LOCAL_STORAGE_KEY,
+  /** The local storage key that specifies the maximum number of alerts to send as context */
+  MAX_ALERTS_LOCAL_STORAGE_KEY,
+  /** The local storage key that specifies whether the settings tour should be shown */
+  SHOW_SETTINGS_TOUR_LOCAL_STORAGE_KEY,
 } from './impl/assistant_context/constants';
 
 export { useLoadConnectors } from './impl/connectorland/use_load_connectors';
@@ -140,3 +147,16 @@ export { mergeBaseWithPersistedConversations } from './impl/assistant/helpers';
 
 export { UpgradeButtons } from './impl/upgrade/upgrade_buttons';
 export { getUserConversations, getPrompts, bulkUpdatePrompts } from './impl/assistant/api';
+
+export {
+  /** A range slider component, typically used to configure the number of alerts sent as context */
+  AlertsRange,
+  /** This event occurs when the `AlertsRange` slider is changed */
+  type SingleRangeChangeEvent,
+} from './impl/knowledge_base/alerts_range';
+export {
+  /** A label instructing the user to send fewer alerts */
+  SELECT_FEWER_ALERTS,
+  /** Your anonymization settings will apply to these alerts (label) */
+  YOUR_ANONYMIZATION_SETTINGS,
+} from './impl/knowledge_base/translations';

x-pack/plugins/elastic_assistant/README.md

@@ -10,15 +10,21 @@ Maintained by the Security Solution team
 
 ## Graph structure
 
+### Default Assistant graph
+
 ![DefaultAssistantGraph](./docs/img/default_assistant_graph.png)
 
+### Default Attack discovery graph
+
+![DefaultAttackDiscoveryGraph](./docs/img/default_attack_discovery_graph.png)
+
 ## Development
 
 ### Generate graph structure
 
 To generate the graph structure, run `yarn draw-graph` from the plugin directory.
-The graph will be generated in the `docs/img` directory of the plugin.
+The graphs will be generated in the `docs/img` directory of the plugin.
 
 ### Testing
 
-To run the tests for this plugin, run `node scripts/jest --watch x-pack/plugins/elastic_assistant/jest.config.js --coverage` from the Kibana root directory.
+To run the tests for this plugin, run `node scripts/jest --watch x-pack/plugins/elastic_assistant/jest.config.js --coverage` from the Kibana root directory.

x-pack/plugins/elastic_assistant/scripts/draw_graph_script.ts

@@ -5,11 +5,13 @@
  * 2.0.
  */
 
+import type { ElasticsearchClient } from '@kbn/core/server';
 import { ToolingLog } from '@kbn/tooling-log';
 import fs from 'fs/promises';
 import path from 'path';
 import {
   ActionsClientChatOpenAI,
+  type ActionsClientLlm,
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server/language_models';
 import type { Logger } from '@kbn/logging';
@@ -17,6 +19,11 @@ import { ChatPromptTemplate } from '@langchain/core/prompts';
 import { FakeLLM } from '@langchain/core/utils/testing';
 import { createOpenAIFunctionsAgent } from 'langchain/agents';
 import { getDefaultAssistantGraph } from '../server/lib/langchain/graphs/default_assistant_graph/graph';
+import { getDefaultAttackDiscoveryGraph } from '../server/lib/attack_discovery/graphs/default_attack_discovery_graph';
+
+interface Drawable {
+  drawMermaidPng: () => Promise<Blob>;
+}
 
 // Just defining some test variables to get the graph to compile..
 const testPrompt = ChatPromptTemplate.fromMessages([
@@ -34,7 +41,7 @@ const createLlmInstance = () => {
   return mockLlm;
 };
 
-async function getGraph(logger: Logger) {
+async function getAssistantGraph(logger: Logger): Promise<Drawable> {
   const agentRunnable = await createOpenAIFunctionsAgent({
     llm: mockLlm,
     tools: [],
@@ -51,16 +58,49 @@ async function getGraph(logger: Logger) {
   return graph.getGraph();
 }
 
-export const draw = async () => {
+async function getAttackDiscoveryGraph(logger: Logger): Promise<Drawable> {
+  const mockEsClient = {} as unknown as ElasticsearchClient;
+
+  const graph = getDefaultAttackDiscoveryGraph({
+    anonymizationFields: [],
+    esClient: mockEsClient,
+    llm: mockLlm as unknown as ActionsClientLlm,
+    logger,
+    replacements: {},
+    size: 20,
+  });
+
+  return graph.getGraph();
+}
+
+export const drawGraph = async ({
+  getGraph,
+  outputFilename,
+}: {
+  getGraph: (logger: Logger) => Promise<Drawable>;
+  outputFilename: string;
+}) => {
   const logger = new ToolingLog({
     level: 'info',
     writeTo: process.stdout,
   }) as unknown as Logger;
   logger.info('Compiling graph');
-  const outputPath = path.join(__dirname, '../docs/img/default_assistant_graph.png');
+  const outputPath = path.join(__dirname, outputFilename);
   const graph = await getGraph(logger);
   const output = await graph.drawMermaidPng();
   const buffer = Buffer.from(await output.arrayBuffer());
   logger.info(`Writing graph to ${outputPath}`);
   await fs.writeFile(outputPath, buffer);
 };
+
+export const draw = async () => {
+  await drawGraph({
+    getGraph: getAssistantGraph,
+    outputFilename: '../docs/img/default_assistant_graph.png',
+  });
+
+  await drawGraph({
+    getGraph: getAttackDiscoveryGraph,
+    outputFilename: '../docs/img/default_attack_discovery_graph.png',
+  });
+};

x-pack/plugins/elastic_assistant/server/__mocks__/attack_discovery_schema.mock.ts

@@ -6,7 +6,7 @@
  */
 
 import { estypes } from '@elastic/elasticsearch';
-import { EsAttackDiscoverySchema } from '../ai_assistant_data_clients/attack_discovery/types';
+import { EsAttackDiscoverySchema } from '../lib/attack_discovery/persistence/types';
 
 export const getAttackDiscoverySearchEsMock = () => {
   const searchResponse: estypes.SearchResponse<EsAttackDiscoverySchema> = {

x-pack/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts

@@ -8,7 +8,7 @@
 import type { PublicMethodsOf } from '@kbn/utility-types';
 import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations';
 import { AIAssistantDataClient } from '../ai_assistant_data_clients';
-import { AttackDiscoveryDataClient } from '../ai_assistant_data_clients/attack_discovery';
+import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence';
 
 type ConversationsDataClientContract = PublicMethodsOf<AIAssistantConversationsDataClient>;
 export type ConversationsDataClientMock = jest.Mocked<ConversationsDataClientContract>;

x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts

@@ -26,7 +26,7 @@ import {
   GetAIAssistantKnowledgeBaseDataClientParams,
 } from '../ai_assistant_data_clients/knowledge_base';
 import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
-import { AttackDiscoveryDataClient } from '../ai_assistant_data_clients/attack_discovery';
+import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence';
 
 export const createMockClients = () => {
   const core = coreMock.createRequestHandlerContext();

x-pack/plugins/elastic_assistant/server/__mocks__/response.ts

@@ -16,7 +16,7 @@ import { getPromptsSearchEsMock } from './prompts_schema.mock';
 import { EsAnonymizationFieldsSchema } from '../ai_assistant_data_clients/anonymization_fields/types';
 import { getAnonymizationFieldsSearchEsMock } from './anonymization_fields_schema.mock';
 import { getAttackDiscoverySearchEsMock } from './attack_discovery_schema.mock';
-import { EsAttackDiscoverySchema } from '../ai_assistant_data_clients/attack_discovery/types';
+import { EsAttackDiscoverySchema } from '../lib/attack_discovery/persistence/types';
 
 export const responseMock = {
   create: httpServerMock.createResponseFactory,

x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts

@@ -11,7 +11,7 @@ import type { AuthenticatedUser, Logger, ElasticsearchClient } from '@kbn/core/s
 import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
 import type { MlPluginSetup } from '@kbn/ml-plugin/server';
 import { Subject } from 'rxjs';
-import { attackDiscoveryFieldMap } from '../ai_assistant_data_clients/attack_discovery/field_maps_configuration';
+import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration';
 import { getDefaultAnonymizationFields } from '../../common/anonymization';
 import { AssistantResourceNames, GetElser } from '../types';
 import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations';
@@ -34,7 +34,7 @@ import {
   AIAssistantKnowledgeBaseDataClient,
   GetAIAssistantKnowledgeBaseDataClientParams,
 } from '../ai_assistant_data_clients/knowledge_base';
-import { AttackDiscoveryDataClient } from '../ai_assistant_data_clients/attack_discovery';
+import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence';
 import { createGetElserId, createPipeline, pipelineExists } from './helpers';
 
 const TOTAL_FIELDS_LIMIT = 2500;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_examples.ts

@@ -0,0 +1,55 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Example } from 'langsmith/schemas';
+
+export const exampleWithReplacements: Example = {
+  id: '5D436078-B2CF-487A-A0FA-7CB46696F54E',
+  created_at: '2024-10-10T23:01:19.350232+00:00',
+  dataset_id: '0DA3497B-B084-4105-AFC0-2D8E05DE4B7C',
+  modified_at: '2024-10-10T23:01:19.350232+00:00',
+  inputs: {},
+  outputs: {
+    attackDiscoveries: [
+      {
+        title: 'Critical Malware and Phishing Alerts on host e1cb3cf0-30f3-4f99-a9c8-518b955c6f90',
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        timestamp: '2024-10-10T22:59:52.749Z',
+        detailsMarkdown:
+          '- On `2023-06-19T00:28:38.061Z` a critical malware detection alert was triggered on host {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} running {{ host.os.name macOS }} version {{ host.os.version 13.4 }}.\n- The malware was identified as {{ file.name unix1 }} with SHA256 hash {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\n- The process {{ process.name My Go Application.app }} was executed with command line {{ process.command_line /private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app }}.\n- The process was not trusted as its code signature failed to satisfy specified code requirements.\n- The user involved was {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}.\n- Another critical alert was triggered for potential credentials phishing via {{ process.name osascript }} on the same host.\n- The phishing attempt involved displaying a dialog to capture the user\'s password.\n- The process {{ process.name osascript }} was executed with command line {{ process.command_line osascript -e display dialog "MacOS wants to access System Preferences\\n\\nPlease enter your password." with title "System Preferences" with icon file "System:Library:CoreServices:CoreTypes.bundle:Contents:Resources:ToolbarAdvanced.icns" default answer "" giving up after 30 with hidden answer ¬ }}.\n- The MITRE ATT&CK tactics involved include Credential Access and Input Capture.',
+        summaryMarkdown:
+          'Critical malware and phishing alerts detected on {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} involving user {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}. Malware identified as {{ file.name unix1 }} and phishing attempt via {{ process.name osascript }}.',
+        mitreAttackTactics: ['Credential Access', 'Input Capture'],
+        entitySummaryMarkdown:
+          'Critical malware and phishing alerts detected on {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} involving user {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}.',
+      },
+    ],
+    replacements: {
+      '039c15c5-3964-43e7-a891-42fe2ceeb9ff': 'james',
+      '0b53f092-96dd-4282-bfb9-4f75a4530b80': 'root',
+      '1123bd7b-3afb-45d1-801a-108f04e7cfb7': 'SRVWIN04',
+      '3b9856bc-2c0d-4f1a-b9ae-32742e15ddd1': 'SRVWIN07',
+      '5306bcfd-2729-49e3-bdf0-678002778ccf': 'SRVWIN01',
+      '55af96a7-69b0-47cf-bf11-29be98a59eb0': 'SRVNIX05',
+      '66919fe3-16a4-4dfe-bc90-713f0b33a2ff': 'Administrator',
+      '9404361f-53fa-484f-adf8-24508256e70e': 'SRVWIN03',
+      'e1cb3cf0-30f3-4f99-a9c8-518b955c6f90': 'SRVMAC08',
+      'f59a00e2-f9c4-4069-8390-fd36ecd16918': 'SRVWIN02',
+      'fc6d07da-5186-4d59-9b79-9382b0c226b3': 'SRVWIN06',
+    },
+  },
+  runs: [],
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_runs.ts

@@ -0,0 +1,53 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Run } from 'langsmith/schemas';
+
+export const runWithReplacements: Run = {
+  id: 'B7B03FEE-9AC4-4823-AEDB-F8EC20EAD5C4',
+  inputs: {},
+  name: 'test',
+  outputs: {
+    attackDiscoveries: [
+      {
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        detailsMarkdown:
+          '- The attack began with the execution of a malicious file named `unix1` on the host `{{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }}` by the user `{{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}`.\n- The file `unix1` was detected at `{{ file.path /Users/james/unix1 }}` with a SHA256 hash of `{{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}`.\n- The process `{{ process.name My Go Application.app }}` was executed multiple times with different arguments, indicating potential persistence mechanisms.\n- The process `{{ process.name chmod }}` was used to change permissions of the file `unix1` to 777, making it executable.\n- A phishing attempt was detected via `osascript` on the same host, attempting to capture user credentials.\n- The attack involved multiple critical alerts, all indicating high-risk malware activity.',
+        entitySummaryMarkdown:
+          'The host `{{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }}` and user `{{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}` were involved in the attack.',
+        mitreAttackTactics: ['Initial Access', 'Execution', 'Persistence', 'Credential Access'],
+        summaryMarkdown:
+          'A series of critical malware alerts were detected on the host `{{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }}` involving the user `{{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}`. The attack included the execution of a malicious file `unix1`, permission changes, and a phishing attempt via `osascript`.',
+        title: 'Critical Malware Attack on macOS Host',
+        timestamp: '2024-10-11T17:55:59.702Z',
+      },
+    ],
+    replacements: {
+      '039c15c5-3964-43e7-a891-42fe2ceeb9ff': 'james',
+      '0b53f092-96dd-4282-bfb9-4f75a4530b80': 'root',
+      '1123bd7b-3afb-45d1-801a-108f04e7cfb7': 'SRVWIN04',
+      '3b9856bc-2c0d-4f1a-b9ae-32742e15ddd1': 'SRVWIN07',
+      '5306bcfd-2729-49e3-bdf0-678002778ccf': 'SRVWIN01',
+      '55af96a7-69b0-47cf-bf11-29be98a59eb0': 'SRVNIX05',
+      '66919fe3-16a4-4dfe-bc90-713f0b33a2ff': 'Administrator',
+      '9404361f-53fa-484f-adf8-24508256e70e': 'SRVWIN03',
+      'e1cb3cf0-30f3-4f99-a9c8-518b955c6f90': 'SRVMAC08',
+      'f59a00e2-f9c4-4069-8390-fd36ecd16918': 'SRVWIN02',
+      'fc6d07da-5186-4d59-9b79-9382b0c226b3': 'SRVWIN06',
+    },
+  },
+  run_type: 'evaluation',
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/constants.ts

@@ -0,0 +1,911 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
+
+export const DEFAULT_EVAL_ANONYMIZATION_FIELDS: AnonymizationFieldResponse[] = [
+  {
+    id: 'Mx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: '_id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'NB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: '@timestamp',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'NR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'cloud.availability_zone',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Nh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'cloud.provider',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Nx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'cloud.region',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'OB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'destination.ip',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'OR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'dns.question.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Oh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'dns.question.type',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Ox09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'event.category',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'PB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'event.dataset',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'PR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'event.module',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Ph09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'event.outcome',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Px09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'file.Ext.original.path',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'QB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'file.hash.sha256',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'QR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'file.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Qh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'file.path',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Qx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'group.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'RB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'group.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'RR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'host.asset.criticality',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Rh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'host.name',
+    allowed: true,
+    anonymized: true,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Rx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'host.os.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'SB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'host.os.version',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'SR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'host.risk.calculated_level',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Sh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'host.risk.calculated_score_norm',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Sx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.original_time',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'TB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.risk_score',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'TR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.description',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Th09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Tx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.references',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'UB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.framework',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'UR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.tactic.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Uh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.tactic.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Ux09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.tactic.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'VB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.technique.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'VR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.technique.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Vh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.technique.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Vx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.technique.subtechnique.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'WB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.technique.subtechnique.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'WR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.rule.threat.technique.subtechnique.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Wh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.severity',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Wx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'kibana.alert.workflow_status',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'XB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'message',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'XR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'network.protocol',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Xh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.args',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Xx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.code_signature.exists',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'YB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.code_signature.signing_id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'YR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.code_signature.status',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Yh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.code_signature.subject_name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Yx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.code_signature.trusted',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ZB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.command_line',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ZR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.executable',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Zh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.exit_code',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'Zx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.Ext.memory_region.bytes_compressed_present',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'aB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.Ext.memory_region.malware_signature.all_names',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'aR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.Ext.memory_region.malware_signature.primary.matches',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ah09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.Ext.memory_region.malware_signature.primary.signature.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ax09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.Ext.token.integrity_level_name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'bB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.hash.md5',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'bR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.hash.sha1',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'bh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.hash.sha256',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'bx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'cB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.args',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'cR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.args_count',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ch09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.code_signature.exists',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'cx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.code_signature.status',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'dB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.code_signature.subject_name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'dR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.code_signature.trusted',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'dh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.command_line',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'dx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.executable',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'eB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.parent.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'eR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.pe.original_file_name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'eh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.pid',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ex09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'process.working_directory',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'fB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.feature',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'fR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.data',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'fh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.entropy',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'fx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.extension',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'gB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.metrics',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'gR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.operation',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'gh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.path',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'gx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.files.score',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'hB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'Ransomware.version',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'hR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'rule.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'hh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'rule.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'hx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'source.ip',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'iB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.framework',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'iR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.tactic.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ih09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.tactic.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'ix09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.tactic.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'jB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.technique.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'jR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.technique.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'jh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.technique.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'jx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.technique.subtechnique.id',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'kB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.technique.subtechnique.name',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'kR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'threat.technique.subtechnique.reference',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'kh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'user.asset.criticality',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'kx09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'user.domain',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'lB09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'user.name',
+    allowed: true,
+    anonymized: true,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'lR09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'user.risk.calculated_level',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+  {
+    id: 'lh09VpEBOiz7eA-eF2fb',
+    timestamp: '2024-08-15T13:32:10.073Z',
+    field: 'user.risk.calculated_score_norm',
+    allowed: true,
+    anonymized: false,
+    createdAt: '2024-08-15T13:32:10.073Z',
+    namespace: 'default',
+  },
+];

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/example_input/index.test.ts

@@ -0,0 +1,75 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ExampleInput, ExampleInputWithOverrides } from '.';
+
+const validInput = {
+  attackDiscoveries: null,
+  attackDiscoveryPrompt: 'prompt',
+  anonymizedAlerts: [{ pageContent: 'content', metadata: { key: 'value' } }],
+  combinedGenerations: 'gen1gen2',
+  combinedRefinements: 'ref1ref2',
+  errors: ['error1', 'error2'],
+  generationAttempts: 1,
+  generations: ['gen1', 'gen2'],
+  hallucinationFailures: 0,
+  maxGenerationAttempts: 5,
+  maxHallucinationFailures: 2,
+  maxRepeatedGenerations: 3,
+  refinements: ['ref1', 'ref2'],
+  refinePrompt: 'refine prompt',
+  replacements: { key: 'replacement' },
+  unrefinedResults: null,
+};
+
+describe('ExampleInput Schema', () => {
+  it('validates a correct ExampleInput object', () => {
+    expect(() => ExampleInput.parse(validInput)).not.toThrow();
+  });
+
+  it('throws given an invalid ExampleInput', () => {
+    const invalidInput = {
+      attackDiscoveries: 'invalid', // should be an array or null
+    };
+
+    expect(() => ExampleInput.parse(invalidInput)).toThrow();
+  });
+
+  it('removes unknown properties', () => {
+    const hasUnknownProperties = {
+      ...validInput,
+      unknownProperty: 'unknown', // <-- should be removed
+    };
+
+    const parsed = ExampleInput.parse(hasUnknownProperties);
+
+    expect(parsed).not.toHaveProperty('unknownProperty');
+  });
+});
+
+describe('ExampleInputWithOverrides Schema', () => {
+  it('validates a correct ExampleInputWithOverrides object', () => {
+    const validInputWithOverrides = {
+      ...validInput,
+      overrides: {
+        attackDiscoveryPrompt: 'ad prompt override',
+        refinePrompt: 'refine prompt override',
+      },
+    };
+
+    expect(() => ExampleInputWithOverrides.parse(validInputWithOverrides)).not.toThrow();
+  });
+
+  it('throws when given an invalid ExampleInputWithOverrides object', () => {
+    const invalidInputWithOverrides = {
+      attackDiscoveries: null,
+      overrides: 'invalid', // should be an object
+    };
+
+    expect(() => ExampleInputWithOverrides.parse(invalidInputWithOverrides)).toThrow();
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/example_input/index.ts

@@ -0,0 +1,52 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AttackDiscovery, Replacements } from '@kbn/elastic-assistant-common';
+import { z } from '@kbn/zod';
+
+const Document = z.object({
+  pageContent: z.string(),
+  metadata: z.record(z.string(), z.any()),
+});
+
+type Document = z.infer<typeof Document>;
+
+/**
+ * Parses the input from an example in a LangSmith dataset
+ */
+export const ExampleInput = z.object({
+  attackDiscoveries: z.array(AttackDiscovery).nullable().optional(),
+  attackDiscoveryPrompt: z.string().optional(),
+  anonymizedAlerts: z.array(Document).optional(),
+  combinedGenerations: z.string().optional(),
+  combinedRefinements: z.string().optional(),
+  errors: z.array(z.string()).optional(),
+  generationAttempts: z.number().optional(),
+  generations: z.array(z.string()).optional(),
+  hallucinationFailures: z.number().optional(),
+  maxGenerationAttempts: z.number().optional(),
+  maxHallucinationFailures: z.number().optional(),
+  maxRepeatedGenerations: z.number().optional(),
+  refinements: z.array(z.string()).optional(),
+  refinePrompt: z.string().optional(),
+  replacements: Replacements.optional(),
+  unrefinedResults: z.array(AttackDiscovery).nullable().optional(),
+});
+
+export type ExampleInput = z.infer<typeof ExampleInput>;
+
+/**
+ * The optional overrides for an example input
+ */
+export const ExampleInputWithOverrides = z.intersection(
+  ExampleInput,
+  z.object({
+    overrides: ExampleInput.optional(),
+  })
+);
+
+export type ExampleInputWithOverrides = z.infer<typeof ExampleInputWithOverrides>;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/get_default_prompt_template/index.test.ts

@@ -0,0 +1,42 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getDefaultPromptTemplate } from '.';
+
+describe('getDefaultPromptTemplate', () => {
+  it('returns the expected prompt template', () => {
+    const expectedTemplate = `Evaluate based on how well the following submission follows the specified rubric. Grade only based on the rubric and "expected response":
+
+[BEGIN rubric]
+1. Is the submission non-empty and not null?
+2. Is the submission well-formed JSON?
+3. Evaluate the value of the "detailsMarkdown" field of all the "attackDiscoveries" in the submission json. Do the values of "detailsMarkdown" in the submission capture the essence of the "expected response", regardless of the order in which they appear, and highlight the same incident(s)?
+4. Evaluate the value of the "entitySummaryMarkdown" field of all the "attackDiscoveries" in the submission json. Does the value of "entitySummaryMarkdown" in the submission mention at least 50% the same entities as in the "expected response"?
+5. Evaluate the value of the "summaryMarkdown" field of all the "attackDiscoveries" in the submission json. Do the values of "summaryMarkdown" in the submission at least partially similar to that of the "expected response", regardless of the order in which they appear, and summarize the same incident(s)?
+6. Evaluate the value of the "title" field of all the "attackDiscoveries" in the submission json. Are the "title" values in the submission at least partially similar to the tile(s) of the "expected response", regardless of the order in which they appear, and mention the same incident(s)?
+7. Evaluate the value of the "alertIds" field of all the "attackDiscoveries" in the submission json. Do they match at least 100% of the "alertIds" in the submission?
+[END rubric]
+
+[BEGIN DATA]
+{input}
+[BEGIN submission]
+{output}
+[END submission]
+[BEGIN expected response]
+{reference}
+[END expected response]
+[END DATA]
+
+{criteria} Base your answer based on all the grading rubric items. If at least 5 of the 7 rubric items are correct, consider the submission correct. Write out your explanation for each criterion in the rubric, first in detail, then as a separate summary on a new line.
+
+Then finally respond with a single character, 'Y' or 'N', on a new line without any preceding or following characters. It's important that only a single character appears on the last line.`;
+
+    const result = getDefaultPromptTemplate();
+
+    expect(result).toBe(expectedTemplate);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/get_default_prompt_template/index.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getDefaultPromptTemplate =
+  () => `Evaluate based on how well the following submission follows the specified rubric. Grade only based on the rubric and "expected response":
+
+[BEGIN rubric]
+1. Is the submission non-empty and not null?
+2. Is the submission well-formed JSON?
+3. Evaluate the value of the "detailsMarkdown" field of all the "attackDiscoveries" in the submission json. Do the values of "detailsMarkdown" in the submission capture the essence of the "expected response", regardless of the order in which they appear, and highlight the same incident(s)?
+4. Evaluate the value of the "entitySummaryMarkdown" field of all the "attackDiscoveries" in the submission json. Does the value of "entitySummaryMarkdown" in the submission mention at least 50% the same entities as in the "expected response"?
+5. Evaluate the value of the "summaryMarkdown" field of all the "attackDiscoveries" in the submission json. Do the values of "summaryMarkdown" in the submission at least partially similar to that of the "expected response", regardless of the order in which they appear, and summarize the same incident(s)?
+6. Evaluate the value of the "title" field of all the "attackDiscoveries" in the submission json. Are the "title" values in the submission at least partially similar to the tile(s) of the "expected response", regardless of the order in which they appear, and mention the same incident(s)?
+7. Evaluate the value of the "alertIds" field of all the "attackDiscoveries" in the submission json. Do they match at least 100% of the "alertIds" in the submission?
+[END rubric]
+
+[BEGIN DATA]
+{input}
+[BEGIN submission]
+{output}
+[END submission]
+[BEGIN expected response]
+{reference}
+[END expected response]
+[END DATA]
+
+{criteria} Base your answer based on all the grading rubric items. If at least 5 of the 7 rubric items are correct, consider the submission correct. Write out your explanation for each criterion in the rubric, first in detail, then as a separate summary on a new line.
+
+Then finally respond with a single character, 'Y' or 'N', on a new line without any preceding or following characters. It's important that only a single character appears on the last line.`;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/get_example_attack_discoveries_with_replacements/index.test.ts

@@ -0,0 +1,125 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { omit } from 'lodash/fp';
+
+import { getExampleAttackDiscoveriesWithReplacements } from '.';
+import { exampleWithReplacements } from '../../../__mocks__/mock_examples';
+
+describe('getExampleAttackDiscoveriesWithReplacements', () => {
+  it('returns attack discoveries with replacements applied to the detailsMarkdown, entitySummaryMarkdown, summaryMarkdown, and title', () => {
+    const result = getExampleAttackDiscoveriesWithReplacements(exampleWithReplacements);
+
+    expect(result).toEqual([
+      {
+        title: 'Critical Malware and Phishing Alerts on host SRVMAC08',
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        timestamp: '2024-10-10T22:59:52.749Z',
+        detailsMarkdown:
+          '- On `2023-06-19T00:28:38.061Z` a critical malware detection alert was triggered on host {{ host.name SRVMAC08 }} running {{ host.os.name macOS }} version {{ host.os.version 13.4 }}.\n- The malware was identified as {{ file.name unix1 }} with SHA256 hash {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\n- The process {{ process.name My Go Application.app }} was executed with command line {{ process.command_line /private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app }}.\n- The process was not trusted as its code signature failed to satisfy specified code requirements.\n- The user involved was {{ user.name james }}.\n- Another critical alert was triggered for potential credentials phishing via {{ process.name osascript }} on the same host.\n- The phishing attempt involved displaying a dialog to capture the user\'s password.\n- The process {{ process.name osascript }} was executed with command line {{ process.command_line osascript -e display dialog "MacOS wants to access System Preferences\\n\\nPlease enter your password." with title "System Preferences" with icon file "System:Library:CoreServices:CoreTypes.bundle:Contents:Resources:ToolbarAdvanced.icns" default answer "" giving up after 30 with hidden answer ¬ }}.\n- The MITRE ATT&CK tactics involved include Credential Access and Input Capture.',
+        summaryMarkdown:
+          'Critical malware and phishing alerts detected on {{ host.name SRVMAC08 }} involving user {{ user.name james }}. Malware identified as {{ file.name unix1 }} and phishing attempt via {{ process.name osascript }}.',
+        mitreAttackTactics: ['Credential Access', 'Input Capture'],
+        entitySummaryMarkdown:
+          'Critical malware and phishing alerts detected on {{ host.name SRVMAC08 }} involving user {{ user.name james }}.',
+      },
+    ]);
+  });
+
+  it('returns an empty entitySummaryMarkdown when the entitySummaryMarkdown is missing', () => {
+    const missingEntitySummaryMarkdown = omit(
+      'entitySummaryMarkdown',
+      exampleWithReplacements.outputs?.attackDiscoveries?.[0]
+    );
+
+    const exampleWithMissingEntitySummaryMarkdown = {
+      ...exampleWithReplacements,
+      outputs: {
+        ...exampleWithReplacements.outputs,
+        attackDiscoveries: [missingEntitySummaryMarkdown],
+      },
+    };
+
+    const result = getExampleAttackDiscoveriesWithReplacements(
+      exampleWithMissingEntitySummaryMarkdown
+    );
+
+    expect(result).toEqual([
+      {
+        title: 'Critical Malware and Phishing Alerts on host SRVMAC08',
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        timestamp: '2024-10-10T22:59:52.749Z',
+        detailsMarkdown:
+          '- On `2023-06-19T00:28:38.061Z` a critical malware detection alert was triggered on host {{ host.name SRVMAC08 }} running {{ host.os.name macOS }} version {{ host.os.version 13.4 }}.\n- The malware was identified as {{ file.name unix1 }} with SHA256 hash {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\n- The process {{ process.name My Go Application.app }} was executed with command line {{ process.command_line /private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app }}.\n- The process was not trusted as its code signature failed to satisfy specified code requirements.\n- The user involved was {{ user.name james }}.\n- Another critical alert was triggered for potential credentials phishing via {{ process.name osascript }} on the same host.\n- The phishing attempt involved displaying a dialog to capture the user\'s password.\n- The process {{ process.name osascript }} was executed with command line {{ process.command_line osascript -e display dialog "MacOS wants to access System Preferences\\n\\nPlease enter your password." with title "System Preferences" with icon file "System:Library:CoreServices:CoreTypes.bundle:Contents:Resources:ToolbarAdvanced.icns" default answer "" giving up after 30 with hidden answer ¬ }}.\n- The MITRE ATT&CK tactics involved include Credential Access and Input Capture.',
+        summaryMarkdown:
+          'Critical malware and phishing alerts detected on {{ host.name SRVMAC08 }} involving user {{ user.name james }}. Malware identified as {{ file.name unix1 }} and phishing attempt via {{ process.name osascript }}.',
+        mitreAttackTactics: ['Credential Access', 'Input Capture'],
+        entitySummaryMarkdown: '',
+      },
+    ]);
+  });
+
+  it('throws when an example is undefined', () => {
+    expect(() => getExampleAttackDiscoveriesWithReplacements(undefined)).toThrowError();
+  });
+
+  it('throws when the example is missing attackDiscoveries', () => {
+    const missingAttackDiscoveries = {
+      ...exampleWithReplacements,
+      outputs: {
+        replacements: { ...exampleWithReplacements.outputs?.replacements },
+      },
+    };
+
+    expect(() =>
+      getExampleAttackDiscoveriesWithReplacements(missingAttackDiscoveries)
+    ).toThrowError();
+  });
+
+  it('throws when attackDiscoveries is null', () => {
+    const nullAttackDiscoveries = {
+      ...exampleWithReplacements,
+      outputs: {
+        attackDiscoveries: null,
+        replacements: { ...exampleWithReplacements.outputs?.replacements },
+      },
+    };
+
+    expect(() => getExampleAttackDiscoveriesWithReplacements(nullAttackDiscoveries)).toThrowError();
+  });
+
+  it('returns the original attack discoveries when replacements are missing', () => {
+    const missingReplacements = {
+      ...exampleWithReplacements,
+      outputs: {
+        attackDiscoveries: [...exampleWithReplacements.outputs?.attackDiscoveries],
+      },
+    };
+
+    const result = getExampleAttackDiscoveriesWithReplacements(missingReplacements);
+
+    expect(result).toEqual(exampleWithReplacements.outputs?.attackDiscoveries);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/get_example_attack_discoveries_with_replacements/index.ts

@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AttackDiscoveries, Replacements } from '@kbn/elastic-assistant-common';
+import type { Example } from 'langsmith/schemas';
+
+import { getDiscoveriesWithOriginalValues } from '../../get_discoveries_with_original_values';
+
+export const getExampleAttackDiscoveriesWithReplacements = (
+  example: Example | undefined
+): AttackDiscoveries => {
+  const exampleAttackDiscoveries = example?.outputs?.attackDiscoveries;
+  const exampleReplacements = example?.outputs?.replacements ?? {};
+
+  // NOTE: calls to `parse` throw an error if the Example input is invalid
+  const validatedAttackDiscoveries = AttackDiscoveries.parse(exampleAttackDiscoveries);
+  const validatedReplacements = Replacements.parse(exampleReplacements);
+
+  const withReplacements = getDiscoveriesWithOriginalValues({
+    attackDiscoveries: validatedAttackDiscoveries,
+    replacements: validatedReplacements,
+  });
+
+  return withReplacements;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/get_run_attack_discoveries_with_replacements/index.test.ts

@@ -0,0 +1,117 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { omit } from 'lodash/fp';
+
+import { getRunAttackDiscoveriesWithReplacements } from '.';
+import { runWithReplacements } from '../../../__mocks__/mock_runs';
+
+describe('getRunAttackDiscoveriesWithReplacements', () => {
+  it('returns attack discoveries with replacements applied to the detailsMarkdown, entitySummaryMarkdown, summaryMarkdown, and title', () => {
+    const result = getRunAttackDiscoveriesWithReplacements(runWithReplacements);
+
+    expect(result).toEqual([
+      {
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        detailsMarkdown:
+          '- The attack began with the execution of a malicious file named `unix1` on the host `{{ host.name SRVMAC08 }}` by the user `{{ user.name james }}`.\n- The file `unix1` was detected at `{{ file.path /Users/james/unix1 }}` with a SHA256 hash of `{{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}`.\n- The process `{{ process.name My Go Application.app }}` was executed multiple times with different arguments, indicating potential persistence mechanisms.\n- The process `{{ process.name chmod }}` was used to change permissions of the file `unix1` to 777, making it executable.\n- A phishing attempt was detected via `osascript` on the same host, attempting to capture user credentials.\n- The attack involved multiple critical alerts, all indicating high-risk malware activity.',
+        entitySummaryMarkdown:
+          'The host `{{ host.name SRVMAC08 }}` and user `{{ user.name james }}` were involved in the attack.',
+        mitreAttackTactics: ['Initial Access', 'Execution', 'Persistence', 'Credential Access'],
+        summaryMarkdown:
+          'A series of critical malware alerts were detected on the host `{{ host.name SRVMAC08 }}` involving the user `{{ user.name james }}`. The attack included the execution of a malicious file `unix1`, permission changes, and a phishing attempt via `osascript`.',
+        title: 'Critical Malware Attack on macOS Host',
+        timestamp: '2024-10-11T17:55:59.702Z',
+      },
+    ]);
+  });
+
+  it("returns an empty entitySummaryMarkdown when it's missing from the attack discovery", () => {
+    const missingEntitySummaryMarkdown = omit(
+      'entitySummaryMarkdown',
+      runWithReplacements.outputs?.attackDiscoveries?.[0]
+    );
+
+    const runWithMissingEntitySummaryMarkdown = {
+      ...runWithReplacements,
+      outputs: {
+        ...runWithReplacements.outputs,
+        attackDiscoveries: [missingEntitySummaryMarkdown],
+      },
+    };
+
+    const result = getRunAttackDiscoveriesWithReplacements(runWithMissingEntitySummaryMarkdown);
+
+    expect(result).toEqual([
+      {
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        detailsMarkdown:
+          '- The attack began with the execution of a malicious file named `unix1` on the host `{{ host.name SRVMAC08 }}` by the user `{{ user.name james }}`.\n- The file `unix1` was detected at `{{ file.path /Users/james/unix1 }}` with a SHA256 hash of `{{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}`.\n- The process `{{ process.name My Go Application.app }}` was executed multiple times with different arguments, indicating potential persistence mechanisms.\n- The process `{{ process.name chmod }}` was used to change permissions of the file `unix1` to 777, making it executable.\n- A phishing attempt was detected via `osascript` on the same host, attempting to capture user credentials.\n- The attack involved multiple critical alerts, all indicating high-risk malware activity.',
+        entitySummaryMarkdown: '',
+        mitreAttackTactics: ['Initial Access', 'Execution', 'Persistence', 'Credential Access'],
+        summaryMarkdown:
+          'A series of critical malware alerts were detected on the host `{{ host.name SRVMAC08 }}` involving the user `{{ user.name james }}`. The attack included the execution of a malicious file `unix1`, permission changes, and a phishing attempt via `osascript`.',
+        title: 'Critical Malware Attack on macOS Host',
+        timestamp: '2024-10-11T17:55:59.702Z',
+      },
+    ]);
+  });
+
+  it('throws when the run is missing attackDiscoveries', () => {
+    const missingAttackDiscoveries = {
+      ...runWithReplacements,
+      outputs: {
+        replacements: { ...runWithReplacements.outputs?.replacements },
+      },
+    };
+
+    expect(() => getRunAttackDiscoveriesWithReplacements(missingAttackDiscoveries)).toThrowError();
+  });
+
+  it('throws when attackDiscoveries is null', () => {
+    const nullAttackDiscoveries = {
+      ...runWithReplacements,
+      outputs: {
+        attackDiscoveries: null,
+        replacements: { ...runWithReplacements.outputs?.replacements },
+      },
+    };
+
+    expect(() => getRunAttackDiscoveriesWithReplacements(nullAttackDiscoveries)).toThrowError();
+  });
+
+  it('returns the original attack discoveries when replacements are missing', () => {
+    const missingReplacements = {
+      ...runWithReplacements,
+      outputs: {
+        attackDiscoveries: [...runWithReplacements.outputs?.attackDiscoveries],
+      },
+    };
+
+    const result = getRunAttackDiscoveriesWithReplacements(missingReplacements);
+
+    expect(result).toEqual(runWithReplacements.outputs?.attackDiscoveries);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/get_run_attack_discoveries_with_replacements/index.ts

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AttackDiscoveries, Replacements } from '@kbn/elastic-assistant-common';
+import type { Run } from 'langsmith/schemas';
+
+import { getDiscoveriesWithOriginalValues } from '../../get_discoveries_with_original_values';
+
+export const getRunAttackDiscoveriesWithReplacements = (run: Run): AttackDiscoveries => {
+  const runAttackDiscoveries = run.outputs?.attackDiscoveries;
+  const runReplacements = run.outputs?.replacements ?? {};
+
+  // NOTE: calls to `parse` throw an error if the Run Input is invalid
+  const validatedAttackDiscoveries = AttackDiscoveries.parse(runAttackDiscoveries);
+  const validatedReplacements = Replacements.parse(runReplacements);
+
+  const withReplacements = getDiscoveriesWithOriginalValues({
+    attackDiscoveries: validatedAttackDiscoveries,
+    replacements: validatedReplacements,
+  });
+
+  return withReplacements;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/index.test.ts

@@ -0,0 +1,98 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { PromptTemplate } from '@langchain/core/prompts';
+import type { ActionsClientLlm } from '@kbn/langchain/server';
+import { loadEvaluator } from 'langchain/evaluation';
+
+import { type GetCustomEvaluatorOptions, getCustomEvaluator } from '.';
+import { getDefaultPromptTemplate } from './get_default_prompt_template';
+import { getExampleAttackDiscoveriesWithReplacements } from './get_example_attack_discoveries_with_replacements';
+import { getRunAttackDiscoveriesWithReplacements } from './get_run_attack_discoveries_with_replacements';
+import { exampleWithReplacements } from '../../__mocks__/mock_examples';
+import { runWithReplacements } from '../../__mocks__/mock_runs';
+
+const mockLlm = jest.fn() as unknown as ActionsClientLlm;
+
+jest.mock('langchain/evaluation', () => ({
+  ...jest.requireActual('langchain/evaluation'),
+  loadEvaluator: jest.fn().mockResolvedValue({
+    evaluateStrings: jest.fn().mockResolvedValue({
+      key: 'correctness',
+      score: 0.9,
+    }),
+  }),
+}));
+
+const options: GetCustomEvaluatorOptions = {
+  criteria: 'correctness',
+  key: 'attack_discovery_correctness',
+  llm: mockLlm,
+  template: getDefaultPromptTemplate(),
+};
+
+describe('getCustomEvaluator', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  it('returns an evaluator function', () => {
+    const evaluator = getCustomEvaluator(options);
+
+    expect(typeof evaluator).toBe('function');
+  });
+
+  it('calls loadEvaluator with the expected arguments', async () => {
+    const evaluator = getCustomEvaluator(options);
+
+    await evaluator(runWithReplacements, exampleWithReplacements);
+
+    expect(loadEvaluator).toHaveBeenCalledWith('labeled_criteria', {
+      criteria: options.criteria,
+      chainOptions: {
+        prompt: PromptTemplate.fromTemplate(options.template),
+      },
+      llm: mockLlm,
+    });
+  });
+
+  it('calls evaluateStrings with the expected arguments', async () => {
+    const mockEvaluateStrings = jest.fn().mockResolvedValue({
+      key: 'correctness',
+      score: 0.9,
+    });
+
+    (loadEvaluator as jest.Mock).mockResolvedValue({
+      evaluateStrings: mockEvaluateStrings,
+    });
+
+    const evaluator = getCustomEvaluator(options);
+
+    await evaluator(runWithReplacements, exampleWithReplacements);
+
+    const prediction = getRunAttackDiscoveriesWithReplacements(runWithReplacements);
+    const reference = getExampleAttackDiscoveriesWithReplacements(exampleWithReplacements);
+
+    expect(mockEvaluateStrings).toHaveBeenCalledWith({
+      input: '',
+      prediction: JSON.stringify(prediction, null, 2),
+      reference: JSON.stringify(reference, null, 2),
+    });
+  });
+
+  it('returns the expected result', async () => {
+    const evaluator = getCustomEvaluator(options);
+
+    const result = await evaluator(runWithReplacements, exampleWithReplacements);
+
+    expect(result).toEqual({ key: 'attack_discovery_correctness', score: 0.9 });
+  });
+
+  it('throws given an undefined example', async () => {
+    const evaluator = getCustomEvaluator(options);
+
+    await expect(async () => evaluator(runWithReplacements, undefined)).rejects.toThrow();
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_custom_evaluator/index.ts

@@ -0,0 +1,69 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClientLlm } from '@kbn/langchain/server';
+import { PromptTemplate } from '@langchain/core/prompts';
+import type { EvaluationResult } from 'langsmith/evaluation';
+import type { Run, Example } from 'langsmith/schemas';
+import { CriteriaLike, loadEvaluator } from 'langchain/evaluation';
+
+import { getExampleAttackDiscoveriesWithReplacements } from './get_example_attack_discoveries_with_replacements';
+import { getRunAttackDiscoveriesWithReplacements } from './get_run_attack_discoveries_with_replacements';
+
+export interface GetCustomEvaluatorOptions {
+  /**
+   * Examples:
+   *  - "conciseness"
+   *  - "relevance"
+   *  - "correctness"
+   *  - "detail"
+   */
+  criteria: CriteriaLike;
+  /**
+   * The evaluation score will use this key
+   */
+  key: string;
+  /**
+   * LLm to use for evaluation
+   */
+  llm: ActionsClientLlm;
+  /**
+   * A prompt template that uses the {input}, {submission}, and {reference} variables
+   */
+  template: string;
+}
+
+export type CustomEvaluator = (
+  rootRun: Run,
+  example: Example | undefined
+) => Promise<EvaluationResult>;
+
+export const getCustomEvaluator =
+  ({ criteria, key, llm, template }: GetCustomEvaluatorOptions): CustomEvaluator =>
+  async (rootRun, example) => {
+    const chain = await loadEvaluator('labeled_criteria', {
+      criteria,
+      chainOptions: {
+        prompt: PromptTemplate.fromTemplate(template),
+      },
+      llm,
+    });
+
+    const exampleAttackDiscoveriesWithReplacements =
+      getExampleAttackDiscoveriesWithReplacements(example);
+
+    const runAttackDiscoveriesWithReplacements = getRunAttackDiscoveriesWithReplacements(rootRun);
+
+    // NOTE: res contains a score, as well as the reasoning for the score
+    const res = await chain.evaluateStrings({
+      input: '', // empty for now, but this could be the alerts, i.e. JSON.stringify(rootRun.outputs?.anonymizedAlerts, null, 2),
+      prediction: JSON.stringify(runAttackDiscoveriesWithReplacements, null, 2),
+      reference: JSON.stringify(exampleAttackDiscoveriesWithReplacements, null, 2),
+    });
+
+    return { key, score: res.score };
+  };

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_discoveries_with_original_values/index.test.ts

@@ -0,0 +1,79 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AttackDiscovery } from '@kbn/elastic-assistant-common';
+import { omit } from 'lodash/fp';
+
+import { getDiscoveriesWithOriginalValues } from '.';
+import { runWithReplacements } from '../../__mocks__/mock_runs';
+
+describe('getDiscoveriesWithOriginalValues', () => {
+  it('returns attack discoveries with replacements applied to the detailsMarkdown, entitySummaryMarkdown, summaryMarkdown, and title', () => {
+    const result = getDiscoveriesWithOriginalValues({
+      attackDiscoveries: runWithReplacements.outputs?.attackDiscoveries,
+      replacements: runWithReplacements.outputs?.replacements,
+    });
+
+    expect(result).toEqual([
+      {
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        detailsMarkdown:
+          '- The attack began with the execution of a malicious file named `unix1` on the host `{{ host.name SRVMAC08 }}` by the user `{{ user.name james }}`.\n- The file `unix1` was detected at `{{ file.path /Users/james/unix1 }}` with a SHA256 hash of `{{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}`.\n- The process `{{ process.name My Go Application.app }}` was executed multiple times with different arguments, indicating potential persistence mechanisms.\n- The process `{{ process.name chmod }}` was used to change permissions of the file `unix1` to 777, making it executable.\n- A phishing attempt was detected via `osascript` on the same host, attempting to capture user credentials.\n- The attack involved multiple critical alerts, all indicating high-risk malware activity.',
+        entitySummaryMarkdown:
+          'The host `{{ host.name SRVMAC08 }}` and user `{{ user.name james }}` were involved in the attack.',
+        mitreAttackTactics: ['Initial Access', 'Execution', 'Persistence', 'Credential Access'],
+        summaryMarkdown:
+          'A series of critical malware alerts were detected on the host `{{ host.name SRVMAC08 }}` involving the user `{{ user.name james }}`. The attack included the execution of a malicious file `unix1`, permission changes, and a phishing attempt via `osascript`.',
+        title: 'Critical Malware Attack on macOS Host',
+        timestamp: '2024-10-11T17:55:59.702Z',
+      },
+    ]);
+  });
+
+  it("returns an empty entitySummaryMarkdown when it's missing from the attack discovery", () => {
+    const missingEntitySummaryMarkdown = omit(
+      'entitySummaryMarkdown',
+      runWithReplacements.outputs?.attackDiscoveries?.[0]
+    ) as unknown as AttackDiscovery;
+
+    const result = getDiscoveriesWithOriginalValues({
+      attackDiscoveries: [missingEntitySummaryMarkdown],
+      replacements: runWithReplacements.outputs?.replacements,
+    });
+    expect(result).toEqual([
+      {
+        alertIds: [
+          '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16',
+          'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b',
+          '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c',
+          '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608',
+          'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac',
+          '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080',
+          '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6',
+          '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998',
+        ],
+        detailsMarkdown:
+          '- The attack began with the execution of a malicious file named `unix1` on the host `{{ host.name SRVMAC08 }}` by the user `{{ user.name james }}`.\n- The file `unix1` was detected at `{{ file.path /Users/james/unix1 }}` with a SHA256 hash of `{{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}`.\n- The process `{{ process.name My Go Application.app }}` was executed multiple times with different arguments, indicating potential persistence mechanisms.\n- The process `{{ process.name chmod }}` was used to change permissions of the file `unix1` to 777, making it executable.\n- A phishing attempt was detected via `osascript` on the same host, attempting to capture user credentials.\n- The attack involved multiple critical alerts, all indicating high-risk malware activity.',
+        entitySummaryMarkdown: '',
+        mitreAttackTactics: ['Initial Access', 'Execution', 'Persistence', 'Credential Access'],
+        summaryMarkdown:
+          'A series of critical malware alerts were detected on the host `{{ host.name SRVMAC08 }}` involving the user `{{ user.name james }}`. The attack included the execution of a malicious file `unix1`, permission changes, and a phishing attempt via `osascript`.',
+        title: 'Critical Malware Attack on macOS Host',
+        timestamp: '2024-10-11T17:55:59.702Z',
+      },
+    ]);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_discoveries_with_original_values/index.ts

@@ -0,0 +1,39 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import {
+  type AttackDiscovery,
+  Replacements,
+  replaceAnonymizedValuesWithOriginalValues,
+} from '@kbn/elastic-assistant-common';
+
+export const getDiscoveriesWithOriginalValues = ({
+  attackDiscoveries,
+  replacements,
+}: {
+  attackDiscoveries: AttackDiscovery[];
+  replacements: Replacements;
+}): AttackDiscovery[] =>
+  attackDiscoveries.map((attackDiscovery) => ({
+    ...attackDiscovery,
+    detailsMarkdown: replaceAnonymizedValuesWithOriginalValues({
+      messageContent: attackDiscovery.detailsMarkdown,
+      replacements,
+    }),
+    entitySummaryMarkdown: replaceAnonymizedValuesWithOriginalValues({
+      messageContent: attackDiscovery.entitySummaryMarkdown ?? '',
+      replacements,
+    }),
+    summaryMarkdown: replaceAnonymizedValuesWithOriginalValues({
+      messageContent: attackDiscovery.summaryMarkdown,
+      replacements,
+    }),
+    title: replaceAnonymizedValuesWithOriginalValues({
+      messageContent: attackDiscovery.title,
+      replacements,
+    }),
+  }));

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_evaluator_llm/index.test.ts

@@ -0,0 +1,161 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClient } from '@kbn/actions-plugin/server';
+import type { Connector } from '@kbn/actions-plugin/server/application/connector/types';
+import { ActionsClientLlm } from '@kbn/langchain/server';
+import { loggerMock } from '@kbn/logging-mocks';
+
+import { getEvaluatorLlm } from '.';
+
+jest.mock('@kbn/langchain/server', () => ({
+  ...jest.requireActual('@kbn/langchain/server'),
+
+  ActionsClientLlm: jest.fn(),
+}));
+
+const connectorTimeout = 1000;
+
+const evaluatorConnectorId = 'evaluator-connector-id';
+const evaluatorConnector = {
+  id: 'evaluatorConnectorId',
+  actionTypeId: '.gen-ai',
+  name: 'GPT-4o',
+  isPreconfigured: true,
+  isSystemAction: false,
+  isDeprecated: false,
+} as Connector;
+
+const experimentConnector: Connector = {
+  name: 'Gemini 1.5 Pro 002',
+  actionTypeId: '.gemini',
+  config: {
+    apiUrl: 'https://example.com',
+    defaultModel: 'gemini-1.5-pro-002',
+    gcpRegion: 'test-region',
+    gcpProjectID: 'test-project-id',
+  },
+  secrets: {
+    credentialsJson: '{}',
+  },
+  id: 'gemini-1-5-pro-002',
+  isPreconfigured: true,
+  isSystemAction: false,
+  isDeprecated: false,
+} as Connector;
+
+const logger = loggerMock.create();
+
+describe('getEvaluatorLlm', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  describe('getting the evaluation connector', () => {
+    it("calls actionsClient.get with the evaluator connector ID when it's provided", async () => {
+      const actionsClient = {
+        get: jest.fn(),
+      } as unknown as ActionsClient;
+
+      await getEvaluatorLlm({
+        actionsClient,
+        connectorTimeout,
+        evaluatorConnectorId,
+        experimentConnector,
+        langSmithApiKey: undefined,
+        logger,
+      });
+
+      expect(actionsClient.get).toHaveBeenCalledWith({
+        id: evaluatorConnectorId,
+        throwIfSystemAction: false,
+      });
+    });
+
+    it("calls actionsClient.get with the experiment connector ID when the evaluator connector ID isn't provided", async () => {
+      const actionsClient = {
+        get: jest.fn().mockResolvedValue(null),
+      } as unknown as ActionsClient;
+
+      await getEvaluatorLlm({
+        actionsClient,
+        connectorTimeout,
+        evaluatorConnectorId: undefined,
+        experimentConnector,
+        langSmithApiKey: undefined,
+        logger,
+      });
+
+      expect(actionsClient.get).toHaveBeenCalledWith({
+        id: experimentConnector.id,
+        throwIfSystemAction: false,
+      });
+    });
+
+    it('falls back to the experiment connector when the evaluator connector is not found', async () => {
+      const actionsClient = {
+        get: jest.fn().mockResolvedValue(null),
+      } as unknown as ActionsClient;
+
+      await getEvaluatorLlm({
+        actionsClient,
+        connectorTimeout,
+        evaluatorConnectorId,
+        experimentConnector,
+        langSmithApiKey: undefined,
+        logger,
+      });
+
+      expect(ActionsClientLlm).toHaveBeenCalledWith(
+        expect.objectContaining({
+          connectorId: experimentConnector.id,
+        })
+      );
+    });
+  });
+
+  it('logs the expected connector names and types', async () => {
+    const actionsClient = {
+      get: jest.fn().mockResolvedValue(evaluatorConnector),
+    } as unknown as ActionsClient;
+
+    await getEvaluatorLlm({
+      actionsClient,
+      connectorTimeout,
+      evaluatorConnectorId,
+      experimentConnector,
+      langSmithApiKey: undefined,
+      logger,
+    });
+
+    expect(logger.info).toHaveBeenCalledWith(
+      `The ${evaluatorConnector.name} (openai) connector will judge output from the ${experimentConnector.name} (gemini) connector`
+    );
+  });
+
+  it('creates a new ActionsClientLlm instance with the expected traceOptions', async () => {
+    const actionsClient = {
+      get: jest.fn().mockResolvedValue(evaluatorConnector),
+    } as unknown as ActionsClient;
+
+    await getEvaluatorLlm({
+      actionsClient,
+      connectorTimeout,
+      evaluatorConnectorId,
+      experimentConnector,
+      langSmithApiKey: 'test-api-key',
+      logger,
+    });
+
+    expect(ActionsClientLlm).toHaveBeenCalledWith(
+      expect.objectContaining({
+        traceOptions: {
+          projectName: 'evaluators',
+          tracers: expect.any(Array),
+        },
+      })
+    );
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_evaluator_llm/index.ts

@@ -0,0 +1,65 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClient } from '@kbn/actions-plugin/server';
+import type { Connector } from '@kbn/actions-plugin/server/application/connector/types';
+import { Logger } from '@kbn/core/server';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
+import { ActionsClientLlm } from '@kbn/langchain/server';
+import { PublicMethodsOf } from '@kbn/utility-types';
+
+import { getLlmType } from '../../../../../routes/utils';
+
+export const getEvaluatorLlm = async ({
+  actionsClient,
+  connectorTimeout,
+  evaluatorConnectorId,
+  experimentConnector,
+  langSmithApiKey,
+  logger,
+}: {
+  actionsClient: PublicMethodsOf<ActionsClient>;
+  connectorTimeout: number;
+  evaluatorConnectorId: string | undefined;
+  experimentConnector: Connector;
+  langSmithApiKey: string | undefined;
+  logger: Logger;
+}): Promise<ActionsClientLlm> => {
+  const evaluatorConnector =
+    (await actionsClient.get({
+      id: evaluatorConnectorId ?? experimentConnector.id, // fallback to the experiment connector if the evaluator connector is not found:
+      throwIfSystemAction: false,
+    })) ?? experimentConnector;
+
+  const evaluatorLlmType = getLlmType(evaluatorConnector.actionTypeId);
+  const experimentLlmType = getLlmType(experimentConnector.actionTypeId);
+
+  logger.info(
+    `The ${evaluatorConnector.name} (${evaluatorLlmType}) connector will judge output from the ${experimentConnector.name} (${experimentLlmType}) connector`
+  );
+
+  const traceOptions = {
+    projectName: 'evaluators',
+    tracers: [
+      ...getLangSmithTracer({
+        apiKey: langSmithApiKey,
+        projectName: 'evaluators',
+        logger,
+      }),
+    ],
+  };
+
+  return new ActionsClientLlm({
+    actionsClient,
+    connectorId: evaluatorConnector.id,
+    llmType: evaluatorLlmType,
+    logger,
+    temperature: 0, // zero temperature for evaluation
+    timeout: connectorTimeout,
+    traceOptions,
+  });
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_graph_input_overrides/index.test.ts

@@ -0,0 +1,121 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { omit } from 'lodash/fp';
+import type { Example } from 'langsmith/schemas';
+
+import { getGraphInputOverrides } from '.';
+import { exampleWithReplacements } from '../../__mocks__/mock_examples';
+
+const exampleWithAlerts: Example = {
+  ...exampleWithReplacements,
+  outputs: {
+    ...exampleWithReplacements.outputs,
+    anonymizedAlerts: [
+      {
+        metadata: {},
+        pageContent:
+          '@timestamp,2024-10-10T21:01:24.148Z\n' +
+          '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' +
+          'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' +
+          'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff',
+      },
+      {
+        metadata: {},
+        pageContent:
+          '@timestamp,2024-10-10T21:01:24.148Z\n' +
+          '_id,c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b\n' +
+          'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' +
+          'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff',
+      },
+    ],
+  },
+};
+
+const exampleWithNoReplacements: Example = {
+  ...exampleWithReplacements,
+  outputs: {
+    ...omit('replacements', exampleWithReplacements.outputs),
+  },
+};
+
+describe('getGraphInputOverrides', () => {
+  describe('root-level outputs overrides', () => {
+    it('returns the anonymizedAlerts from the root level of the outputs when present', () => {
+      const overrides = getGraphInputOverrides(exampleWithAlerts.outputs);
+
+      expect(overrides.anonymizedAlerts).toEqual(exampleWithAlerts.outputs?.anonymizedAlerts);
+    });
+
+    it('does NOT populate the anonymizedAlerts key when it does NOT exist in the outputs', () => {
+      const overrides = getGraphInputOverrides(exampleWithReplacements.outputs);
+
+      expect(overrides).not.toHaveProperty('anonymizedAlerts');
+    });
+
+    it('returns replacements from the root level of the outputs when present', () => {
+      const overrides = getGraphInputOverrides(exampleWithReplacements.outputs);
+
+      expect(overrides.replacements).toEqual(exampleWithReplacements.outputs?.replacements);
+    });
+
+    it('does NOT populate the replacements key when it does NOT exist in the outputs', () => {
+      const overrides = getGraphInputOverrides(exampleWithNoReplacements.outputs);
+
+      expect(overrides).not.toHaveProperty('replacements');
+    });
+
+    it('removes unknown properties', () => {
+      const withUnknownProperties = {
+        ...exampleWithReplacements,
+        outputs: {
+          ...exampleWithReplacements.outputs,
+          unknownProperty: 'unknown',
+        },
+      };
+
+      const overrides = getGraphInputOverrides(withUnknownProperties.outputs);
+
+      expect(overrides).not.toHaveProperty('unknownProperty');
+    });
+  });
+
+  describe('overrides', () => {
+    it('returns all overrides at the root level', () => {
+      const exampleWithOverrides = {
+        ...exampleWithAlerts,
+        outputs: {
+          ...exampleWithAlerts.outputs,
+          overrides: {
+            attackDiscoveries: [],
+            attackDiscoveryPrompt: 'prompt',
+            anonymizedAlerts: [],
+            combinedGenerations: 'combinedGenerations',
+            combinedRefinements: 'combinedRefinements',
+            errors: ['error'],
+            generationAttempts: 1,
+            generations: ['generation'],
+            hallucinationFailures: 2,
+            maxGenerationAttempts: 3,
+            maxHallucinationFailures: 4,
+            maxRepeatedGenerations: 5,
+            refinements: ['refinement'],
+            refinePrompt: 'refinePrompt',
+            replacements: {},
+            unrefinedResults: [],
+          },
+        },
+      };
+
+      const overrides = getGraphInputOverrides(exampleWithOverrides.outputs);
+
+      expect(overrides).toEqual({
+        ...exampleWithOverrides.outputs?.overrides,
+      });
+    });
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_graph_input_overrides/index.ts

@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { pick } from 'lodash/fp';
+
+import { ExampleInputWithOverrides } from '../../example_input';
+import { GraphState } from '../../../graphs/default_attack_discovery_graph/types';
+
+/**
+ * Parses input from an LangSmith dataset example to get the graph input overrides
+ */
+export const getGraphInputOverrides = (outputs: unknown): Partial<GraphState> => {
+  const validatedInput = ExampleInputWithOverrides.safeParse(outputs).data ?? {}; // safeParse removes unknown properties
+
+  const { overrides } = validatedInput;
+
+  // return all overrides at the root level:
+  return {
+    // pick extracts just the anonymizedAlerts and replacements from the root level of the input,
+    // and only adds the anonymizedAlerts key if it exists in the input
+    ...pick('anonymizedAlerts', validatedInput),
+    ...pick('replacements', validatedInput),
+    ...overrides, // bring all other overrides to the root level
+  };
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.ts

@@ -0,0 +1,122 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClient } from '@kbn/actions-plugin/server';
+import type { Connector } from '@kbn/actions-plugin/server/application/connector/types';
+import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
+import { Logger } from '@kbn/core/server';
+import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
+import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
+import { ActionsClientLlm } from '@kbn/langchain/server';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
+import { asyncForEach } from '@kbn/std';
+import { PublicMethodsOf } from '@kbn/utility-types';
+
+import { DEFAULT_EVAL_ANONYMIZATION_FIELDS } from './constants';
+import { AttackDiscoveryGraphMetadata } from '../../langchain/graphs';
+import { DefaultAttackDiscoveryGraph } from '../graphs/default_attack_discovery_graph';
+import { getLlmType } from '../../../routes/utils';
+import { runEvaluations } from './run_evaluations';
+
+export const evaluateAttackDiscovery = async ({
+  actionsClient,
+  attackDiscoveryGraphs,
+  alertsIndexPattern,
+  anonymizationFields = DEFAULT_EVAL_ANONYMIZATION_FIELDS, // determines which fields are included in the alerts
+  connectors,
+  connectorTimeout,
+  datasetName,
+  esClient,
+  evaluationId,
+  evaluatorConnectorId,
+  langSmithApiKey,
+  langSmithProject,
+  logger,
+  runName,
+  size,
+}: {
+  actionsClient: PublicMethodsOf<ActionsClient>;
+  attackDiscoveryGraphs: AttackDiscoveryGraphMetadata[];
+  alertsIndexPattern: string;
+  anonymizationFields?: AnonymizationFieldResponse[];
+  connectors: Connector[];
+  connectorTimeout: number;
+  datasetName: string;
+  esClient: ElasticsearchClient;
+  evaluationId: string;
+  evaluatorConnectorId: string | undefined;
+  langSmithApiKey: string | undefined;
+  langSmithProject: string | undefined;
+  logger: Logger;
+  runName: string;
+  size: number;
+}): Promise<void> => {
+  await asyncForEach(attackDiscoveryGraphs, async ({ getDefaultAttackDiscoveryGraph }) => {
+    // create a graph for every connector:
+    const graphs: Array<{
+      connector: Connector;
+      graph: DefaultAttackDiscoveryGraph;
+      llmType: string | undefined;
+      name: string;
+      traceOptions: {
+        projectName: string | undefined;
+        tracers: LangChainTracer[];
+      };
+    }> = connectors.map((connector) => {
+      const llmType = getLlmType(connector.actionTypeId);
+
+      const traceOptions = {
+        projectName: langSmithProject,
+        tracers: [
+          ...getLangSmithTracer({
+            apiKey: langSmithApiKey,
+            projectName: langSmithProject,
+            logger,
+          }),
+        ],
+      };
+
+      const llm = new ActionsClientLlm({
+        actionsClient,
+        connectorId: connector.id,
+        llmType,
+        logger,
+        temperature: 0, // zero temperature for attack discovery, because we want structured JSON output
+        timeout: connectorTimeout,
+        traceOptions,
+      });
+
+      const graph = getDefaultAttackDiscoveryGraph({
+        alertsIndexPattern,
+        anonymizationFields,
+        esClient,
+        llm,
+        logger,
+        size,
+      });
+
+      return {
+        connector,
+        graph,
+        llmType,
+        name: `${runName} - ${connector.name} - ${evaluationId} - Attack discovery`,
+        traceOptions,
+      };
+    });
+
+    // run the evaluations for each graph:
+    await runEvaluations({
+      actionsClient,
+      connectorTimeout,
+      evaluatorConnectorId,
+      datasetName,
+      graphs,
+      langSmithApiKey,
+      logger,
+    });
+  });
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/run_evaluations/index.ts

@@ -0,0 +1,113 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClient } from '@kbn/actions-plugin/server';
+import type { Connector } from '@kbn/actions-plugin/server/application/connector/types';
+import { Logger } from '@kbn/core/server';
+import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
+import { asyncForEach } from '@kbn/std';
+import { PublicMethodsOf } from '@kbn/utility-types';
+import { Client } from 'langsmith';
+import { evaluate } from 'langsmith/evaluation';
+
+import { getEvaluatorLlm } from '../helpers/get_evaluator_llm';
+import { getCustomEvaluator } from '../helpers/get_custom_evaluator';
+import { getDefaultPromptTemplate } from '../helpers/get_custom_evaluator/get_default_prompt_template';
+import { getGraphInputOverrides } from '../helpers/get_graph_input_overrides';
+import { DefaultAttackDiscoveryGraph } from '../../graphs/default_attack_discovery_graph';
+import { GraphState } from '../../graphs/default_attack_discovery_graph/types';
+
+/**
+ * Runs an evaluation for each graph so they show up separately (resulting in
+ * each dataset run grouped by connector)
+ */
+export const runEvaluations = async ({
+  actionsClient,
+  connectorTimeout,
+  evaluatorConnectorId,
+  datasetName,
+  graphs,
+  langSmithApiKey,
+  logger,
+}: {
+  actionsClient: PublicMethodsOf<ActionsClient>;
+  connectorTimeout: number;
+  evaluatorConnectorId: string | undefined;
+  datasetName: string;
+  graphs: Array<{
+    connector: Connector;
+    graph: DefaultAttackDiscoveryGraph;
+    llmType: string | undefined;
+    name: string;
+    traceOptions: {
+      projectName: string | undefined;
+      tracers: LangChainTracer[];
+    };
+  }>;
+  langSmithApiKey: string | undefined;
+  logger: Logger;
+}): Promise<void> =>
+  asyncForEach(graphs, async ({ connector, graph, llmType, name, traceOptions }) => {
+    const subject = `connector "${connector.name}" (${llmType}), running experiment "${name}"`;
+
+    try {
+      logger.info(
+        () =>
+          `Evaluating ${subject} with dataset "${datasetName}" and evaluator "${evaluatorConnectorId}"`
+      );
+
+      const predict = async (input: unknown): Promise<GraphState> => {
+        logger.debug(() => `Raw example Input for ${subject}":\n ${input}`);
+
+        // The example `Input` may have overrides for the initial state of the graph:
+        const overrides = getGraphInputOverrides(input);
+
+        return graph.invoke(
+          {
+            ...overrides,
+          },
+          {
+            callbacks: [...(traceOptions.tracers ?? [])],
+            runName: name,
+            tags: ['evaluation', llmType ?? ''],
+          }
+        );
+      };
+
+      const llm = await getEvaluatorLlm({
+        actionsClient,
+        connectorTimeout,
+        evaluatorConnectorId,
+        experimentConnector: connector,
+        langSmithApiKey,
+        logger,
+      });
+
+      const customEvaluator = getCustomEvaluator({
+        criteria: 'correctness',
+        key: 'attack_discovery_correctness',
+        llm,
+        template: getDefaultPromptTemplate(),
+      });
+
+      const evalOutput = await evaluate(predict, {
+        client: new Client({ apiKey: langSmithApiKey }),
+        data: datasetName ?? '',
+        evaluators: [customEvaluator],
+        experimentPrefix: name,
+        maxConcurrency: 5, // prevents rate limiting
+      });
+
+      logger.info(() => `Evaluation complete for ${subject}`);
+
+      logger.debug(
+        () => `Evaluation output for ${subject}:\n ${JSON.stringify(evalOutput, null, 2)}`
+      );
+    } catch (e) {
+      logger.error(`Error evaluating ${subject}: ${e}`);
+    }
+  });

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/constants.ts

@@ -0,0 +1,21 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+// LangGraph metadata
+export const ATTACK_DISCOVERY_GRAPH_RUN_NAME = 'Attack discovery';
+export const ATTACK_DISCOVERY_TAG = 'attack-discovery';
+
+// Limits
+export const DEFAULT_MAX_GENERATION_ATTEMPTS = 10;
+export const DEFAULT_MAX_HALLUCINATION_FAILURES = 5;
+export const DEFAULT_MAX_REPEATED_GENERATIONS = 3;
+
+export const NodeType = {
+  GENERATE_NODE: 'generate',
+  REFINE_NODE: 'refine',
+  RETRIEVE_ANONYMIZED_ALERTS_NODE: 'retrieve_anonymized_alerts',
+} as const;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.test.ts

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getGenerateOrEndDecision } from '.';
+
+describe('getGenerateOrEndDecision', () => {
+  it('returns "end" when hasZeroAlerts is true', () => {
+    const result = getGenerateOrEndDecision(true);
+
+    expect(result).toEqual('end');
+  });
+
+  it('returns "generate" when hasZeroAlerts is false', () => {
+    const result = getGenerateOrEndDecision(false);
+
+    expect(result).toEqual('generate');
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.ts

@@ -0,0 +1,9 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getGenerateOrEndDecision = (hasZeroAlerts: boolean): 'end' | 'generate' =>
+  hasZeroAlerts ? 'end' : 'generate';

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_end/index.test.ts

@@ -0,0 +1,72 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { loggerMock } from '@kbn/logging-mocks';
+
+import { getGenerateOrEndEdge } from '.';
+import type { GraphState } from '../../types';
+
+const logger = loggerMock.create();
+
+const graphState: GraphState = {
+  attackDiscoveries: null,
+  attackDiscoveryPrompt: 'prompt',
+  anonymizedAlerts: [
+    {
+      metadata: {},
+      pageContent:
+        '@timestamp,2024-10-10T21:01:24.148Z\n' +
+        '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' +
+        'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' +
+        'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff',
+    },
+    {
+      metadata: {},
+      pageContent:
+        '@timestamp,2024-10-10T21:01:24.148Z\n' +
+        '_id,c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b\n' +
+        'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' +
+        'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff',
+    },
+  ],
+  combinedGenerations: 'generations',
+  combinedRefinements: 'refinements',
+  errors: [],
+  generationAttempts: 0,
+  generations: [],
+  hallucinationFailures: 0,
+  maxGenerationAttempts: 10,
+  maxHallucinationFailures: 5,
+  maxRepeatedGenerations: 10,
+  refinements: [],
+  refinePrompt: 'refinePrompt',
+  replacements: {},
+  unrefinedResults: null,
+};
+
+describe('getGenerateOrEndEdge', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  it("returns 'end' when there are zero alerts", () => {
+    const state: GraphState = {
+      ...graphState,
+      anonymizedAlerts: [], // <-- zero alerts
+    };
+
+    const edge = getGenerateOrEndEdge(logger);
+    const result = edge(state);
+
+    expect(result).toEqual('end');
+  });
+
+  it("returns 'generate' when there are alerts", () => {
+    const edge = getGenerateOrEndEdge(logger);
+    const result = edge(graphState);
+
+    expect(result).toEqual('generate');
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_end/index.ts

@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Logger } from '@kbn/core/server';
+
+import { getGenerateOrEndDecision } from './helpers/get_generate_or_end_decision';
+import { getHasZeroAlerts } from '../helpers/get_has_zero_alerts';
+import type { GraphState } from '../../types';
+
+export const getGenerateOrEndEdge = (logger?: Logger) => {
+  const edge = (state: GraphState): 'end' | 'generate' => {
+    logger?.debug(() => '---GENERATE OR END---');
+    const { anonymizedAlerts } = state;
+
+    const hasZeroAlerts = getHasZeroAlerts(anonymizedAlerts);
+
+    const decision = getGenerateOrEndDecision(hasZeroAlerts);
+
+    logger?.debug(
+      () => `generatOrEndEdge evaluated the following (derived) state:\n${JSON.stringify(
+        {
+          anonymizedAlerts: anonymizedAlerts.length,
+          hasZeroAlerts,
+        },
+        null,
+        2
+      )}
+\n---GENERATE OR END: ${decision}---`
+    );
+    return decision;
+  };
+
+  return edge;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.test.ts

@@ -0,0 +1,43 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getGenerateOrRefineOrEndDecision } from '.';
+
+describe('getGenerateOrRefineOrEndDecision', () => {
+  it("returns 'end' if getShouldEnd returns true", () => {
+    const result = getGenerateOrRefineOrEndDecision({
+      hasUnrefinedResults: false,
+      hasZeroAlerts: true,
+      maxHallucinationFailuresReached: true,
+      maxRetriesReached: true,
+    });
+
+    expect(result).toEqual('end');
+  });
+
+  it("returns 'refine' if hasUnrefinedResults is true and getShouldEnd returns false", () => {
+    const result = getGenerateOrRefineOrEndDecision({
+      hasUnrefinedResults: true,
+      hasZeroAlerts: false,
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: false,
+    });
+
+    expect(result).toEqual('refine');
+  });
+
+  it("returns 'generate' if hasUnrefinedResults is false and getShouldEnd returns false", () => {
+    const result = getGenerateOrRefineOrEndDecision({
+      hasUnrefinedResults: false,
+      hasZeroAlerts: false,
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: false,
+    });
+
+    expect(result).toEqual('generate');
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.ts

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getShouldEnd } from '../get_should_end';
+
+export const getGenerateOrRefineOrEndDecision = ({
+  hasUnrefinedResults,
+  hasZeroAlerts,
+  maxHallucinationFailuresReached,
+  maxRetriesReached,
+}: {
+  hasUnrefinedResults: boolean;
+  hasZeroAlerts: boolean;
+  maxHallucinationFailuresReached: boolean;
+  maxRetriesReached: boolean;
+}): 'end' | 'generate' | 'refine' => {
+  if (getShouldEnd({ hasZeroAlerts, maxHallucinationFailuresReached, maxRetriesReached })) {
+    return 'end';
+  } else if (hasUnrefinedResults) {
+    return 'refine';
+  } else {
+    return 'generate';
+  }
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.test.ts

@@ -0,0 +1,60 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getShouldEnd } from '.';
+
+describe('getShouldEnd', () => {
+  it('returns true if hasZeroAlerts is true', () => {
+    const result = getShouldEnd({
+      hasZeroAlerts: true, // <-- true
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: false,
+    });
+
+    expect(result).toBe(true);
+  });
+
+  it('returns true if maxHallucinationFailuresReached is true', () => {
+    const result = getShouldEnd({
+      hasZeroAlerts: false,
+      maxHallucinationFailuresReached: true, // <-- true
+      maxRetriesReached: false,
+    });
+
+    expect(result).toBe(true);
+  });
+
+  it('returns true if maxRetriesReached is true', () => {
+    const result = getShouldEnd({
+      hasZeroAlerts: false,
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: true, // <-- true
+    });
+
+    expect(result).toBe(true);
+  });
+
+  it('returns false if all conditions are false', () => {
+    const result = getShouldEnd({
+      hasZeroAlerts: false,
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: false,
+    });
+
+    expect(result).toBe(false);
+  });
+
+  it('returns true if all conditions are true', () => {
+    const result = getShouldEnd({
+      hasZeroAlerts: true,
+      maxHallucinationFailuresReached: true,
+      maxRetriesReached: true,
+    });
+
+    expect(result).toBe(true);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.ts

@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getShouldEnd = ({
+  hasZeroAlerts,
+  maxHallucinationFailuresReached,
+  maxRetriesReached,
+}: {
+  hasZeroAlerts: boolean;
+  maxHallucinationFailuresReached: boolean;
+  maxRetriesReached: boolean;
+}): boolean => hasZeroAlerts || maxRetriesReached || maxHallucinationFailuresReached;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_refine_or_end/index.test.ts

@@ -0,0 +1,118 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { loggerMock } from '@kbn/logging-mocks';
+
+import { getGenerateOrRefineOrEndEdge } from '.';
+import type { GraphState } from '../../types';
+
+const logger = loggerMock.create();
+
+const graphState: GraphState = {
+  attackDiscoveries: null,
+  attackDiscoveryPrompt: 'prompt',
+  anonymizedAlerts: [
+    {
+      metadata: {},
+      pageContent:
+        '@timestamp,2024-10-10T21:01:24.148Z\n' +
+        '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' +
+        'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' +
+        'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff',
+    },
+    {
+      metadata: {},
+      pageContent:
+        '@timestamp,2024-10-10T21:01:24.148Z\n' +
+        '_id,c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b\n' +
+        'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' +
+        'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff',
+    },
+  ],
+  combinedGenerations: '',
+  combinedRefinements: '',
+  errors: [],
+  generationAttempts: 0,
+  generations: [],
+  hallucinationFailures: 0,
+  maxGenerationAttempts: 10,
+  maxHallucinationFailures: 5,
+  maxRepeatedGenerations: 3,
+  refinements: [],
+  refinePrompt: 'refinePrompt',
+  replacements: {},
+  unrefinedResults: null,
+};
+
+describe('getGenerateOrRefineOrEndEdge', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  it('returns "end" when there are zero alerts', () => {
+    const withZeroAlerts: GraphState = {
+      ...graphState,
+      anonymizedAlerts: [], // <-- zero alerts
+    };
+
+    const edge = getGenerateOrRefineOrEndEdge(logger);
+    const result = edge(withZeroAlerts);
+
+    expect(result).toEqual('end');
+  });
+
+  it('returns "end" when max hallucination failures are reached', () => {
+    const withMaxHallucinationFailures: GraphState = {
+      ...graphState,
+      hallucinationFailures: 5,
+    };
+
+    const edge = getGenerateOrRefineOrEndEdge(logger);
+    const result = edge(withMaxHallucinationFailures);
+
+    expect(result).toEqual('end');
+  });
+
+  it('returns "end" when max retries are reached', () => {
+    const withMaxRetries: GraphState = {
+      ...graphState,
+      generationAttempts: 10,
+    };
+
+    const edge = getGenerateOrRefineOrEndEdge(logger);
+    const result = edge(withMaxRetries);
+
+    expect(result).toEqual('end');
+  });
+
+  it('returns refine when there are unrefined results', () => {
+    const withUnrefinedResults: GraphState = {
+      ...graphState,
+      unrefinedResults: [
+        {
+          alertIds: [],
+          id: 'test-id',
+          detailsMarkdown: 'test-details',
+          entitySummaryMarkdown: 'test-summary',
+          summaryMarkdown: 'test-summary',
+          title: 'test-title',
+          timestamp: '2024-10-10T21:01:24.148Z',
+        },
+      ],
+    };
+
+    const edge = getGenerateOrRefineOrEndEdge(logger);
+    const result = edge(withUnrefinedResults);
+
+    expect(result).toEqual('refine');
+  });
+
+  it('return generate when there are no unrefined results', () => {
+    const edge = getGenerateOrRefineOrEndEdge(logger);
+    const result = edge(graphState);
+
+    expect(result).toEqual('generate');
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_refine_or_end/index.ts

@@ -0,0 +1,66 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Logger } from '@kbn/core/server';
+
+import { getGenerateOrRefineOrEndDecision } from './helpers/get_generate_or_refine_or_end_decision';
+import { getHasResults } from '../helpers/get_has_results';
+import { getHasZeroAlerts } from '../helpers/get_has_zero_alerts';
+import { getMaxHallucinationFailuresReached } from '../../helpers/get_max_hallucination_failures_reached';
+import { getMaxRetriesReached } from '../../helpers/get_max_retries_reached';
+import type { GraphState } from '../../types';
+
+export const getGenerateOrRefineOrEndEdge = (logger?: Logger) => {
+  const edge = (state: GraphState): 'end' | 'generate' | 'refine' => {
+    logger?.debug(() => '---GENERATE OR REFINE OR END---');
+    const {
+      anonymizedAlerts,
+      generationAttempts,
+      hallucinationFailures,
+      maxGenerationAttempts,
+      maxHallucinationFailures,
+      unrefinedResults,
+    } = state;
+
+    const hasZeroAlerts = getHasZeroAlerts(anonymizedAlerts);
+    const hasUnrefinedResults = getHasResults(unrefinedResults);
+    const maxRetriesReached = getMaxRetriesReached({ generationAttempts, maxGenerationAttempts });
+    const maxHallucinationFailuresReached = getMaxHallucinationFailuresReached({
+      hallucinationFailures,
+      maxHallucinationFailures,
+    });
+
+    const decision = getGenerateOrRefineOrEndDecision({
+      hasUnrefinedResults,
+      hasZeroAlerts,
+      maxHallucinationFailuresReached,
+      maxRetriesReached,
+    });
+
+    logger?.debug(
+      () =>
+        `generatOrRefineOrEndEdge evaluated the following (derived) state:\n${JSON.stringify(
+          {
+            anonymizedAlerts: anonymizedAlerts.length,
+            generationAttempts,
+            hallucinationFailures,
+            hasUnrefinedResults,
+            hasZeroAlerts,
+            maxHallucinationFailuresReached,
+            maxRetriesReached,
+            unrefinedResults: unrefinedResults?.length ?? 0,
+          },
+          null,
+          2
+        )}
+        \n---GENERATE OR REFINE OR END: ${decision}---`
+    );
+    return decision;
+  };
+
+  return edge;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_results/index.ts

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AttackDiscovery } from '@kbn/elastic-assistant-common';
+
+export const getHasResults = (attackDiscoveries: AttackDiscovery[] | null): boolean =>
+  attackDiscoveries !== null;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_zero_alerts/index.ts

@@ -0,0 +1,12 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Document } from '@langchain/core/documents';
+import { isEmpty } from 'lodash/fp';
+
+export const getHasZeroAlerts = (anonymizedAlerts: Document[]): boolean =>
+  isEmpty(anonymizedAlerts);

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.ts

@@ -0,0 +1,25 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getShouldEnd } from '../get_should_end';
+
+export const getRefineOrEndDecision = ({
+  hasFinalResults,
+  maxHallucinationFailuresReached,
+  maxRetriesReached,
+}: {
+  hasFinalResults: boolean;
+  maxHallucinationFailuresReached: boolean;
+  maxRetriesReached: boolean;
+}): 'refine' | 'end' =>
+  getShouldEnd({
+    hasFinalResults,
+    maxHallucinationFailuresReached,
+    maxRetriesReached,
+  })
+    ? 'end'
+    : 'refine';

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_should_end/index.ts

@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getShouldEnd = ({
+  hasFinalResults,
+  maxHallucinationFailuresReached,
+  maxRetriesReached,
+}: {
+  hasFinalResults: boolean;
+  maxHallucinationFailuresReached: boolean;
+  maxRetriesReached: boolean;
+}): boolean => hasFinalResults || maxRetriesReached || maxHallucinationFailuresReached;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/index.ts

@@ -0,0 +1,61 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Logger } from '@kbn/core/server';
+
+import { getRefineOrEndDecision } from './helpers/get_refine_or_end_decision';
+import { getHasResults } from '../helpers/get_has_results';
+import { getMaxHallucinationFailuresReached } from '../../helpers/get_max_hallucination_failures_reached';
+import { getMaxRetriesReached } from '../../helpers/get_max_retries_reached';
+import type { GraphState } from '../../types';
+
+export const getRefineOrEndEdge = (logger?: Logger) => {
+  const edge = (state: GraphState): 'end' | 'refine' => {
+    logger?.debug(() => '---REFINE OR END---');
+    const {
+      attackDiscoveries,
+      generationAttempts,
+      hallucinationFailures,
+      maxGenerationAttempts,
+      maxHallucinationFailures,
+    } = state;
+
+    const hasFinalResults = getHasResults(attackDiscoveries);
+    const maxRetriesReached = getMaxRetriesReached({ generationAttempts, maxGenerationAttempts });
+    const maxHallucinationFailuresReached = getMaxHallucinationFailuresReached({
+      hallucinationFailures,
+      maxHallucinationFailures,
+    });
+
+    const decision = getRefineOrEndDecision({
+      hasFinalResults,
+      maxHallucinationFailuresReached,
+      maxRetriesReached,
+    });
+
+    logger?.debug(
+      () =>
+        `refineOrEndEdge evaluated the following (derived) state:\n${JSON.stringify(
+          {
+            attackDiscoveries: attackDiscoveries?.length ?? 0,
+            generationAttempts,
+            hallucinationFailures,
+            hasFinalResults,
+            maxHallucinationFailuresReached,
+            maxRetriesReached,
+          },
+          null,
+          2
+        )}
+        \n---REFINE OR END: ${decision}---`
+    );
+
+    return decision;
+  };
+
+  return edge;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/get_retrieve_or_generate/index.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Document } from '@langchain/core/documents';
+
+export const getRetrieveOrGenerate = (
+  anonymizedAlerts: Document[]
+): 'retrieve_anonymized_alerts' | 'generate' =>
+  anonymizedAlerts.length === 0 ? 'retrieve_anonymized_alerts' : 'generate';

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/index.ts

@@ -0,0 +1,36 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Logger } from '@kbn/core/server';
+
+import { getRetrieveOrGenerate } from './get_retrieve_or_generate';
+import type { GraphState } from '../../types';
+
+export const getRetrieveAnonymizedAlertsOrGenerateEdge = (logger?: Logger) => {
+  const edge = (state: GraphState): 'retrieve_anonymized_alerts' | 'generate' => {
+    logger?.debug(() => '---RETRIEVE ANONYMIZED ALERTS OR GENERATE---');
+    const { anonymizedAlerts } = state;
+
+    const decision = getRetrieveOrGenerate(anonymizedAlerts);
+
+    logger?.debug(
+      () =>
+        `retrieveAnonymizedAlertsOrGenerateEdge evaluated the following (derived) state:\n${JSON.stringify(
+          {
+            anonymizedAlerts: anonymizedAlerts.length,
+          },
+          null,
+          2
+        )}
+        \n---RETRIEVE ANONYMIZED ALERTS OR GENERATE: ${decision}---`
+    );
+
+    return decision;
+  };
+
+  return edge;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_hallucination_failures_reached/index.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getMaxHallucinationFailuresReached = ({
+  hallucinationFailures,
+  maxHallucinationFailures,
+}: {
+  hallucinationFailures: number;
+  maxHallucinationFailures: number;
+}): boolean => hallucinationFailures >= maxHallucinationFailures;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_retries_reached/index.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getMaxRetriesReached = ({
+  generationAttempts,
+  maxGenerationAttempts,
+}: {
+  generationAttempts: number;
+  maxGenerationAttempts: number;
+}): boolean => generationAttempts >= maxGenerationAttempts;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/index.ts

@@ -0,0 +1,122 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ElasticsearchClient, Logger } from '@kbn/core/server';
+import { Replacements } from '@kbn/elastic-assistant-common';
+import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
+import type { ActionsClientLlm } from '@kbn/langchain/server';
+import type { CompiledStateGraph } from '@langchain/langgraph';
+import { END, START, StateGraph } from '@langchain/langgraph';
+
+import { NodeType } from './constants';
+import { getGenerateOrEndEdge } from './edges/generate_or_end';
+import { getGenerateOrRefineOrEndEdge } from './edges/generate_or_refine_or_end';
+import { getRefineOrEndEdge } from './edges/refine_or_end';
+import { getRetrieveAnonymizedAlertsOrGenerateEdge } from './edges/retrieve_anonymized_alerts_or_generate';
+import { getDefaultGraphState } from './state';
+import { getGenerateNode } from './nodes/generate';
+import { getRefineNode } from './nodes/refine';
+import { getRetrieveAnonymizedAlertsNode } from './nodes/retriever';
+import type { GraphState } from './types';
+
+export interface GetDefaultAttackDiscoveryGraphParams {
+  alertsIndexPattern?: string;
+  anonymizationFields: AnonymizationFieldResponse[];
+  esClient: ElasticsearchClient;
+  llm: ActionsClientLlm;
+  logger?: Logger;
+  onNewReplacements?: (replacements: Replacements) => void;
+  replacements?: Replacements;
+  size: number;
+}
+
+export type DefaultAttackDiscoveryGraph = ReturnType<typeof getDefaultAttackDiscoveryGraph>;
+
+/**
+ * This function returns a compiled state graph that represents the default
+ * Attack discovery graph.
+ *
+ * Refer to the following diagram for this graph:
+ * x-pack/plugins/elastic_assistant/docs/img/default_attack_discovery_graph.png
+ */
+export const getDefaultAttackDiscoveryGraph = ({
+  alertsIndexPattern,
+  anonymizationFields,
+  esClient,
+  llm,
+  logger,
+  onNewReplacements,
+  replacements,
+  size,
+}: GetDefaultAttackDiscoveryGraphParams): CompiledStateGraph<
+  GraphState,
+  Partial<GraphState>,
+  'generate' | 'refine' | 'retrieve_anonymized_alerts' | '__start__'
+> => {
+  try {
+    const graphState = getDefaultGraphState();
+
+    // get nodes:
+    const retrieveAnonymizedAlertsNode = getRetrieveAnonymizedAlertsNode({
+      alertsIndexPattern,
+      anonymizationFields,
+      esClient,
+      logger,
+      onNewReplacements,
+      replacements,
+      size,
+    });
+
+    const generateNode = getGenerateNode({
+      llm,
+      logger,
+    });
+
+    const refineNode = getRefineNode({
+      llm,
+      logger,
+    });
+
+    // get edges:
+    const generateOrEndEdge = getGenerateOrEndEdge(logger);
+
+    const generatOrRefineOrEndEdge = getGenerateOrRefineOrEndEdge(logger);
+
+    const refineOrEndEdge = getRefineOrEndEdge(logger);
+
+    const retrieveAnonymizedAlertsOrGenerateEdge =
+      getRetrieveAnonymizedAlertsOrGenerateEdge(logger);
+
+    // create the graph:
+    const graph = new StateGraph<GraphState>({ channels: graphState })
+      .addNode(NodeType.RETRIEVE_ANONYMIZED_ALERTS_NODE, retrieveAnonymizedAlertsNode)
+      .addNode(NodeType.GENERATE_NODE, generateNode)
+      .addNode(NodeType.REFINE_NODE, refineNode)
+      .addConditionalEdges(START, retrieveAnonymizedAlertsOrGenerateEdge, {
+        generate: NodeType.GENERATE_NODE,
+        retrieve_anonymized_alerts: NodeType.RETRIEVE_ANONYMIZED_ALERTS_NODE,
+      })
+      .addConditionalEdges(NodeType.RETRIEVE_ANONYMIZED_ALERTS_NODE, generateOrEndEdge, {
+        end: END,
+        generate: NodeType.GENERATE_NODE,
+      })
+      .addConditionalEdges(NodeType.GENERATE_NODE, generatOrRefineOrEndEdge, {
+        end: END,
+        generate: NodeType.GENERATE_NODE,
+        refine: NodeType.REFINE_NODE,
+      })
+      .addConditionalEdges(NodeType.REFINE_NODE, refineOrEndEdge, {
+        end: END,
+        refine: NodeType.REFINE_NODE,
+      });
+
+    // compile the graph:
+    return graph.compile();
+  } catch (e) {
+    throw new Error(`Unable to compile AttackDiscoveryGraph\n${e}`);
+  }
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/mock/mock_empty_open_and_acknowledged_alerts_qery_results.ts

@@ -0,0 +1,25 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const mockEmptyOpenAndAcknowledgedAlertsQueryResults = {
+  took: 0,
+  timed_out: false,
+  _shards: {
+    total: 1,
+    successful: 1,
+    skipped: 0,
+    failed: 0,
+  },
+  hits: {
+    total: {
+      value: 0,
+      relation: 'eq',
+    },
+    max_score: null,
+    hits: [],
+  },
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/discard_previous_generations/index.ts

@@ -0,0 +1,30 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { GraphState } from '../../../../types';
+
+export const discardPreviousGenerations = ({
+  generationAttempts,
+  hallucinationFailures,
+  isHallucinationDetected,
+  state,
+}: {
+  generationAttempts: number;
+  hallucinationFailures: number;
+  isHallucinationDetected: boolean;
+  state: GraphState;
+}): GraphState => {
+  return {
+    ...state,
+    combinedGenerations: '', // <-- reset the combined generations
+    generationAttempts: generationAttempts + 1,
+    generations: [], // <-- reset the generations
+    hallucinationFailures: isHallucinationDetected
+      ? hallucinationFailures + 1
+      : hallucinationFailures,
+  };
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_alerts_context_prompt/index.test.ts

@@ -4,15 +4,17 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import { getAttackDiscoveryPrompt } from './get_attack_discovery_prompt';
 
-describe('getAttackDiscoveryPrompt', () => {
-  it('should generate the correct attack discovery prompt', () => {
+import { getAlertsContextPrompt } from '.';
+import { getDefaultAttackDiscoveryPrompt } from '../../../helpers/get_default_attack_discovery_prompt';
+
+describe('getAlertsContextPrompt', () => {
+  it('generates the correct prompt', () => {
     const anonymizedAlerts = ['Alert 1', 'Alert 2', 'Alert 3'];
 
-    const expected = `You are a cyber security analyst tasked with analyzing security events from Elastic Security to identify and report on potential cyber attacks or progressions. Your report should focus on high-risk incidents that could severely impact the organization, rather than isolated alerts. Present your findings in a way that can be easily understood by anyone, regardless of their technical expertise, as if you were briefing the CISO. Break down your response into sections based on timing, hosts, and users involved. When correlating alerts, use kibana.alert.original_time when it's available, otherwise use @timestamp. Include appropriate context about the affected hosts and users. Describe how the attack progression might have occurred and, if feasible, attribute it to known threat groups. Prioritize high and critical alerts, but include lower-severity alerts if desired. In the description field, provide as much detail as possible, in a bulleted list explaining any attack progressions. Accuracy is of utmost importance. Escape backslashes to respect JSON validation. New lines must always be escaped with double backslashes, i.e. \\\\n to ensure valid JSON. Only return JSON output, as described above. Do not add any additional text to describe your output.
+    const expected = `You are a cyber security analyst tasked with analyzing security events from Elastic Security to identify and report on potential cyber attacks or progressions. Your report should focus on high-risk incidents that could severely impact the organization, rather than isolated alerts. Present your findings in a way that can be easily understood by anyone, regardless of their technical expertise, as if you were briefing the CISO. Break down your response into sections based on timing, hosts, and users involved. When correlating alerts, use kibana.alert.original_time when it's available, otherwise use @timestamp. Include appropriate context about the affected hosts and users. Describe how the attack progression might have occurred and, if feasible, attribute it to known threat groups. Prioritize high and critical alerts, but include lower-severity alerts if desired. In the description field, provide as much detail as possible, in a bulleted list explaining any attack progressions. Accuracy is of utmost importance. You MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds).
 
-Use context from the following open and acknowledged alerts to provide insights:
+Use context from the following alerts to provide insights:
 
 """
 Alert 1
@@ -23,7 +25,10 @@ Alert 3
 """
 `;
 
-    const prompt = getAttackDiscoveryPrompt({ anonymizedAlerts });
+    const prompt = getAlertsContextPrompt({
+      anonymizedAlerts,
+      attackDiscoveryPrompt: getDefaultAttackDiscoveryPrompt(),
+    });
 
     expect(prompt).toEqual(expected);
   });

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_alerts_context_prompt/index.ts

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+// NOTE: we ask the LLM to `provide insights`. We do NOT use the feature name, `AttackDiscovery`, in the prompt.
+export const getAlertsContextPrompt = ({
+  anonymizedAlerts,
+  attackDiscoveryPrompt,
+}: {
+  anonymizedAlerts: string[];
+  attackDiscoveryPrompt: string;
+}) => `${attackDiscoveryPrompt}
+
+Use context from the following alerts to provide insights:
+
+"""
+${anonymizedAlerts.join('\n\n')}
+"""
+`;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_anonymized_alerts_from_state/index.ts

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { GraphState } from '../../../../types';
+
+export const getAnonymizedAlertsFromState = (state: GraphState): string[] =>
+  state.anonymizedAlerts.map((doc) => doc.pageContent);

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_use_unrefined_results/index.ts

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AttackDiscovery } from '@kbn/elastic-assistant-common';
+
+import { getMaxRetriesReached } from '../../../../helpers/get_max_retries_reached';
+
+export const getUseUnrefinedResults = ({
+  generationAttempts,
+  maxGenerationAttempts,
+  unrefinedResults,
+}: {
+  generationAttempts: number;
+  maxGenerationAttempts: number;
+  unrefinedResults: AttackDiscovery[] | null;
+}): boolean => {
+  const nextAttemptWouldExcedLimit = getMaxRetriesReached({
+    generationAttempts: generationAttempts + 1, // + 1, because we just used an attempt
+    maxGenerationAttempts,
+  });
+
+  return nextAttemptWouldExcedLimit && unrefinedResults != null && unrefinedResults.length > 0;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.ts

@@ -0,0 +1,154 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClientLlm } from '@kbn/langchain/server';
+import type { Logger } from '@kbn/core/server';
+
+import { discardPreviousGenerations } from './helpers/discard_previous_generations';
+import { extractJson } from '../helpers/extract_json';
+import { getAnonymizedAlertsFromState } from './helpers/get_anonymized_alerts_from_state';
+import { getChainWithFormatInstructions } from '../helpers/get_chain_with_format_instructions';
+import { getCombined } from '../helpers/get_combined';
+import { getCombinedAttackDiscoveryPrompt } from '../helpers/get_combined_attack_discovery_prompt';
+import { generationsAreRepeating } from '../helpers/generations_are_repeating';
+import { getUseUnrefinedResults } from './helpers/get_use_unrefined_results';
+import { parseCombinedOrThrow } from '../helpers/parse_combined_or_throw';
+import { responseIsHallucinated } from '../helpers/response_is_hallucinated';
+import type { GraphState } from '../../types';
+
+export const getGenerateNode = ({
+  llm,
+  logger,
+}: {
+  llm: ActionsClientLlm;
+  logger?: Logger;
+}): ((state: GraphState) => Promise<GraphState>) => {
+  const generate = async (state: GraphState): Promise<GraphState> => {
+    logger?.debug(() => `---GENERATE---`);
+
+    const anonymizedAlerts: string[] = getAnonymizedAlertsFromState(state);
+
+    const {
+      attackDiscoveryPrompt,
+      combinedGenerations,
+      generationAttempts,
+      generations,
+      hallucinationFailures,
+      maxGenerationAttempts,
+      maxRepeatedGenerations,
+    } = state;
+
+    let combinedResponse = ''; // mutable, because it must be accessed in the catch block
+    let partialResponse = ''; // mutable, because it must be accessed in the catch block
+
+    try {
+      const query = getCombinedAttackDiscoveryPrompt({
+        anonymizedAlerts,
+        attackDiscoveryPrompt,
+        combinedMaybePartialResults: combinedGenerations,
+      });
+
+      const { chain, formatInstructions, llmType } = getChainWithFormatInstructions(llm);
+
+      logger?.debug(
+        () => `generate node is invoking the chain (${llmType}), attempt ${generationAttempts}`
+      );
+
+      const rawResponse = (await chain.invoke({
+        format_instructions: formatInstructions,
+        query,
+      })) as unknown as string;
+
+      // LOCAL MUTATION:
+      partialResponse = extractJson(rawResponse); // remove the surrounding ```json```
+
+      // if the response is hallucinated, discard previous generations and start over:
+      if (responseIsHallucinated(partialResponse)) {
+        logger?.debug(
+          () =>
+            `generate node detected a hallucination (${llmType}), on attempt ${generationAttempts}; discarding the accumulated generations and starting over`
+        );
+
+        return discardPreviousGenerations({
+          generationAttempts,
+          hallucinationFailures,
+          isHallucinationDetected: true,
+          state,
+        });
+      }
+
+      // if the generations are repeating, discard previous generations and start over:
+      if (
+        generationsAreRepeating({
+          currentGeneration: partialResponse,
+          previousGenerations: generations,
+          sampleLastNGenerations: maxRepeatedGenerations,
+        })
+      ) {
+        logger?.debug(
+          () =>
+            `generate node detected (${llmType}), detected ${maxRepeatedGenerations} repeated generations on attempt ${generationAttempts}; discarding the accumulated results and starting over`
+        );
+
+        // discard the accumulated results and start over:
+        return discardPreviousGenerations({
+          generationAttempts,
+          hallucinationFailures,
+          isHallucinationDetected: false,
+          state,
+        });
+      }
+
+      // LOCAL MUTATION:
+      combinedResponse = getCombined({ combinedGenerations, partialResponse }); // combine the new response with the previous ones
+
+      const unrefinedResults = parseCombinedOrThrow({
+        combinedResponse,
+        generationAttempts,
+        llmType,
+        logger,
+        nodeName: 'generate',
+      });
+
+      // use the unrefined results if we already reached the max number of retries:
+      const useUnrefinedResults = getUseUnrefinedResults({
+        generationAttempts,
+        maxGenerationAttempts,
+        unrefinedResults,
+      });
+
+      if (useUnrefinedResults) {
+        logger?.debug(
+          () =>
+            `generate node is using unrefined results response (${llm._llmType()}) from attempt ${generationAttempts}, because all attempts have been used`
+        );
+      }
+
+      return {
+        ...state,
+        attackDiscoveries: useUnrefinedResults ? unrefinedResults : null, // optionally skip the refinement step by returning the final answer
+        combinedGenerations: combinedResponse,
+        generationAttempts: generationAttempts + 1,
+        generations: [...generations, partialResponse],
+        unrefinedResults,
+      };
+    } catch (error) {
+      const parsingError = `generate node is unable to parse (${llm._llmType()}) response from attempt ${generationAttempts}; (this may be an incomplete response from the model): ${error}`;
+      logger?.debug(() => parsingError); // logged at debug level because the error is expected when the model returns an incomplete response
+
+      return {
+        ...state,
+        combinedGenerations: combinedResponse,
+        errors: [...state.errors, parsingError],
+        generationAttempts: generationAttempts + 1,
+        generations: [...generations, partialResponse],
+      };
+    }
+  };
+
+  return generate;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/schema/index.ts

@@ -0,0 +1,84 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { z } from '@kbn/zod';
+
+export const SYNTAX = '{{ field.name fieldValue1 fieldValue2 fieldValueN }}';
+const GOOD_SYNTAX_EXAMPLES =
+  'Examples of CORRECT syntax (includes field names and values): {{ host.name hostNameValue }} {{ user.name userNameValue }} {{ source.ip sourceIpValue }}';
+
+const BAD_SYNTAX_EXAMPLES =
+  'Examples of INCORRECT syntax (bad, because the field names are not included): {{ hostNameValue }} {{ userNameValue }} {{ sourceIpValue }}';
+
+const RECONNAISSANCE = 'Reconnaissance';
+const INITIAL_ACCESS = 'Initial Access';
+const EXECUTION = 'Execution';
+const PERSISTENCE = 'Persistence';
+const PRIVILEGE_ESCALATION = 'Privilege Escalation';
+const DISCOVERY = 'Discovery';
+const LATERAL_MOVEMENT = 'Lateral Movement';
+const COMMAND_AND_CONTROL = 'Command and Control';
+const EXFILTRATION = 'Exfiltration';
+
+const MITRE_ATTACK_TACTICS = [
+  RECONNAISSANCE,
+  INITIAL_ACCESS,
+  EXECUTION,
+  PERSISTENCE,
+  PRIVILEGE_ESCALATION,
+  DISCOVERY,
+  LATERAL_MOVEMENT,
+  COMMAND_AND_CONTROL,
+  EXFILTRATION,
+] as const;
+
+export const AttackDiscoveriesGenerationSchema = z.object({
+  insights: z
+    .array(
+      z.object({
+        alertIds: z.string().array().describe(`The alert IDs that the insight is based on.`),
+        detailsMarkdown: z
+          .string()
+          .describe(
+            `A detailed insight with markdown, where each markdown bullet contains a description of what happened that reads like a story of the attack as it played out and always uses special ${SYNTAX} syntax for field names and values from the source data. ${GOOD_SYNTAX_EXAMPLES} ${BAD_SYNTAX_EXAMPLES}`
+          ),
+        entitySummaryMarkdown: z
+          .string()
+          .optional()
+          .describe(
+            `A short (no more than a sentence) summary of the insight featuring only the host.name and user.name fields (when they are applicable), using the same ${SYNTAX} syntax`
+          ),
+        mitreAttackTactics: z
+          .string()
+          .array()
+          .optional()
+          .describe(
+            `An array of MITRE ATT&CK tactic for the insight, using one of the following values: ${MITRE_ATTACK_TACTICS.join(
+              ','
+            )}`
+          ),
+        summaryMarkdown: z
+          .string()
+          .describe(`A markdown summary of insight, using the same ${SYNTAX} syntax`),
+        title: z
+          .string()
+          .describe(
+            'A short, no more than 7 words, title for the insight, NOT formatted with special syntax or markdown. This must be as brief as possible.'
+          ),
+      })
+    )
+    .describe(
+      `Insights with markdown that always uses special ${SYNTAX} syntax for field names and values from the source data. ${GOOD_SYNTAX_EXAMPLES} ${BAD_SYNTAX_EXAMPLES}`
+    ),
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.ts

@@ -0,0 +1,20 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const addTrailingBackticksIfNecessary = (text: string): string => {
+  const leadingJSONpattern = /^\w*```json(.*?)/s;
+  const trailingBackticksPattern = /(.*?)```\w*$/s;
+
+  const hasLeadingJSONWrapper = leadingJSONpattern.test(text);
+  const hasTrailingBackticks = trailingBackticksPattern.test(text);
+
+  if (hasLeadingJSONWrapper && !hasTrailingBackticks) {
+    return `${text}\n\`\`\``;
+  }
+
+  return text;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.ts

@@ -0,0 +1,17 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const extractJson = (input: string): string => {
+  const regex = /```json\s*([\s\S]*?)(?:\s*```|$)/;
+  const match = input.match(regex);
+
+  if (match && match[1]) {
+    return match[1].trim();
+  }
+
+  return input;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/generations_are_repeating/index.test.tsx

@@ -0,0 +1,90 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { generationsAreRepeating } from '.';
+
+describe('getIsGenerationRepeating', () => {
+  it('returns true when all previous generations are the same as the current generation', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: ['gen1', 'gen1', 'gen1'], // <-- all the same, length 3
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(true);
+  });
+
+  it('returns false when some of the previous generations are NOT the same as the current generation', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: ['gen1', 'gen2', 'gen1'], // <-- some are different, length 3
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(false);
+  });
+
+  it('returns true when all *sampled* generations are the same as the current generation, and there are older samples past the last N', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: [
+        'gen2', // <-- older sample will be ignored
+        'gen1',
+        'gen1',
+        'gen1',
+      ],
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(true);
+  });
+
+  it('returns false when some of the *sampled* generations are NOT the same as the current generation, and there are additional samples past the last N', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: [
+        'gen1', // <-- older sample will be ignored
+        'gen1',
+        'gen1',
+        'gen2',
+      ],
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(false);
+  });
+
+  it('returns false when sampling fewer generations than sampleLastNGenerations, and all are the same as the current generation', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: ['gen1', 'gen1'], // <-- same, but only 2 generations
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(false);
+  });
+
+  it('returns false when sampling fewer generations than sampleLastNGenerations, and some are different from the current generation', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: ['gen1', 'gen2'], // <-- different, but only 2 generations
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(false);
+  });
+
+  it('returns false when there are no previous generations to sample', () => {
+    const result = generationsAreRepeating({
+      currentGeneration: 'gen1',
+      previousGenerations: [],
+      sampleLastNGenerations: 3,
+    });
+
+    expect(result).toBe(false);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/generations_are_repeating/index.tsx

@@ -0,0 +1,25 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+/** Returns true if the last n generations are repeating the same output */
+export const generationsAreRepeating = ({
+  currentGeneration,
+  previousGenerations,
+  sampleLastNGenerations,
+}: {
+  currentGeneration: string;
+  previousGenerations: string[];
+  sampleLastNGenerations: number;
+}): boolean => {
+  const generationsToSample = previousGenerations.slice(-sampleLastNGenerations);
+
+  if (generationsToSample.length < sampleLastNGenerations) {
+    return false; // Not enough generations to sample
+  }
+
+  return generationsToSample.every((generation) => generation === currentGeneration);
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_chain_with_format_instructions/index.ts

@@ -0,0 +1,34 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ActionsClientLlm } from '@kbn/langchain/server';
+import { ChatPromptTemplate } from '@langchain/core/prompts';
+import { Runnable } from '@langchain/core/runnables';
+
+import { getOutputParser } from '../get_output_parser';
+
+interface GetChainWithFormatInstructions {
+  chain: Runnable;
+  formatInstructions: string;
+  llmType: string;
+}
+
+export const getChainWithFormatInstructions = (
+  llm: ActionsClientLlm
+): GetChainWithFormatInstructions => {
+  const outputParser = getOutputParser();
+  const formatInstructions = outputParser.getFormatInstructions();
+
+  const prompt = ChatPromptTemplate.fromTemplate(
+    `Answer the user's question as best you can:\n{format_instructions}\n{query}`
+  );
+
+  const chain = prompt.pipe(llm);
+  const llmType = llm._llmType();
+
+  return { chain, formatInstructions, llmType };
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined/index.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getCombined = ({
+  combinedGenerations,
+  partialResponse,
+}: {
+  combinedGenerations: string;
+  partialResponse: string;
+}): string => `${combinedGenerations}${partialResponse}`;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined_attack_discovery_prompt/index.ts

@@ -0,0 +1,43 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { isEmpty } from 'lodash/fp';
+
+import { getAlertsContextPrompt } from '../../generate/helpers/get_alerts_context_prompt';
+import { getContinuePrompt } from '../get_continue_prompt';
+
+/**
+ * Returns the the initial query, or the initial query combined with a
+ * continuation prompt and partial results
+ */
+export const getCombinedAttackDiscoveryPrompt = ({
+  anonymizedAlerts,
+  attackDiscoveryPrompt,
+  combinedMaybePartialResults,
+}: {
+  anonymizedAlerts: string[];
+  attackDiscoveryPrompt: string;
+  /** combined results that may contain incomplete JSON */
+  combinedMaybePartialResults: string;
+}): string => {
+  const alertsContextPrompt = getAlertsContextPrompt({
+    anonymizedAlerts,
+    attackDiscoveryPrompt,
+  });
+
+  return isEmpty(combinedMaybePartialResults)
+    ? alertsContextPrompt // no partial results yet
+    : `${alertsContextPrompt}
+
+${getContinuePrompt()}
+
+"""
+${combinedMaybePartialResults}
+"""
+
+`;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_continue_prompt/index.ts

@@ -0,0 +1,15 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getContinuePrompt =
+  (): string => `Continue exactly where you left off in the JSON output below, generating only the additional JSON output when it's required to complete your work. The additional JSON output MUST ALWAYS follow these rules:
+1) it MUST conform to the schema above, because it will be checked against the JSON schema
+2) it MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds), because it will be parsed as JSON
+3) it MUST NOT repeat any the previous output, because that would prevent partial results from being combined
+4) it MUST NOT restart from the beginning, because that would prevent partial results from being combined
+5) it MUST NOT be prefixed or suffixed with additional text outside of the JSON, because that would prevent it from being combined and parsed as JSON:
+`;

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_default_attack_discovery_prompt/index.ts

@@ -0,0 +1,9 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getDefaultAttackDiscoveryPrompt = (): string =>
+  "You are a cyber security analyst tasked with analyzing security events from Elastic Security to identify and report on potential cyber attacks or progressions. Your report should focus on high-risk incidents that could severely impact the organization, rather than isolated alerts. Present your findings in a way that can be easily understood by anyone, regardless of their technical expertise, as if you were briefing the CISO. Break down your response into sections based on timing, hosts, and users involved. When correlating alerts, use kibana.alert.original_time when it's available, otherwise use @timestamp. Include appropriate context about the affected hosts and users. Describe how the attack progression might have occurred and, if feasible, attribute it to known threat groups. Prioritize high and critical alerts, but include lower-severity alerts if desired. In the description field, provide as much detail as possible, in a bulleted list explaining any attack progressions. Accuracy is of utmost importance. You MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds).";

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_output_parser/index.test.ts

@@ -0,0 +1,31 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { getOutputParser } from '.';
+
+describe('getOutputParser', () => {
+  it('returns a structured output parser with the expected format instructions', () => {
+    const outputParser = getOutputParser();
+
+    const expected = `You must format your output as a JSON value that adheres to a given \"JSON Schema\" instance.
+
+\"JSON Schema\" is a declarative language that allows you to annotate and validate JSON documents.
+
+For example, the example \"JSON Schema\" instance {{\"properties\": {{\"foo\": {{\"description\": \"a list of test words\", \"type\": \"array\", \"items\": {{\"type\": \"string\"}}}}}}, \"required\": [\"foo\"]}}}}
+would match an object with one required property, \"foo\". The \"type\" property specifies \"foo\" must be an \"array\", and the \"description\" property semantically describes it as \"a list of test words\". The items within \"foo\" must be strings.
+Thus, the object {{\"foo\": [\"bar\", \"baz\"]}} is a well-formatted instance of this example \"JSON Schema\". The object {{\"properties\": {{\"foo\": [\"bar\", \"baz\"]}}}} is not well-formatted.
+
+Your output will be parsed and type-checked according to the provided schema instance, so make sure all fields in your output match the schema exactly and there are no trailing commas!
+
+Here is the JSON Schema instance your output must adhere to. Include the enclosing markdown codeblock:
+\`\`\`json
+{"type":"object","properties":{"insights":{\"type\":\"array\",\"items\":{\"type\":\"object\",\"properties\":{\"alertIds\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"description\":\"The alert IDs that the insight is based on.\"},\"detailsMarkdown\":{\"type\":\"string\",\"description\":\"A detailed insight with markdown, where each markdown bullet contains a description of what happened that reads like a story of the attack as it played out and always uses special {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax for field names and values from the source data. Examples of CORRECT syntax (includes field names and values): {{ host.name hostNameValue }} {{ user.name userNameValue }} {{ source.ip sourceIpValue }} Examples of INCORRECT syntax (bad, because the field names are not included): {{ hostNameValue }} {{ userNameValue }} {{ sourceIpValue }}\"},\"entitySummaryMarkdown\":{\"type\":\"string\",\"description\":\"A short (no more than a sentence) summary of the insight featuring only the host.name and user.name fields (when they are applicable), using the same {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax\"},\"mitreAttackTactics\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"description\":\"An array of MITRE ATT&CK tactic for the insight, using one of the following values: Reconnaissance,Initial Access,Execution,Persistence,Privilege Escalation,Discovery,Lateral Movement,Command and Control,Exfiltration\"},\"summaryMarkdown\":{\"type\":\"string\",\"description\":\"A markdown summary of insight, using the same {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax\"},\"title\":{\"type\":\"string\",\"description\":\"A short, no more than 7 words, title for the insight, NOT formatted with special syntax or markdown. This must be as brief as possible.\"}},\"required\":[\"alertIds\",\"detailsMarkdown\",\"summaryMarkdown\",\"title\"],\"additionalProperties\":false},\"description\":\"Insights with markdown that always uses special {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax for field names and values from the source data. Examples of CORRECT syntax (includes field names and values): {{ host.name hostNameValue }} {{ user.name userNameValue }} {{ source.ip sourceIpValue }} Examples of INCORRECT syntax (bad, because the field names are not included): {{ hostNameValue }} {{ userNameValue }} {{ sourceIpValue }}\"}},\"required\":[\"insights\"],\"additionalProperties":false,\"$schema\":\"http://json-schema.org/draft-07/schema#\"}
+\`\`\`
+`;
+
+    expect(outputParser.getFormatInstructions()).toEqual(expected);
+  });
+});

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_output_parser/index.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { StructuredOutputParser } from 'langchain/output_parsers';
+
+import { AttackDiscoveriesGenerationSchema } from '../../generate/schema';
+
+export const getOutputParser = () =>
+  StructuredOutputParser.fromZodSchema(AttackDiscoveriesGenerationSchema);

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/parse_combined_or_throw/index.ts

@@ -0,0 +1,53 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Logger } from '@kbn/core/server';
+import type { AttackDiscovery } from '@kbn/elastic-assistant-common';
+
+import { addTrailingBackticksIfNecessary } from '../add_trailing_backticks_if_necessary';
+import { extractJson } from '../extract_json';
+import { AttackDiscoveriesGenerationSchema } from '../../generate/schema';
+
+export const parseCombinedOrThrow = ({
+  combinedResponse,
+  generationAttempts,
+  llmType,
+  logger,
+  nodeName,
+}: {
+  /** combined responses that maybe valid JSON */
+  combinedResponse: string;
+  generationAttempts: number;
+  nodeName: string;
+  llmType: string;
+  logger?: Logger;
+}): AttackDiscovery[] => {
+  const timestamp = new Date().toISOString();
+
+  const extractedJson = extractJson(addTrailingBackticksIfNecessary(combinedResponse));
+
+  logger?.debug(
+    () =>
+      `${nodeName} node is parsing extractedJson (${llmType}) from attempt ${generationAttempts}`
+  );
+
+  const unvalidatedParsed = JSON.parse(extractedJson);
+
+  logger?.debug(
+    () =>
+      `${nodeName} node is validating combined response (${llmType}) from attempt ${generationAttempts}`
+  );
+
+  const validatedResponse = AttackDiscoveriesGenerationSchema.parse(unvalidatedParsed);
+
+  logger?.debug(
+    () =>
+      `${nodeName} node successfully validated Attack discoveries response (${llmType}) from attempt ${generationAttempts}`
+  );
+
+  return [...validatedResponse.insights.map((insight) => ({ ...insight, timestamp }))];
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/response_is_hallucinated/index.ts

@@ -0,0 +1,9 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const responseIsHallucinated = (result: string): boolean =>
+  result.includes('{{ host.name hostNameValue }}');

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/discard_previous_refinements/index.ts

@@ -0,0 +1,30 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { GraphState } from '../../../../types';
+
+export const discardPreviousRefinements = ({
+  generationAttempts,
+  hallucinationFailures,
+  isHallucinationDetected,
+  state,
+}: {
+  generationAttempts: number;
+  hallucinationFailures: number;
+  isHallucinationDetected: boolean;
+  state: GraphState;
+}): GraphState => {
+  return {
+    ...state,
+    combinedRefinements: '', // <-- reset the combined refinements
+    generationAttempts: generationAttempts + 1,
+    refinements: [], // <-- reset the refinements
+    hallucinationFailures: isHallucinationDetected
+      ? hallucinationFailures + 1
+      : hallucinationFailures,
+  };
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_combined_refine_prompt/index.ts

@@ -0,0 +1,48 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { AttackDiscovery } from '@kbn/elastic-assistant-common';
+import { isEmpty } from 'lodash/fp';
+
+import { getContinuePrompt } from '../../../helpers/get_continue_prompt';
+
+/**
+ * Returns a prompt that combines the initial query, a refine prompt, and partial results
+ */
+export const getCombinedRefinePrompt = ({
+  attackDiscoveryPrompt,
+  combinedRefinements,
+  refinePrompt,
+  unrefinedResults,
+}: {
+  attackDiscoveryPrompt: string;
+  combinedRefinements: string;
+  refinePrompt: string;
+  unrefinedResults: AttackDiscovery[] | null;
+}): string => {
+  const baseQuery = `${attackDiscoveryPrompt}
+
+${refinePrompt}
+
+"""
+${JSON.stringify(unrefinedResults, null, 2)}
+"""
+
+`;
+
+  return isEmpty(combinedRefinements)
+    ? baseQuery // no partial results yet
+    : `${baseQuery}
+
+${getContinuePrompt()}
+
+"""
+${combinedRefinements}
+"""
+
+`;
+};

x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_default_refine_prompt/index.ts

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const getDefaultRefinePrompt =
+  (): string => `You previously generated the following insights, but sometimes they represent the same attack.
+
+Combine the insights below, when they represent the same attack; leave any insights that are not combined unchanged:`;