diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts index b78dd82cc48e0..6596ca667dee5 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts +++ b/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); }, }); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts index 10b2dc9511fea..89a7ce88d7312 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts +++ b/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 }; }; diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/create_or_update_knowledge_base_index_assets.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/create_or_update_knowledge_base_index_assets.ts index 7358903685c19..be55b5cbf59b3 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/create_or_update_knowledge_base_index_assets.ts +++ b/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; @@ -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) { diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts index 4285a69c6f1ee..af96f8819c1f0 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts +++ b/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 & { 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}`); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts new file mode 100644 index 0000000000000..b3057c666d1d3 --- /dev/null +++ b/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) + ); +} diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts index 58b7b98d1e31d..8dd7cae3902fd 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts +++ b/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 }); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts index 3d893d6c2a6f3..44d91626a7343 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts +++ b/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,6 +26,11 @@ 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 } }, @@ -32,7 +38,14 @@ export async function updateKnowledgeBaseWriteIndexAlias({ ], }); } 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.` + ); } diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts index c37baa03c7c85..662024d3a0940 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts +++ b/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 }); }); diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts index 48d141b46e8a5..570469f3f2af7 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts +++ b/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>[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); }); diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts index 7e0a5c9466386..389bf18a952a9 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts @@ -65,11 +65,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon let body: Awaited>['body']; before(async () => { - // setup KB initially - await deployTinyElserAndSetupKb(getService); - - // setup KB with custom inference endpoint await createTinyElserInferenceEndpoint(getService, { inferenceId: CUSTOM_TINY_ELSER_INFERENCE_ID, }); @@ -120,7 +116,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon await createTinyElserInferenceEndpoint(getService, { inferenceId: customInferenceId, }); - await setupKnowledgeBase(observabilityAIAssistantAPIClient, customInferenceId); + await setupKnowledgeBase(getService, customInferenceId); await waitForKnowledgeBaseReady(getService); }); @@ -154,19 +150,14 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon } function setupKbAsAdmin(inferenceId: string) { - return observabilityAIAssistantAPIClient.admin({ - endpoint: 'POST /internal/observability_ai_assistant/kb/setup', - params: { - query: { inference_id: inferenceId }, - }, - }); + return setupKnowledgeBase(getService, inferenceId); } function setupKbAsViewer(inferenceId: string) { return observabilityAIAssistantAPIClient.viewer({ endpoint: 'POST /internal/observability_ai_assistant/kb/setup', params: { - query: { inference_id: inferenceId }, + query: { inference_id: inferenceId, wait_until_complete: true }, }, }); } diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts index 4cd1d9e4b0f3f..14ff7b8286b01 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts @@ -14,6 +14,7 @@ import { import { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context'; import type { ObservabilityAIAssistantApiClient } from '../../../../services/observability_ai_assistant_api'; import { TINY_ELSER_INFERENCE_ID } from './model_and_inference'; +import { getConcreteWriteIndexFromAlias } from './knowledge_base'; export async function runStartupMigrations( observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient @@ -67,9 +68,16 @@ export async function restoreIndexAssets( getService: DeploymentAgnosticFtrProviderContext['getService'] ) { const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); + const retry = getService('retry'); + const es = getService('es'); + const log = getService('log'); - await deleteIndexAssets(getService); - await createOrUpdateIndexAssets(observabilityAIAssistantAPIClient); + await retry.try(async () => { + log.debug('Restoring index assets'); + await deleteIndexAssets(getService); + await createOrUpdateIndexAssets(observabilityAIAssistantAPIClient); + expect(await getConcreteWriteIndexFromAlias(es)).to.be(resourceNames.concreteWriteIndexName.kb); + }); } export async function getComponentTemplate(es: Client) { diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts index 66a59bc93bc57..6d07ffaf2f224 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts @@ -40,6 +40,14 @@ export async function waitForKnowledgeBaseIndex( }); } +export async function getKnowledgeBaseStatus( + observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient +) { + return observabilityAIAssistantAPIClient.editor({ + endpoint: 'GET /internal/observability_ai_assistant/kb/status', + }); +} + export async function waitForKnowledgeBaseReady( getService: DeploymentAgnosticFtrProviderContext['getService'] ) { @@ -49,26 +57,51 @@ export async function waitForKnowledgeBaseReady( await retry.tryForTime(5 * 60 * 1000, async () => { log.debug(`Waiting for knowledge base to be ready...`); - const res = await observabilityAIAssistantAPIClient.editor({ - endpoint: 'GET /internal/observability_ai_assistant/kb/status', - }); - 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.`); }); } export async function setupKnowledgeBase( - observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient, + getService: DeploymentAgnosticFtrProviderContext['getService'], inferenceId: string ) { - return observabilityAIAssistantAPIClient.admin({ + const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); + const log = getService('log'); + + const statusResult = await getKnowledgeBaseStatus(observabilityAIAssistantAPIClient); + + 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}` + ); + const { body, status } = await observabilityAIAssistantAPIClient.admin({ endpoint: 'POST /internal/observability_ai_assistant/kb/setup', params: { - query: { inference_id: inferenceId }, + 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( @@ -77,6 +110,8 @@ export async function addSampleDocsToInternalKb( ) { const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); const es = getService('es'); + const log = getService('log'); + const retry = getService('retry'); await observabilityAIAssistantAPIClient.editor({ endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import', @@ -88,6 +123,14 @@ export async function addSampleDocsToInternalKb( }); await refreshKbIndex(es); + + await retry.try(async () => { + const itemsInKb = await getKnowledgeBaseEntriesFromEs(es); + log.debug( + `Waiting for at least ${sampleDocs.length} docs to be available for search in KB. Currently ${itemsInKb.length} docs available.` + ); + expect(itemsInKb.length >= sampleDocs.length).to.be(true); + }); } // refresh the index to make sure the documents are searchable diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts index 2a610e0e45296..52dcf7d36a4b0 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts +++ b/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'] ) { @@ -98,16 +131,9 @@ export function createTinyTextEmbeddingInferenceEndpoint( export async function deployTinyElserAndSetupKb( getService: DeploymentAgnosticFtrProviderContext['getService'] ) { - const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); - const log = getService('log'); - await setupTinyElserModelAndInferenceEndpoint(getService); - log.debug(`Setting up knowledge base with inference endpoint "${TINY_ELSER_INFERENCE_ID}"`); - const { status, body } = await setupKnowledgeBase( - observabilityAIAssistantAPIClient, - TINY_ELSER_INFERENCE_ID - ); + const { status, body } = await setupKnowledgeBase(getService, TINY_ELSER_INFERENCE_ID); await waitForKnowledgeBaseReady(getService); return { status, body }; @@ -151,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(