[Security Assistant] AI Assistant - Better Solution for OSS models (#10416)

[e40pud]

Summary

This PR improves OSS models integration via OpenAI connector (https://github.com/elastic/security-team/issues/10416).

OSS models do not support "native" tool calling, thus we need to do that via prompting. To do that we use structured chat agent (LangChain.createStructuredChatAgent) and appropriate prompt (see STRUCTURED_SYSTEM_PROMPT in x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts file).

One important note, since we are using OpenAI connector to connect to OSS model we need to recognise and distinguish that which is done in a new utils function isOpenSourceModel in x-pack/plugins/elastic_assistant/server/routes/utils.ts file.

To avoid tool's output parsing in streaming mode, we do enable response step for the OSS models - similar to what we do in case of Bedrock. This simplifies parsing during streaming and safes us headaches from fighting with JSON formatting issues returned by OSS models (especially Llama). I noticed that in case of Llama, the final answer returned with extra escape characters which breaks a markdown.

Tested OSS models

• Mistral Large 2407
• LLama 3.1 70B
• Llama 3.1 405B (NOTE: for some reason Azure deployment does not work right now. Working on it.)

Both models work well with existing ES tools. Mistral shown slightly better results than Llama.

Credentials

Use these credentials for testing LLama 3.1 70B and Mistral Large 2407:

https://p.elstc.co/paste/ofqxUGiW#RR1Pedserj9hWKm3vOw5oVEHWuwbX7i9Jl2rS7q7MRP

Notes about ES|QL generation

By default works with the ESQLKnowledgeBaseTool tool. The new NaturalLanguageESQLTool provides better results. To enable it use:

xpack.securitySolution.enableExperimental:
  - 'assistantNaturalLanguageESQLTool'

Example of ES|QL generation using Mistral model and NaturalLanguageESQLTool tool

mistral.mov

Evaluation

Experiment 1

Suite: ES|QL Generation Regression Suite
Models: Llama 3.1 405B and Mistral Large 2407
Tool: ESQLKnowledgeBaseTool
Correctness results:

• Llama 3.1 405B
• Mistral Large 2407

Experiment 2

Suite: ES|QL Generation Regression Suite
Models: Llama 3.1 405B and Mistral Large 2407
Tool: NaturalLanguageESQLTool
Results:

• Llama 3.1 405B
• Mistral Large 2407

Checklist

Delete any items that are not applicable to this PR.

• Unit or functional tests were updated or added to match the most common scenarios

[e40pud]

@elasticmachine merge upstream

[e40pud]

Removing this await solves the issue where different OSS models fail multiple tests due to error:

Error: ActionsClientChatOpenAI: an error occurred while running the action - Status code: undefined. Message: Unexpected API Error: ERR_CANCELED - canceled

This mostly happens with Llama model, but occasionally with Mistral as well.

Here is the example of those failures https://smith.langchain.com/o/b739bf24-7ba4-4994-b632-65dd677ac74e/datasets/261dcc59-fbe7-4397-a662-ff94042f666c/compare?selectedSessions=f26a13b9-73ba-4e6d-992b-3b4eda61bb20&baseline=undefined&activeSession=f26a13b9-73ba-4e6d-992b-3b4eda61bb20

I feel it has something to do with the slowness of the model (not proved yet) which leads to timeouts.

@spong How important for us to make sure that all experiments are finished and only then respond back? Can we just schedule those experiments and return and log evaluation results and handle errors in a different way here? Something like this:

evaluate(predict, {
  data: datasetName ?? '',
  evaluators: [], // Evals to be managed in LangSmith for now
  experimentPrefix: name,
  client: new Client({ apiKey: langSmithApiKey }),
  // prevent rate limiting and unexpected multiple experiment runs
  maxConcurrency: 5,
}).then((evalOutput) => {
  logger.debug(`runResp:\n ${JSON.stringify(evalOutput, null, 2)}`);
}).catch(() => {
  // handle error here
});

also we would need to add better error handling in this case as well.

[spong]

Yeah that should be fine. We used to use the result output and write that to the local ES, but don't do that anymore, and so just debug log the output since we've been using LangSmith for analysis after the fact. Really we just need to know if the eval was successful from the API perspective.

That said, I'm curious why this await would cause internal issues with the chat clients. That shouldn't change anything with the concurrency or executions, but seems with it that previous ones are getting cancelled? Perhaps there's some shared state between the chat clients like @stephmilovic fixed with createLlmInstance over in #190004?

[spong]

If you turn down maxConcurrency to 1 and run with the await do you still see it? Curious if rate limiting or something might be cancelling the previous runs?

[e40pud]

So, I was able to run simple stress testing against our llama deployment and I was able to get this error which I guess is the reason our evaluations are failing sometimes

{
    "statusCode": 429,
    "message": "Rate Limit on Number of Tokens is exceeded. Retry after 8 seconds."
}

I don't think, this is something we cannot control in code and will need to be configured in llm deployment if we need it for evaluations. I found this article https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/quota?tabs=rest#assign-quota which shows how to address that in deployment settings (sent to James as well). There we have an option to assign Tokens-Per-Minute (TPM) to llm deployment.

Btw, lowering maxConcurrency did not help with the issue. Setting it to 1 failed requests even faster for whatever reasons.

[e40pud]

To lower chances of hitting rate limit failure, I removed the await and restructured code in order to log results and errors at the end of execution. This will allow users to avoid 429 errors in most cases.

After talking to @jamesspi, we agreed that this is something user can fix on their end by increasing the quota and resources. We would prefer to keep things simple instead of finding a good solution to deal with 429 during evaluations and make our code complicated.

@spong what do you think?

[e40pud]

@elasticmachine merge upstream

[e40pud]

@elasticmachine merge upstream

[stephmilovic]

Did you consider adding a new apiProvider of Open Source and relying on that for the isOpenSourceModel determination?

[e40pud]

Yes, my second PR is about that #194831

We will add another provide to cover "other" OpenAI compatible services. It will work for all newly added connectors, though we would like to support already created connectors - OSS models via OpenAI provide.

[e40pud]

@elasticmachine merge upstream

[kibana-ci]

💛 Build succeeded, but was flaky

• Buildkite Build
• Commit: 0f96231

Failed CI Steps

Metrics [docs]

Module Count

Fewer modules leads to a faster build time

id	before	after	diff
integrationAssistant	562	563	+1
securitySolution	5925	5926	+1
total			+2

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	37	38	+1

Async chunks

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

id	before	after	diff
securitySolution	20.6MB	20.6MB	+255.0B

Unknown metric groups

API count

id	before	after	diff
elasticAssistant	52	53	+1

History

• 💛 Build #239831 was flaky 26cb4c3
• 💛 Build #239495 was flaky d0458eb
• 💚 Build #239333 succeeded db9fd74
• 💚 Build #239091 succeeded 5c7bf7f
• 💛 Build #238864 was flaky 69543ce
• 💛 Build #238775 was flaky c0a00eb

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

cc @e40pud

[stephmilovic]

LGTM! Thank you for the improvement and for your patience fixing the timeout error. Nice work @e40pud

[elasticmachine]

💛 Build succeeded, but was flaky

• Buildkite Build
• Commit: 59fb225

Failed CI Steps

• Jest Integration Tests #2

Metrics [docs]

Module Count

Fewer modules leads to a faster build time

id	before	after	diff
integrationAssistant	562	563	+1
securitySolution	5925	5926	+1
total			+2

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	37	38	+1

Async chunks

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

id	before	after	diff
securitySolution	20.6MB	20.6MB	+285.0B

Unknown metric groups

API count

id	before	after	diff
elasticAssistant	52	53	+1

cc @e40pud

[kibanamachine]

Starting backport for target branches: 8.x

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

[kibanamachine]

💚 All backports created successfully

Status	Branch	Result
✅	8.x

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

Questions ?

Please refer to the Backport tool documentation