[Obs AI Assistant] Do not attempt to overwrite index mappings if inference id in component template has changed

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/create_or_update_knowledge_base_index_assets.ts

@@ -10,11 +10,12 @@ import { createConcreteWriteIndex, getDataStreamAdapter } from '@kbn/alerting-pl
 import type { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
 import { getComponentTemplate } from './templates/kb_component_template';
 import { resourceNames } from '..';
+import { getInferenceIdFromWriteIndex } from '../knowledge_base_service/get_inference_id_from_write_index';
 
 export async function createOrUpdateKnowledgeBaseIndexAssets({
   logger,
   core,
-  inferenceId,
+  inferenceId: componentTemplateInferenceId,
 }: {
   logger: Logger;
   core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
@@ -23,13 +24,14 @@ export async function createOrUpdateKnowledgeBaseIndexAssets({
   try {
     logger.debug('Setting up knowledge base index assets');
     const [coreStart] = await core.getStartServices();
-    const { asInternalUser } = coreStart.elasticsearch.client;
+    const esClient = coreStart.elasticsearch.client;
+    const { asInternalUser } = esClient;
 
     // Knowledge base: component template
     await asInternalUser.cluster.putComponentTemplate({
       create: false,
       name: resourceNames.componentTemplate.kb,
-      template: getComponentTemplate(inferenceId),
+      template: getComponentTemplate(componentTemplateInferenceId),
     });
 
     // Knowledge base: index template
@@ -47,21 +49,29 @@ export async function createOrUpdateKnowledgeBaseIndexAssets({
       },
     });
 
+    const writeIndexInferenceId = await getInferenceIdFromWriteIndex(esClient).catch(
+      () => undefined
+    );
+
     // Knowledge base: write index
-    const kbAliasName = resourceNames.writeIndexAlias.kb;
-    await createConcreteWriteIndex({
-      esClient: asInternalUser,
-      logger,
-      totalFieldsLimit: 10000,
-      indexPatterns: {
-        alias: kbAliasName,
-        pattern: `${kbAliasName}*`,
-        basePattern: `${kbAliasName}*`,
-        name: resourceNames.concreteWriteIndexName.kb,
-        template: resourceNames.indexTemplate.kb,
-      },
-      dataStreamAdapter: getDataStreamAdapter({ useDataStreamForAlerts: false }),
-    });
+    // `createConcreteWriteIndex` will create the write index, or update the index mappings if the index already exists
+    // only invoke `createConcreteWriteIndex` if the write index does not exist or the inferenceId in the component template is the same as the one in the write index
+    if (!writeIndexInferenceId || writeIndexInferenceId === componentTemplateInferenceId) {
+      const kbAliasName = resourceNames.writeIndexAlias.kb;
+      await createConcreteWriteIndex({
+        esClient: asInternalUser,
+        logger,
+        totalFieldsLimit: 10000,
+        indexPatterns: {
+          alias: kbAliasName,
+          pattern: `${kbAliasName}*`,
+          basePattern: `${kbAliasName}*`,
+          name: resourceNames.concreteWriteIndexName.kb,
+          template: resourceNames.indexTemplate.kb,
+        },
+        dataStreamAdapter: getDataStreamAdapter({ useDataStreamForAlerts: false }),
+      });
+    }
 
     logger.info('Successfully set up knowledge base index assets');
   } catch (error) {

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts

@@ -64,7 +64,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     it('has an index created in 8.18', async () => {
       await retry.try(async () => {
         const indexVersion = await getKbIndexCreatedVersion(es);
-        expect(indexVersion).to.be('8.18.0');
+        expect(indexVersion).to.contain('8.18.0'); // should match both '8.18.0-8.18.1' and '8.18.0': https://github.com/elastic/kibana/issues/220599
       });
     });
 

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts

@@ -31,13 +31,14 @@ import {
   deleteInferenceEndpoint,
   deleteModel,
   importModel,
+  startModelDeployment,
 } from '../utils/model_and_inference';
 import { animalSampleDocs } from '../utils/sample_docs';
 
 export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
   const es = getService('es');
-  const ml = getService('ml');
   const log = getService('log');
+  const retry = getService('retry');
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
 
   type KnowledgeBaseEsEntry = Awaited<ReturnType<typeof getKnowledgeBaseEntriesFromEs>>[0];
@@ -54,40 +55,51 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     let e5WriteIndex: string;
 
     before(async () => {
-      await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
-      await createTinyElserInferenceEndpoint(getService, { inferenceId: TINY_ELSER_INFERENCE_ID });
-      await setupKnowledgeBase(getService, TINY_ELSER_INFERENCE_ID);
-      await waitForKnowledgeBaseReady(getService);
-
-      // ingest documents
-      await addSampleDocsToInternalKb(getService, animalSampleDocs);
-
-      elserEntriesFromApi = (
-        await getKnowledgeBaseEntriesFromApi({ observabilityAIAssistantAPIClient })
-      ).body.entries;
-
-      elserEntriesFromEs = await getKnowledgeBaseEntriesFromEs(es);
-      elserInferenceId = await getInferenceIdFromWriteIndex({ asInternalUser: es });
-      elserWriteIndex = await getConcreteWriteIndexFromAlias(es);
-
-      // setup KB with E5-like model
-      await importModel(getService, { modelId: TINY_TEXT_EMBEDDING_MODEL_ID });
-      await ml.api.startTrainedModelDeploymentES(TINY_TEXT_EMBEDDING_MODEL_ID);
-      await createTinyTextEmbeddingInferenceEndpoint(getService, {
-        inferenceId: TINY_TEXT_EMBEDDING_INFERENCE_ID,
-      });
-      await setupKnowledgeBase(getService, TINY_TEXT_EMBEDDING_INFERENCE_ID);
+      await retry.try(async () => {
+        await restoreIndexAssets(getService);
+        await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
+        await createTinyElserInferenceEndpoint(getService, {
+          inferenceId: TINY_ELSER_INFERENCE_ID,
+        });
+        await setupKnowledgeBase(getService, TINY_ELSER_INFERENCE_ID);
+        await waitForKnowledgeBaseReady(getService);
+
+        // ingest documents
+        await addSampleDocsToInternalKb(getService, animalSampleDocs);
+
+        elserEntriesFromApi = (
+          await getKnowledgeBaseEntriesFromApi({ observabilityAIAssistantAPIClient })
+        ).body.entries;
 
-      await waitForKnowledgeBaseIndex(getService, '.kibana-observability-ai-assistant-kb-000002');
-      await waitForKnowledgeBaseReady(getService);
+        elserEntriesFromEs = await getKnowledgeBaseEntriesFromEs(es);
+        elserInferenceId = await getInferenceIdFromWriteIndex({ asInternalUser: es });
+        elserWriteIndex = await getConcreteWriteIndexFromAlias(es);
+
+        // setup KB with E5-like model
+        await importModel(getService, { modelId: TINY_TEXT_EMBEDDING_MODEL_ID });
+        await startModelDeployment(getService, { modelId: TINY_TEXT_EMBEDDING_MODEL_ID });
+
+        await createTinyTextEmbeddingInferenceEndpoint(getService, {
+          inferenceId: TINY_TEXT_EMBEDDING_INFERENCE_ID,
+        });
+        await setupKnowledgeBase(getService, TINY_TEXT_EMBEDDING_INFERENCE_ID);
 
-      e5EntriesFromApi = (
-        await getKnowledgeBaseEntriesFromApi({ observabilityAIAssistantAPIClient })
-      ).body.entries;
+        await waitForKnowledgeBaseIndex(getService, '.kibana-observability-ai-assistant-kb-000002');
+        await waitForKnowledgeBaseReady(getService);
 
-      e5EntriesFromEs = await getKnowledgeBaseEntriesFromEs(es);
-      e5InferenceId = await getInferenceIdFromWriteIndex({ asInternalUser: es });
-      e5WriteIndex = await getConcreteWriteIndexFromAlias(es);
+        e5EntriesFromApi = (
+          await getKnowledgeBaseEntriesFromApi({ observabilityAIAssistantAPIClient })
+        ).body.entries;
+
+        e5EntriesFromEs = await getKnowledgeBaseEntriesFromEs(es);
+        e5InferenceId = await getInferenceIdFromWriteIndex({ asInternalUser: es });
+        e5WriteIndex = await getConcreteWriteIndexFromAlias(es);
+
+        // retry until the following assertions pass
+        expect(elserWriteIndex).to.be(`${resourceNames.writeIndexAlias.kb}-000001`);
+        expect(e5WriteIndex).to.be(`${resourceNames.writeIndexAlias.kb}-000002`);
+        expect(e5InferenceId).to.be(TINY_TEXT_EMBEDDING_INFERENCE_ID);
+      });
     });
 
     after(async () => {
@@ -135,15 +147,15 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
 
     describe('when model is changed to E5', () => {
-      it('has increments the index name', async () => {
+      it('increments the index name', async () => {
         expect(e5WriteIndex).to.be(`${resourceNames.writeIndexAlias.kb}-000002`);
       });
 
-      it('returns the same entries from the API', async () => {
+      it('still returns the same entries from the API', async () => {
         expect(e5EntriesFromApi).to.eql(elserEntriesFromApi);
       });
 
-      it('has updates the inference id', async () => {
+      it('updates the inference id', async () => {
         expect(e5InferenceId).to.be(TINY_TEXT_EMBEDDING_INFERENCE_ID);
       });
 

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts

@@ -57,10 +57,21 @@ export async function waitForKnowledgeBaseReady(
 
   await retry.tryForTime(5 * 60 * 1000, async () => {
     log.debug(`Waiting for knowledge base to be ready...`);
-    const res = await getKnowledgeBaseStatus(observabilityAIAssistantAPIClient);
-    expect(res.status).to.be(200);
-    expect(res.body.kbState).to.be(KnowledgeBaseState.READY);
-    expect(res.body.isReIndexing).to.be(false);
+    const { body, status } = await getKnowledgeBaseStatus(observabilityAIAssistantAPIClient);
+
+    const { kbState, isReIndexing, concreteWriteIndex, currentInferenceId } = body;
+    if (status !== 200) {
+      log.warning(`Knowledge base is not ready yet:
+        Status code: ${status}
+        State: ${kbState}
+        isReIndexing: ${isReIndexing}
+        concreteWriteIndex: ${concreteWriteIndex}
+        currentInferenceId: ${currentInferenceId}`);
+    }
+
+    expect(status).to.be(200);
+    expect(kbState).to.be(KnowledgeBaseState.READY);
+    expect(isReIndexing).to.be(false);
     log.debug(`Knowledge base is in ready state.`);
   });
 }
@@ -77,12 +88,20 @@ export async function setupKnowledgeBase(
   log.debug(
     `Setting up knowledge base with inference endpoint = "${TINY_ELSER_INFERENCE_ID}", concreteWriteIndex = ${statusResult.body.concreteWriteIndex}, currentInferenceId = ${statusResult.body.currentInferenceId}, isReIndexing = ${statusResult.body.isReIndexing}`
   );
-  return observabilityAIAssistantAPIClient.admin({
+  const { body, status } = await observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
     params: {
       query: { inference_id: inferenceId, wait_until_complete: true },
     },
   });
+
+  if (status !== 200) {
+    log.warning(`Failed to setup knowledge base:
+      Status code: ${status}
+      Body: ${JSON.stringify(body, null, 2)}`);
+  }
+
+  return { body, status };
 }
 
 export async function addSampleDocsToInternalKb(

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts

@@ -8,6 +8,7 @@
 import { Client, errors } from '@elastic/elasticsearch';
 import { ToolingLog } from '@kbn/tooling-log';
 import { InferenceTaskType } from '@elastic/elasticsearch/lib/api/types';
+import pRetry from 'p-retry';
 import { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
 import { SUPPORTED_TRAINED_MODELS } from '../../../../../../functional/services/ml/api';
 import { setupKnowledgeBase, waitForKnowledgeBaseReady } from './knowledge_base';
@@ -39,7 +40,12 @@ export async function importModel(
   try {
     await ml.api.importTrainedModel(modelId, modelId, config);
   } catch (error) {
-    if (error.message.includes('resource_already_exists_exception')) {
+    if (
+      error.message.includes('resource_already_exists_exception') ||
+      error.message.includes(
+        'the model id is the same as the deployment id of a current model deployment'
+      )
+    ) {
       log.info(`Model "${modelId}" is already imported. Skipping import.`);
       return;
     }
@@ -49,6 +55,33 @@ export async function importModel(
   }
 }
 
+export async function startModelDeployment(
+  getService: DeploymentAgnosticFtrProviderContext['getService'],
+  {
+    modelId,
+  }: {
+    modelId: typeof TINY_ELSER_MODEL_ID | typeof TINY_TEXT_EMBEDDING_MODEL_ID;
+  }
+) {
+  const ml = getService('ml');
+  const log = getService('log');
+
+  try {
+    await ml.api.startTrainedModelDeploymentES(modelId);
+  } catch (error) {
+    if (
+      error.message.includes(
+        'Could not start model deployment because an existing deployment with the same id'
+      )
+    ) {
+      log.info(`Model "${modelId}" is already started. Skipping starting deployment.`);
+      return;
+    }
+
+    throw error;
+  }
+}
+
 export async function setupTinyElserModelAndInferenceEndpoint(
   getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
@@ -144,36 +177,41 @@ export async function createInferenceEndpoint({
   modelId: string;
   taskType?: InferenceTaskType;
 }) {
-  try {
-    const res = await es.inference.put({
-      inference_id: inferenceId,
-      task_type: taskType,
-      inference_config: {
-        service: 'elasticsearch',
-        service_settings: {
-          model_id: modelId,
-          adaptive_allocations: { enabled: true, min_number_of_allocations: 1 },
-          num_threads: 1,
-        },
-        task_settings: {},
-      },
-    });
-
-    log.info(`Inference endpoint ${inferenceId} created.`);
-    return res;
-  } catch (error) {
-    if (
-      error instanceof errors.ResponseError &&
-      (error.body?.error?.type === 'resource_not_found_exception' ||
-        error.body?.error?.type === 'status_exception')
-    ) {
-      log.debug(`Inference endpoint "${inferenceId}" already exists. Skipping creation.`);
-      return;
-    }
-
-    log.error(`Error creating inference endpoint "${inferenceId}": ${error}`);
-    throw error;
-  }
+  return pRetry(
+    async () => {
+      try {
+        const res = await es.inference.put({
+          inference_id: inferenceId,
+          task_type: taskType,
+          inference_config: {
+            service: 'elasticsearch',
+            service_settings: {
+              model_id: modelId,
+              adaptive_allocations: { enabled: true, min_number_of_allocations: 1 },
+              num_threads: 1,
+            },
+            task_settings: {},
+          },
+        });
+
+        log.info(`Inference endpoint ${inferenceId} created.`);
+        return res;
+      } catch (error) {
+        if (
+          error instanceof errors.ResponseError &&
+          (error.body?.error?.type === 'resource_not_found_exception' ||
+            error.body?.error?.type === 'status_exception')
+        ) {
+          log.debug(`Inference endpoint "${inferenceId}" already exists. Skipping creation.`);
+          return;
+        }
+
+        log.error(`Error creating inference endpoint "${inferenceId}": ${error}`);
+        throw error;
+      }
+    },
+    { retries: 2 }
+  );
 }
 
 export async function deleteModel(