[Security Assistant] Don't cache MLClient (#217408)

## Summary

Summarize your PR. If it involves visual changes include a screenshot or
gif.


### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [ ] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [ ] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [ ] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [ ] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [ ] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [ ] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

### Identify risks

Does this PR introduce any risks? For example, consider risks like hard
to test bugs, performance regression, potential of data loss.

Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging.

- [ ] [See some risk
examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx)
- [ ] ...
This commit is contained in:
Patryk Kopyciński 2025-04-08 01:25:45 +02:00 committed by GitHub
parent 902bca98f1
commit 17440896f9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 13 additions and 12 deletions

View file

@ -69,7 +69,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
ingestPipelineResourceName: 'something',
setIsKBSetupInProgress: jest.fn().mockImplementation(() => {}),
manageGlobalKnowledgeBaseAIAssistant: true,
trainedModelsProvider: trainedModelsProviderMock,
getTrainedModelsProvider: () => trainedModelsProviderMock,
};
esClientMock.search.mockReturnValue(
// @ts-expect-error not full response interface

View file

@ -90,7 +90,7 @@ export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientPara
ingestPipelineResourceName: string;
setIsKBSetupInProgress: (spaceId: string, isInProgress: boolean) => void;
manageGlobalKnowledgeBaseAIAssistant: boolean;
trainedModelsProvider: ReturnType<TrainedModelsProvider['trainedModelsProvider']>;
getTrainedModelsProvider: () => ReturnType<TrainedModelsProvider['trainedModelsProvider']>;
modelIdOverride: boolean;
}
export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
@ -131,7 +131,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
this.options.logger.debug(`Installing ELSER model '${elserId}'...`);
try {
await this.options.trainedModelsProvider.installElasticModel(elserId);
await this.options.getTrainedModelsProvider().installElasticModel(elserId);
} catch (error) {
this.options.logger.error(`Error installing ELSER model '${elserId}':\n${error}`);
}
@ -147,7 +147,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
this.options.logger.debug(`Checking if ELSER model '${elserId}' is installed...`);
try {
const getResponse = await this.options.trainedModelsProvider.getTrainedModels({
const getResponse = await this.options.getTrainedModelsProvider().getTrainedModels({
model_id: elserId,
include: 'definition_status',
});
@ -209,7 +209,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
let getResponse;
try {
getResponse = await this.options.trainedModelsProvider.getTrainedModelsStats({
getResponse = await this.options.getTrainedModelsProvider().getTrainedModelsStats({
model_id: elserId,
});
} catch (e) {
@ -250,7 +250,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
});
this.options.logger.debug(`Dry run for ELSER model '${elserId}' successfully deployed!`);
await this.options.trainedModelsProvider.stopTrainedModelDeployment({
await this.options.getTrainedModelsProvider().stopTrainedModelDeployment({
model_id: elserId,
deployment_id: dryRunId.assignment.task_parameters.deployment_id,
});

View file

@ -547,7 +547,7 @@ export class AIAssistantService {
public async createAIAssistantKnowledgeBaseDataClient(
opts: CreateAIAssistantClientParams &
GetAIAssistantKnowledgeBaseDataClientParams & {
trainedModelsProvider: ReturnType<TrainedModelsProvider['trainedModelsProvider']>;
getTrainedModelsProvider: () => ReturnType<TrainedModelsProvider['trainedModelsProvider']>;
}
): Promise<AIAssistantKnowledgeBaseDataClient | null> {
// If modelIdOverride is set, swap getElserId(), and ensure the pipeline is re-created with the correct model
@ -587,7 +587,7 @@ export class AIAssistantService {
setIsKBSetupInProgress: this.setIsKBSetupInProgress.bind(this),
spaceId: opts.spaceId,
manageGlobalKnowledgeBaseAIAssistant: opts.manageGlobalKnowledgeBaseAIAssistant ?? false,
trainedModelsProvider: opts.trainedModelsProvider,
getTrainedModelsProvider: opts.getTrainedModelsProvider,
});
}

View file

@ -124,10 +124,11 @@ export class RequestContextFactory implements IRequestContextFactory {
manageGlobalKnowledgeBaseAIAssistant:
securitySolutionAssistant.manageGlobalKnowledgeBaseAIAssistant as boolean,
// uses internal user to interact with ML API
trainedModelsProvider: plugins.ml.trainedModelsProvider(
{} as KibanaRequest,
coreStart.savedObjects.createInternalRepository()
),
getTrainedModelsProvider: () =>
plugins.ml.trainedModelsProvider(
{} as KibanaRequest,
coreStart.savedObjects.createInternalRepository()
),
});
}),