[Security AI] Bedrock prompt tuning and inference corrections

[stephmilovic]

Summary

@jamesspi noticed some issues when invoking inference with an ESQL query, specifically:

Adversaries leverage persistence to keep their foothold within a network. A common way to establish persistence is to add Windows Registry keys that autostart applications after reboot.

Generate an ES|QL query for me to hunt for the addition of registry keys to autorun applications made by suspicious processes. I need to use the information found to name the name of the post-exploitation rootkit used by the bad actor.

Windows Registry Keys are stored in the winlog.event_data.TargetObject ECS field name. Use logs-* as the index pattern.

1. Bedrock needs a Respond step

We noticed the model starts streaming back the conversation with the tool. Meaning, we saw this "Certainly..." line coming into the stream:

To address this, I needed to add the same "Respond" step we have for Bedrock. There is no distinction in Bedrock when the stream is communicating with the tool or with the final answer, which is why we needed the "Respond" step. Since inference is currently using different providers behind the scene, I added a check for the provider and if it's Bedrock we need the "Respond" step.

2. The ESQL query was not in the response

Despite the system prompt including Always return value from NaturalLanguageESQLTool as is., this is not the behavior we were seeing with both Bedrock and Inference. Updating the prompt to be even more strict and specific fixed the issue for both Bedrock and Inference. This is the new portion of the prompt that enforces the results from the ESQL tool: ALWAYS return the exact response from NaturalLanguageESQLTool verbatim in the final response, without adding further description.

3. Evals not configured for inference

I tried to run evals to confirm this change and realized that the post_evaluate route was not properly configured for inference. To correct this, I updated the agent from the deprecated createOpenAIFunctionsAgent to createOpenAIToolsAgent and added a condition to ensure llmType === 'inference' uses the tool calling agent. This update was already made in the other place we have the agent selection logic, callAssistantGraph.

Running the evals after the prompt change improved the accuracy from 94% to 97%

Testing

Testing prompt:

Adversaries leverage persistence to keep their foothold within a network. A common way to establish persistence is to add Windows Registry keys that autostart applications after reboot.

Generate an ES|QL query for me to hunt for the addition of registry keys to autorun applications made by suspicious processes. I need to use the information found to name the name of the post-exploitation rootkit used by the bad actor.

Windows Registry Keys are stored in the winlog.event_data.TargetObject ECS field name. Use logs-* as the index pattern.

1. Have a Bedrock connector and an Inference connector configured with EIS (ping Steph if you need help setting that up)
2. With Bedrock selected, send the testing prompt. Ensure the response contains an ESQL query.
3. With Inference selected, send the testing prompt. Ensure the response does not start streaming early with the tool invocation (Certainly...) and contains an ESQL query.
4. To confirm the eval fix, run the ESQL eval with Inference connector selected. Previously, an error showed on the securityAiAssistantManagement?tab=evaluation view after hitting "Perform evaluation..."

[stephmilovic]

this is actually an important param for OpenAI, removing my comment

[elasticmachine]

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

[KDKHD]

Looks great! I was able to see the issues before this PR and confirm they are now resolved. The only thing I was unable to test was this section of the test guide:

To confirm the eval fix, run the ESQL eval with Inference connector selected. Previously, an error showed on the securityAiAssistantManagement?tab=evaluation view after hitting "Perform evaluation..."

because the inference connector did not appear on the evaluation page.

[stephmilovic]

the inference connector did not appear on the evaluation page.

2e87e12

This commit ensures inference shows on eval page

[elasticmachine]

💛 Build succeeded, but was flaky

• Buildkite Build
• Commit: 2e87e12

Failed CI Steps

• FTR Configs #118

Test Failures

• [job] [logs] FTR Configs #118 / console app console autocomplete feature Autocomplete behavior JSON autocompletion with placeholder fields

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
@kbn/security-ai-prompts	31	38	+7

Async chunks

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

id	before	after	diff
securitySolution	21.4MB	21.4MB	+40.0B

Unknown metric groups

API count

id	before	after	diff
@kbn/security-ai-prompts	33	40	+7

History

• 💚 Build #272168 succeeded 52ff53e
• 💛 Build #272103 was flaky c2a2361

[kibanamachine]

Starting backport for target branches: 8.18, 8.x, 9.0

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

[kibanamachine]

💔 Some backports could not be created

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

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

Manual backport

To create the backport manually run:

node scripts/backport --pr 209011

Questions ?

Please refer to the Backport tool documentation

[stephmilovic]

💚 All backports created successfully

Status	Branch	Result
✅	8.18

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

Questions ?

Please refer to the Backport tool documentation

[kibanamachine]

Looks like this PR has backport PRs but they still haven't been merged. Please merge them ASAP to keep the branches relatively in sync.