[Obs AI Assistant] knowledge base integration tests (#189000)

Closes #188999

- integration tests for knowledge base api
- adds new config field `modelId`, for internal use, to override elser
model id
- refactors `knowledgeBaseService.setup()` to fix bug where if the model
failed to install when calling ml.putTrainedModel, we dont get stuck
polling and retrying the install. We were assuming that the first error
that gets throw when the model is exists would only happen once and the
return true or false and poll for whether its done installing. But the
installation could fail itself causing getTrainedModelsStats to
continuously throw and try to install the model. Now user immediately
gets error if model fails to install and polling does not happen.

---------

Co-authored-by: James Gowdy <jgowdy@elastic.co>

x-pack/plugins/observability_solution/observability_ai_assistant/server/config.ts

@@ -9,6 +9,7 @@ import { schema, type TypeOf } from '@kbn/config-schema';
 
 export const config = schema.object({
   enabled: schema.boolean({ defaultValue: true }),
+  modelId: schema.maybe(schema.string()),
 });
 
 export type ObservabilityAIAssistantConfig = TypeOf<typeof config>;

x-pack/plugins/observability_solution/observability_ai_assistant/server/plugin.ts

@@ -47,10 +47,12 @@ export class ObservabilityAIAssistantPlugin
     >
 {
   logger: Logger;
+  config: ObservabilityAIAssistantConfig;
   service: ObservabilityAIAssistantService | undefined;
 
   constructor(context: PluginInitializerContext<ObservabilityAIAssistantConfig>) {
     this.logger = context.logger.get();
+    this.config = context.config.get<ObservabilityAIAssistantConfig>();
     initLangtrace();
   }
   public setup(
@@ -112,10 +114,14 @@ export class ObservabilityAIAssistantPlugin
 
     // Using once to make sure the same model ID is used during service init and Knowledge base setup
     const getModelId = once(async () => {
+      const configModelId = this.config.modelId;
+      if (configModelId) {
+        return configModelId;
+      }
       const defaultModelId = '.elser_model_2';
       const [_, pluginsStart] = await core.getStartServices();
+      // Wait for the license to be available so the ML plugin's guards pass once we ask for ELSER stats
       const license = await firstValueFrom(pluginsStart.licensing.license$);
-
       if (!license.hasAtLeast('enterprise')) {
         return defaultModelId;
       }

x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 import { errors } from '@elastic/elasticsearch';
-import { serverUnavailable, gatewayTimeout } from '@hapi/boom';
+import { serverUnavailable, gatewayTimeout, badRequest } from '@hapi/boom';
 import type { ElasticsearchClient, IUiSettingsClient } from '@kbn/core/server';
 import type { Logger } from '@kbn/logging';
 import type { TaskManagerStartContract } from '@kbn/task-manager-plugin/server';
@@ -39,14 +39,20 @@ export interface RecalledEntry {
   labels?: Record<string, string>;
 }
 
-function isAlreadyExistsError(error: Error) {
+function isModelMissingOrUnavailableError(error: Error) {
   return (
     error instanceof errors.ResponseError &&
     (error.body.error.type === 'resource_not_found_exception' ||
       error.body.error.type === 'status_exception')
   );
 }
-
+function isCreateModelValidationError(error: Error) {
+  return (
+    error instanceof errors.ResponseError &&
+    error.statusCode === 400 &&
+    error.body?.error?.type === 'action_request_validation_exception'
+  );
+}
 function throwKnowledgeBaseNotReady(body: any) {
   throw serverUnavailable(`Knowledge base is not ready yet`, body);
 }
@@ -84,52 +90,73 @@ export class KnowledgeBaseService {
     const elserModelId = await this.dependencies.getModelId();
 
     const retryOptions = { factor: 1, minTimeout: 10000, retries: 12 };
-
-    const installModel = async () => {
-      this.dependencies.logger.info('Installing ELSER model');
-      await this.dependencies.esClient.asInternalUser.ml.putTrainedModel(
-        {
-          model_id: elserModelId,
-          input: {
-            field_names: ['text_field'],
-          },
-          wait_for_completion: true,
-        },
-        { requestTimeout: '20m' }
-      );
-      this.dependencies.logger.info('Finished installing ELSER model');
-    };
-
-    const getIsModelInstalled = async () => {
-      const getResponse = await this.dependencies.esClient.asInternalUser.ml.getTrainedModels({
+    const getModelInfo = async () => {
+      return await this.dependencies.esClient.asInternalUser.ml.getTrainedModels({
         model_id: elserModelId,
         include: 'definition_status',
       });
+    };
 
-      this.dependencies.logger.debug(
-        () => 'Model definition status:\n' + JSON.stringify(getResponse.trained_model_configs[0])
-      );
+    const isModelInstalledAndReady = async () => {
+      try {
+        const getResponse = await getModelInfo();
+        this.dependencies.logger.debug(
+          () => 'Model definition status:\n' + JSON.stringify(getResponse.trained_model_configs[0])
+        );
 
-      return Boolean(getResponse.trained_model_configs[0]?.fully_defined);
+        return Boolean(getResponse.trained_model_configs[0]?.fully_defined);
+      } catch (error) {
+        if (isModelMissingOrUnavailableError(error)) {
+          return false;
+        } else {
+          throw error;
+        }
+      }
     };
 
-    await pRetry(async () => {
-      let isModelInstalled: boolean = false;
+    const installModelIfDoesNotExist = async () => {
+      const modelInstalledAndReady = await isModelInstalledAndReady();
+      if (!modelInstalledAndReady) {
+        await installModel();
+      }
+    };
+
+    const installModel = async () => {
+      this.dependencies.logger.info('Installing ELSER model');
       try {
-        isModelInstalled = await getIsModelInstalled();
+        await this.dependencies.esClient.asInternalUser.ml.putTrainedModel(
+          {
+            model_id: elserModelId,
+            input: {
+              field_names: ['text_field'],
+            },
+            wait_for_completion: true,
+          },
+          { requestTimeout: '20m' }
+        );
       } catch (error) {
-        if (isAlreadyExistsError(error)) {
-          await installModel();
-          isModelInstalled = await getIsModelInstalled();
+        if (isCreateModelValidationError(error)) {
+          throw badRequest(error);
+        } else {
+          throw error;
         }
       }
+      this.dependencies.logger.info('Finished installing ELSER model');
+    };
 
-      if (!isModelInstalled) {
-        throwKnowledgeBaseNotReady({
-          message: 'Model is not fully defined',
-        });
-      }
-    }, retryOptions);
+    const pollForModelInstallCompleted = async () => {
+      await pRetry(async () => {
+        this.dependencies.logger.info('Polling installation of ELSER model');
+        const modelInstalledAndReady = await isModelInstalledAndReady();
+        if (!modelInstalledAndReady) {
+          throwKnowledgeBaseNotReady({
+            message: 'Model is not fully defined',
+          });
+        }
+      }, retryOptions);
+    };
+    await installModelIfDoesNotExist();
+    await pollForModelInstallCompleted();
 
     try {
       await this.dependencies.esClient.asInternalUser.ml.startTrainedModelDeployment({
@@ -139,7 +166,7 @@ export class KnowledgeBaseService {
     } catch (error) {
       this.dependencies.logger.debug('Error starting model deployment');
       this.dependencies.logger.debug(error);
-      if (!isAlreadyExistsError(error)) {
+      if (!isModelMissingOrUnavailableError(error)) {
         throw error;
       }
     }
@@ -380,7 +407,7 @@ export class KnowledgeBaseService {
         namespace,
         modelId,
       }).catch((error) => {
-        if (isAlreadyExistsError(error)) {
+        if (isModelMissingOrUnavailableError(error)) {
           throwKnowledgeBaseNotReady(error.body);
         }
         throw error;
@@ -521,7 +548,7 @@ export class KnowledgeBaseService {
         })),
       };
     } catch (error) {
-      if (isAlreadyExistsError(error)) {
+      if (isModelMissingOrUnavailableError(error)) {
         throwKnowledgeBaseNotReady(error.body);
       }
       throw error;
@@ -588,7 +615,7 @@ export class KnowledgeBaseService {
 
       return Promise.resolve();
     } catch (error) {
-      if (isAlreadyExistsError(error)) {
+      if (isModelMissingOrUnavailableError(error)) {
         throwKnowledgeBaseNotReady(error.body);
       }
       throw error;

x-pack/test/functional/services/ml/api.ts

@@ -1473,11 +1473,13 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       }
     },
 
-    async stopTrainedModelDeploymentES(deploymentId: string) {
+    async stopTrainedModelDeploymentES(deploymentId: string, force: boolean = false) {
       log.debug(`Stopping trained model deployment with id "${deploymentId}"`);
-      const { body, status } = await esSupertest.post(
-        `/_ml/trained_models/${deploymentId}/deployment/_stop`
-      );
+      const url = `/_ml/trained_models/${deploymentId}/deployment/_stop${
+        force ? '?force=true' : ''
+      }`;
+
+      const { body, status } = await esSupertest.post(url);
       this.assertResponseStatusCode(200, status, body);
 
       log.debug('> Trained model deployment stopped');
@@ -1570,8 +1572,13 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       );
     },
 
-    async importTrainedModel(modelId: string, modelName: SupportedTrainedModelNamesType) {
-      await this.createTrainedModel(modelId, this.getTrainedModelConfig(modelName));
+    async importTrainedModel(
+      modelId: string,
+      modelName: SupportedTrainedModelNamesType,
+      config?: PutTrainedModelConfig
+    ) {
+      const trainedModelConfig = config ?? this.getTrainedModelConfig(modelName);
+      await this.createTrainedModel(modelId, trainedModelConfig);
       await this.createTrainedModelVocabularyES(modelId, this.getTrainedModelVocabulary(modelName));
       await this.uploadTrainedModelDefinitionES(
         modelId,

x-pack/test/observability_ai_assistant_api_integration/configs/index.ts

@@ -8,6 +8,7 @@
 import { mapValues } from 'lodash';
 import path from 'path';
 import { createTestConfig, CreateTestConfig } from '../common/config';
+import { SUPPORTED_TRAINED_MODELS } from '../../functional/services/ml/api';
 
 export const observabilityAIAssistantDebugLogger = {
   name: 'plugins.observabilityAIAssistant',
@@ -30,6 +31,7 @@ export const observabilityAIAssistantFtrConfigs = {
         __dirname,
         '../../../../test/analytics/plugins/analytics_ftr_helpers'
       ),
+      'xpack.observabilityAIAssistant.modelId': SUPPORTED_TRAINED_MODELS.TINY_ELSER.name,
     },
   },
 };

x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/helpers.ts

@@ -0,0 +1,31 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MachineLearningProvider } from '../../../api_integration/services/ml';
+import { SUPPORTED_TRAINED_MODELS } from '../../../functional/services/ml/api';
+
+export const TINY_ELSER = {
+  ...SUPPORTED_TRAINED_MODELS.TINY_ELSER,
+  id: SUPPORTED_TRAINED_MODELS.TINY_ELSER.name,
+};
+
+export async function createKnowledgeBaseModel(ml: ReturnType<typeof MachineLearningProvider>) {
+  const config = {
+    ...ml.api.getTrainedModelConfig(TINY_ELSER.name),
+    input: {
+      field_names: ['text_field'],
+    },
+  };
+  await ml.api.importTrainedModel(TINY_ELSER.name, TINY_ELSER.id, config);
+  await ml.api.assureMlStatsIndexExists();
+}
+export async function deleteKnowledgeBaseModel(ml: ReturnType<typeof MachineLearningProvider>) {
+  await ml.api.stopTrainedModelDeploymentES(TINY_ELSER.id, true);
+  await ml.api.deleteTrainedModelES(TINY_ELSER.id);
+  await ml.api.cleanMlIndices();
+  await ml.testResources.cleanMLSavedObjects();
+}