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/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 4eac156b73868..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(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); }); 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 6abc77856a1e5..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 @@ -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( 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 1d28bec819214..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'] ) { @@ -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(