[8.19] [Obs AI Assistant] Improve flaky recall tests (#220638) | [Obs AI Assistant] Do not attempt to overwrite index mappings if inference id in component template has changed (#220739)

x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts

@@ -50,9 +50,10 @@ const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
 const setupKnowledgeBase = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
   params: t.type({
-    query: t.type({
-      inference_id: t.string,
-    }),
+    query: t.intersection([
+      t.type({ inference_id: t.string }),
+      t.partial({ wait_until_complete: toBooleanRt }),
+    ]),
   }),
   security: {
     authz: {
@@ -67,8 +68,9 @@ const setupKnowledgeBase = createObservabilityAIAssistantServerRoute({
     nextInferenceId: string;
   }> => {
     const client = await resources.service.getClient({ request: resources.request });
-    const { inference_id: inferenceId } = resources.params.query;
-    return client.setupKnowledgeBase(inferenceId);
+    const { inference_id: inferenceId, wait_until_complete: waitUntilComplete } =
+      resources.params.query;
+    return client.setupKnowledgeBase(inferenceId, waitUntilComplete);
   },
 });
 

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts

@@ -668,7 +668,8 @@ export class ObservabilityAIAssistantClient {
   };
 
   setupKnowledgeBase = async (
-    nextInferenceId: string
+    nextInferenceId: string,
+    waitUntilComplete: boolean = false
   ): Promise<{
     reindex: boolean;
     currentInferenceId: string | undefined;
@@ -697,7 +698,7 @@ export class ObservabilityAIAssistantClient {
       inferenceId: nextInferenceId,
     });
 
-    waitForKbModel({
+    const kbSetupPromise = waitForKbModel({
       core: this.dependencies.core,
       esClient,
       logger,
@@ -732,6 +733,10 @@ export class ObservabilityAIAssistantClient {
         }
       });
 
+    if (waitUntilComplete) {
+      await kbSetupPromise;
+    }
+
     return { reindex: true, currentInferenceId, nextInferenceId };
   };
 

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/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts

@@ -416,7 +416,7 @@ export class KnowledgeBaseService {
     }
 
     try {
-      await this.dependencies.esClient.asInternalUser.index<
+      const indexResult = await this.dependencies.esClient.asInternalUser.index<
         Omit<KnowledgeBaseEntry, 'id'> & { namespace: string }
       >({
         index: resourceNames.writeIndexAlias.kb,
@@ -432,7 +432,7 @@ export class KnowledgeBaseService {
       });
 
       this.dependencies.logger.debug(
-        `Entry added to knowledge base. title = "${doc.title}", user = "${user?.name}, namespace = "${namespace}"`
+        `Entry added to knowledge base. title = "${doc.title}", user = "${user?.name}, namespace = "${namespace}", index = ${indexResult._index}, id = ${indexResult._id}`
       );
     } catch (error) {
       this.dependencies.logger.debug(`Failed to add entry to knowledge base ${error}`);

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts

@@ -0,0 +1,41 @@
+/*
+ * 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 { errors } from '@elastic/elasticsearch';
+import { ElasticsearchClient } from '@kbn/core/server';
+import { resourceNames } from '..';
+
+export async function addIndexWriteBlock({
+  esClient,
+  index,
+}: {
+  esClient: { asInternalUser: ElasticsearchClient };
+  index: string;
+}) {
+  await esClient.asInternalUser.indices.addBlock({ index, block: 'write' });
+}
+
+export function removeIndexWriteBlock({
+  esClient,
+  index,
+}: {
+  esClient: { asInternalUser: ElasticsearchClient };
+  index: string;
+}) {
+  return esClient.asInternalUser.indices.putSettings({
+    index,
+    body: { 'index.blocks.write': false },
+  });
+}
+
+export function isKnowledgeBaseIndexWriteBlocked(error: any) {
+  return (
+    error instanceof errors.ResponseError &&
+    error.message.includes(`cluster_block_exception`) &&
+    error.message.includes(resourceNames.writeIndexAlias.kb)
+  );
+}

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts

@@ -59,13 +59,6 @@ async function reIndexKnowledgeBase({
     `Re-indexing knowledge base from "${currentWriteIndexName}" to index "${nextWriteIndexName}"...`
   );
 
-  const reindexResponse = await esClient.asInternalUser.reindex({
-    source: { index: currentWriteIndexName },
-    dest: { index: nextWriteIndexName },
-    refresh: true,
-    wait_for_completion: false,
-  });
-
   // Point write index alias to the new index
   await updateKnowledgeBaseWriteIndexAlias({
     esClient,
@@ -74,6 +67,13 @@ async function reIndexKnowledgeBase({
     currentWriteIndexName,
   });
 
+  const reindexResponse = await esClient.asInternalUser.reindex({
+    source: { index: currentWriteIndexName },
+    dest: { index: nextWriteIndexName },
+    refresh: true,
+    wait_for_completion: false,
+  });
+
   const taskId = reindexResponse.task?.toString();
   if (taskId) {
     await waitForReIndexTaskToComplete({ esClient, taskId, logger });

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts

@@ -8,6 +8,7 @@
 import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
 import { Logger } from '@kbn/logging';
 import { resourceNames } from '..';
+import { addIndexWriteBlock, removeIndexWriteBlock } from './index_write_block_utils';
 
 export async function updateKnowledgeBaseWriteIndexAlias({
   esClient,
@@ -25,14 +26,26 @@ export async function updateKnowledgeBaseWriteIndexAlias({
   );
   const alias = resourceNames.writeIndexAlias.kb;
   try {
+    await addIndexWriteBlock({ esClient, index: currentWriteIndexName });
+    logger.debug(
+      `Added write block to "${currentWriteIndexName}". It is now read-only and writes are temporarily blocked.`
+    );
+
     await esClient.asInternalUser.indices.updateAliases({
       actions: [
         { remove: { index: currentWriteIndexName, alias } },
         { add: { index: nextWriteIndexName, alias, is_write_index: true } },
       ],
     });
   } catch (error) {
-    logger.error(`Failed to update write index alias: ${error.message}`);
+    await removeIndexWriteBlock({ esClient, index: currentWriteIndexName });
+    logger.error(
+      `Failed to update write index alias: ${error.message}. Reverting back to ${currentWriteIndexName}`
+    );
     throw error;
   }
+
+  logger.debug(
+    `Successfully updated write index alias to "${nextWriteIndexName}". Writes are now enabled again.`
+  );
 }

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