[Security solution] naturalLanguageToEsql Tool added to default assistant graph

[stephmilovic]

Summary

Adds a new tool to AI Assistant and Evaluator - NaturalLanguageESQLTool, which utilizes the new inference NL-to-ESQL task. It is behind a feature flag. When the feature flag is enabled, the new NaturalLanguageESQLTool is used instead of the ESQLKnowledgeBaseTool. To enable, add this to your kibana.dev.yml:

xpack.securitySolution.enableExperimental: ['assistantNaturalLanguageESQLTool']

Test runs

To see the tool in use, check out the ES|QL Generation Regression Suite. As Gemini is still in progress, we are only comparing Bedrock and OpenAI runs. Please compare the ESQLKnowledgeBaseTool runs (91-94) compared to the NaturalLanguageESQLTool runs (111-114). As you can see, OpenAI runs remained consistent in correctness and Bedrock saw improvements

ESQLKnowledgeBaseTool

NaturalLanguageESQLTool

Note:

The correctness difference between the 2 NaturalLanguageESQLTool Bedrock runs is due to Bedrock server errors and we can assume this is a server issue, not an issue with the tool

[elasticmachine]

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

[spong]

Great stuff @stephmilovic! 🎉 And good to see where we landed on evals between the two, so thanks for including those!

[jloleysens]

kibana jsonc lgtm

[stephmilovic]

@elasticmachine merge upstream

[stephmilovic]

@elasticmachine merge upstream

[spong]

nit: consider making naturalLanguageESQLToolEnabled optional and default to false

[spong]

We'll probably want to update this to be less suggestive of the additional o11y functions like dataset info, visualize data, etc.

I'll keep an eye out while testing for any flake here, but fine for now while behind a feature flag 👍

[spong]

