From 17440896f98dc76ea532f771caa6b82b9cc853a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Patryk=20Kopyci=C5=84ski?= Date: Tue, 8 Apr 2025 01:25:45 +0200 Subject: [PATCH] [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) - [ ] ... --- .../knowledge_base/index.test.ts | 2 +- .../ai_assistant_data_clients/knowledge_base/index.ts | 10 +++++----- .../server/ai_assistant_service/index.ts | 4 ++-- .../server/routes/request_context_factory.ts | 9 +++++---- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts index aa8fcb0c1922..1d10c49e0446 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts @@ -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 diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index e54ff6cbf84e..2b68dbb2ab24 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -90,7 +90,7 @@ export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientPara ingestPipelineResourceName: string; setIsKBSetupInProgress: (spaceId: string, isInProgress: boolean) => void; manageGlobalKnowledgeBaseAIAssistant: boolean; - trainedModelsProvider: ReturnType; + getTrainedModelsProvider: () => ReturnType; 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, }); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 2a2f4f5e8520..c5fc659db565 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -547,7 +547,7 @@ export class AIAssistantService { public async createAIAssistantKnowledgeBaseDataClient( opts: CreateAIAssistantClientParams & GetAIAssistantKnowledgeBaseDataClientParams & { - trainedModelsProvider: ReturnType; + getTrainedModelsProvider: () => ReturnType; } ): Promise { // 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, }); } diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts index 6c9b373e94ca..cdc2559e477b 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -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() + ), }); }),