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

# Backport This will backport the following commits from `main` to `8.18`: - [[Security Assistant] Don't cache MLClient (#217408)](https://github.com/elastic/kibana/pull/217408)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport)  Co-authored-by: Patryk Kopyciński <contact@patrykkopycinski.com>
2025-04-23 17:28:26 -04:00 · 2025-04-08 03:10:18 +02:00 · 2025-04-08 03:10:18 +02:00 · b59647e5ef
commit b59647e5ef
parent 6d142b805e
4 changed files with 13 additions and 12 deletions
--- 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
--- 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<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,
      });
--- 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<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,
    });
  }

--- 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()
+            ),
        });
      }),