Just making a note from my other recent works, but isEnabledKnowledgeBase && modelExists can be collapsed into a single param now (in #192665 I'm updating them to be backed by the same value). No change needed here, just noting this upcoming refactor.

[pgayvallet]

What exactly question is here? Is that only the user's last message?

The task behaves better when additional context is provided (info about which index / index pattern is being targeted, the index's schema or relevant fields to avoid hallucinating fields, and so on). Is that kind of user query rewriting / enhancement performed before calling the task, or not at all? (as this is not something that can be done in this black box, unfortunately)

[stephmilovic]

What exactly question is here? Is that only the user's last message?

This is going to be however the LLM interprets how it needs to phrase the question to the tool. So for example, take a look at this trace. Within this trace, click into the first orange graph step ActionsClientChatOpenAI.

Here you can see the user's question is:

Generate an ES|QL query that will search for the following use case: 
Identifying Rare User Agent

The LLM is then given this message along with tools, tool descriptions, and tool schemas. The LLM then determines that it needs to create an input of question with a question formatted about ESQL. So in this caes, the LLM returns an instruction for a tool call:

  "function_call": {
    "arguments": "{\"question\":\"Generate an ES|QL query to identify rare user agents.\"}",
    "name": "NaturalLanguageESQLTool"
  }

Right now the tool schema is:

      schema: z.object({
        question: z.string().describe(`The user's exact question about ESQL`),
      }),

We can try to add prompting for optional fields if an index is specified, but if the user do not provide it we will not have it. Should I try that and run evaluations to see if there is an improvement?

      schema: z.object({
        question: z.string().describe(`The user's exact question about ESQL`),
        index: z.string().optional().describe(`The index the user is referencing.`),
      }),

Patryk did have some code in his tool before the NL to ESQL task was available that validated the query alongside the available data views: https://github.com/elastic/kibana/pull/186489/files#diff-bf442ff72176edbab83ea0e5c13d7d23b9273d851bda705fbc6a000afd19232aR209

I did not include that in this PR. I was going to let him expand on that when he returns from PTO on Sept 23.

[spong]

@pgayvallet, IMO the NL->ESQL task should be handling all those abstractions. From query re-writing the original prompt, to supplying any dataset info/mappings, and finally both syntax and functional validation.

This would be extremely beneficial for both consistency and ease of use on the consumer side. Soon we'll be bundling this task into other/more complex tasks like generalized retrievers, visualizations, etc, so the less initial input and custom context packing required the better here. Looks like we'll need to update the interface to take some authorizedUser or scoped request so that data access can be taken into account, but looks like it should have everything it needs after that.

[pgayvallet]

Yeah, I mostly agree, except for that part:

to supplying any dataset info/mappings

Retrieving the mapping from an index is fine, but I'm wondering about the "resolving/deducing which index to use" part:

How would you see that being done in the task itself? The task is a black box, it has no knowledge of the current "context" of the user. How do you see it being able to deduce which index to target?

FWIW, the naturalLanguageToEsql accepts a messages parameter that can be used instead of input, to provide the full conversation (including tool calls) to the task. This is how the o11y assistant is providing the context, as they have a context and dataset tools that the LLM always call because calling the task. That way, the conversation passed to naturalLanguageToEsql has those tool calls, and the agent utilize it to know which index is being targeted

[spong]

Checked out, tested locally and overall implementation LGTM! Thanks for wiring this up @stephmilovic!

Few notes from testing:

1. As discussed in this thread, I think we should expand the responsibilities of the NL->ESQL task in the inference plugin to include query re-writing, inclusion of dataset info, validation, etc. This will ensure better consistency in responses across assistants/consumers, while making the task easier to consume and compose into other tasks.
2. Before release/final testing, we'll need to update the main tool description as mentioned here since it references o11y specific functions that are not yet available in the inference plugin.
3. When testing with Gemini-1.5-pro/Sonnet-3.5, I was seeing the summarizer take over after the tool call and wipe out the query. See this Sonnet trace where there's two more LLM hops after the tool response and this Gemini trace where it thanks the tool for its response 😅. The responses from the tool itself seem mostly 'user-ready', so perhaps we try and return straight from the tool call?

[stephmilovic]

@spong thanks for the review!

1. As discussed in this thread, I think we should expand the responsibilities of the NL->ESQL task in the inference plugin to include query re-writing, inclusion of dataset info, validation, etc. This will ensure better consistency in responses across assistants/consumers, while making the task easier to consume and compose into other tasks.

Sounds good. I saw in the thread that @pgayvallet would add an indices field to the naturalLanguageToEsql method arguments. I will look out for this addition and pass indices once it has been incorporated

2. Before release/final testing, we'll need to update the main tool description as mentioned here since it references o11y specific functions that are not yet available in the inference plugin.

Here is that change: 340649c

3. When testing with Gemini-1.5-pro/Sonnet-3.5, I was seeing the summarizer take over after the tool call and wipe out the query. See this Sonnet trace where there's two more LLM hops after the tool response and this Gemini trace where it thanks the tool for its response 😅. The responses from the tool itself seem mostly 'user-ready', so perhaps we try and return straight from the tool call?

With LangSmith currently down, I can't see these responses. I'll take a look when I am back from PTO. I can say, I too have seen Gemini talking to itself. I have not quite figured that one out yet but am getting closer. Before LangSmith went down today, I made great progress to improve prompts in the VertexChatAI pr. I plan on iterating further in Gemini, and coming back to Bedrock as well as I saw at least one client issue about performance with that model. I'll reference this PR when working on both of those models.

[kibana-ci]

💛 Build succeeded, but was flaky

• Buildkite Build
• Commit: 340649c

Failed CI Steps

• Jest Tests #4

Test Failures

• [job] [logs] Jest Tests #4 / When using the ResponseActionFileDownloadLink component should display custom button label

Metrics [docs]

Public APIs missing comments

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

id	before	after	diff
elasticAssistant	35	38	+3

Page load bundle

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

id	before	after	diff
securitySolution	87.0KB	87.0KB	+36.0B

Unknown metric groups

API count

id	before	after	diff
elasticAssistant	49	53	+4

History

• 💚 Build #234875 succeeded 96ae4eb
• 💔 Build #234581 failed cc06eab
• 💔 Build #234455 failed dd40a8b
• 💔 Build #234241 failed 7330c3b
• 💚 Build #233878 succeeded 10a87ff
• 💔 Build #233868 failed 62e72b3

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

[kibanamachine]

💔 All backports failed

Status	Branch	Result
❌	8.x	Backport failed because of merge conflicts

Manual backport

To create the backport manually run:

node scripts/backport --pr 192042

Questions ?

Please refer to the Backport tool documentation

[stephmilovic]

💚 All backports created successfully

Status	Branch	Result
✅	8.x

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

Questions ?

Please refer to the Backport tool documentation