elastic · sorenlouv · Feb 17, 2025
@@ -11,6 +11,7 @@ export const config = schema.object({
   enabled: schema.boolean({ defaultValue: true }),
   scope: schema.maybe(schema.oneOf([schema.literal('observability'), schema.literal('search')])),
   enableKnowledgeBase: schema.boolean({ defaultValue: true }),
+  disableKbSemanticTextMigration: schema.boolean({ defaultValue: false }),
 });
 
 export type ObservabilityAIAssistantConfig = TypeOf<typeof config>;
@@ -31,7 +31,8 @@ import { registerFunctions } from './functions';
 import { recallRankingEvent } from './analytics/recall_ranking';
 import { initLangtrace } from './service/client/instrumentation/init_langtrace';
 import { aiAssistantCapabilities } from '../common/capabilities';
-import { registerMigrateKnowledgeBaseEntriesTask } from './service/task_manager_definitions/register_migrate_knowledge_base_entries_task';
+import { registerAndScheduleKbSemanticTextMigrationTask } from './service/task_manager_definitions/register_kb_semantic_text_migration_task';
+import { updateExistingIndexAssets } from './service/create_or_update_index_assets';
 
 export class ObservabilityAIAssistantPlugin
   implements
@@ -128,14 +129,22 @@ export class ObservabilityAIAssistantPlugin
       config: this.config,
     }));
 
-    registerMigrateKnowledgeBaseEntriesTask({
+    // Update existing index assets (mappings, templates, etc). This will not create assets if they do not exist.
+    updateExistingIndexAssets({ logger: this.logger.get('index_assets'), core }).catch((e) =>
+      this.logger.error(`Index assets could not be updated: ${e.message}`)
+    );
+
+    // register task to migrate knowledge base entries to include semantic_text field
+    registerAndScheduleKbSemanticTextMigrationTask({
       core,
       taskManager: plugins.taskManager,
-      logger: this.logger,
+      logger: this.logger.get('kb_semantic_text_migration_task'),
       config: this.config,
-    }).catch((e) => {
-      this.logger.error(`Knowledge base migration was not successfully: ${e.message}`);
-    });
+    }).catch((e) =>
+      this.logger.error(
+        `Knowledge base semantic_text migration task could not be registered: ${e.message}`
+      )
+    );
 
     service.register(registerFunctions);
 

@@ -10,9 +10,11 @@ import { connectorRoutes } from './connectors/route';
 import { conversationRoutes } from './conversations/route';
 import { functionRoutes } from './functions/route';
 import { knowledgeBaseRoutes } from './knowledge_base/route';
+import { topLevelRoutes } from './top_level/route';
 
 export function getGlobalObservabilityAIAssistantServerRouteRepository() {
   return {
+    ...topLevelRoutes,
     ...chatRoutes,
     ...conversationRoutes,
     ...connectorRoutes,

@@ -101,7 +101,7 @@ const resetKnowledgeBase = createObservabilityAIAssistantServerRoute({
 });
 
 const semanticTextMigrationKnowledgeBase = createObservabilityAIAssistantServerRoute({
-  endpoint: 'POST /internal/observability_ai_assistant/kb/semantic_text_migration',
+  endpoint: 'POST /internal/observability_ai_assistant/kb/migrations/kb_semantic_text',
   security: {
     authz: {
       requiredPrivileges: ['ai_assistant'],
@@ -114,7 +114,7 @@ const semanticTextMigrationKnowledgeBase = createObservabilityAIAssistantServerR
       throw notImplemented();
     }
 
-    return client.migrateKnowledgeBaseToSemanticText();
+    return client.reIndexKnowledgeBaseAndPopulateSemanticTextField();
   },
 });
 

@@ -0,0 +1,28 @@
+/*
+ * 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 { createOrUpdateIndexAssets } from '../../service/create_or_update_index_assets';
+import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
+
+const createOrUpdateIndexAssetsRoute = createObservabilityAIAssistantServerRoute({
+  endpoint: 'POST /internal/observability_ai_assistant/index_assets',
+  security: {
+    authz: {
+      requiredPrivileges: ['ai_assistant'],
+    },
+  },
+  handler: async (resources): Promise<void> => {
+    return createOrUpdateIndexAssets({
+      logger: resources.logger,
+      core: resources.plugins.core.setup,
+    });
+  },
+});
+
+export const topLevelRoutes = {
+  ...createOrUpdateIndexAssetsRoute,
+};
@@ -72,9 +72,9 @@ import { extractTokenCount } from './operators/extract_token_count';
 import { getGeneratedTitle } from './operators/get_generated_title';
 import { instrumentAndCountTokens } from './operators/instrument_and_count_tokens';
 import {
-  runSemanticTextKnowledgeBaseMigration,
-  scheduleSemanticTextMigration,
-} from '../task_manager_definitions/register_migrate_knowledge_base_entries_task';
+  reIndexKnowledgeBaseAndPopulateSemanticTextField,
+  scheduleKbSemanticTextMigrationTask,
+} from '../task_manager_definitions/register_kb_semantic_text_migration_task';
 import { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
 import { ObservabilityAIAssistantConfig } from '../../config';
 import { getElserModelId } from '../knowledge_base_service/get_elser_model_id';
@@ -660,12 +660,11 @@ export class ObservabilityAIAssistantClient {
 
     core
       .getStartServices()
-      .then(([_, pluginsStart]) => {
-        logger.debug('Schedule semantic text migration task');
-        return scheduleSemanticTextMigration(pluginsStart);
-      })
+      .then(([_, pluginsStart]) =>
+        scheduleKbSemanticTextMigrationTask({ taskManager: pluginsStart.taskManager, logger })
+      )
       .catch((error) => {
-        logger.error(`Failed to run semantic text migration task: ${error}`);
+        logger.error(`Failed to schedule semantic text migration task: ${error}`);
       });
 
     return res;
@@ -676,8 +675,8 @@ export class ObservabilityAIAssistantClient {
     return this.dependencies.knowledgeBaseService.reset(esClient);
   };
 
-  migrateKnowledgeBaseToSemanticText = () => {
-    return runSemanticTextKnowledgeBaseMigration({
+  reIndexKnowledgeBaseAndPopulateSemanticTextField = () => {
+    return reIndexKnowledgeBaseAndPopulateSemanticTextField({
       esClient: this.dependencies.esClient,
       logger: this.dependencies.logger,
       config: this.dependencies.config,

@@ -6,13 +6,39 @@
  */
 
 import { createConcreteWriteIndex, getDataStreamAdapter } from '@kbn/alerting-plugin/server';
-import type { CoreSetup, Logger } from '@kbn/core/server';
+import type { CoreSetup, ElasticsearchClient, Logger } from '@kbn/core/server';
 import type { ObservabilityAIAssistantPluginStartDependencies } from '../types';
 import { conversationComponentTemplate } from './conversation_component_template';
 import { kbComponentTemplate } from './kb_component_template';
 import { resourceNames } from '.';
 
-export async function setupConversationAndKbIndexAssets({
+export async function updateExistingIndexAssets({
+  logger,
+  core,
+}: {
+  logger: Logger;
+  core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
+}) {
+  const [coreStart] = await core.getStartServices();
+  const { asInternalUser } = coreStart.elasticsearch.client;
+
+  const hasKbIndex = await asInternalUser.indices.exists({
+    index: resourceNames.aliases.kb,
+  });
+
+  const hasConversationIndex = await asInternalUser.indices.exists({
+    index: resourceNames.aliases.conversations,
+  });
+
+  if (!hasKbIndex && !hasConversationIndex) {
+    logger.debug('Index assets do not exist. Aborting updating index assets');
+    return;
+  }
+
+  await createOrUpdateIndexAssets({ logger, core });
+}
+
+export async function createOrUpdateIndexAssets({
   logger,
   core,
 }: {
@@ -56,7 +82,7 @@ export async function setupConversationAndKbIndexAssets({
         alias: conversationAliasName,
         pattern: `${conversationAliasName}*`,
         basePattern: `${conversationAliasName}*`,
-        name: `${conversationAliasName}-000001`,
+        name: resourceNames.concreteIndexName.conversations,
         template: resourceNames.indexTemplate.conversations,
       },
       dataStreamAdapter: getDataStreamAdapter({ useDataStreamForAlerts: false }),
@@ -86,24 +112,36 @@ export async function setupConversationAndKbIndexAssets({
     });
 
     // Knowledge base: write index
-    const kbAliasName = resourceNames.aliases.kb;
-    await createConcreteWriteIndex({
-      esClient: asInternalUser,
-      logger,
-      totalFieldsLimit: 10000,
-      indexPatterns: {
-        alias: kbAliasName,
-        pattern: `${kbAliasName}*`,
-        basePattern: `${kbAliasName}*`,
-        name: `${kbAliasName}-000001`,
-        template: resourceNames.indexTemplate.kb,
-      },
-      dataStreamAdapter: getDataStreamAdapter({ useDataStreamForAlerts: false }),
-    });
+    await createKbConcreteIndex({ logger, esClient: coreStart.elasticsearch.client });
 
     logger.info('Successfully set up index assets');
   } catch (error) {
     logger.error(`Failed setting up index assets: ${error.message}`);
     logger.debug(error);
   }
 }
+
+export async function createKbConcreteIndex({
+  logger,
+  esClient,
+}: {
+  logger: Logger;
+  esClient: {
+    asInternalUser: ElasticsearchClient;
+  };
+}) {
+  const kbAliasName = resourceNames.aliases.kb;
+  return createConcreteWriteIndex({
+    esClient: esClient.asInternalUser,
+    logger,
+    totalFieldsLimit: 10000,
+    indexPatterns: {
+      alias: kbAliasName,
+      pattern: `${kbAliasName}*`,
+      basePattern: `${kbAliasName}*`,
+      name: resourceNames.concreteIndexName.kb,
+      template: resourceNames.indexTemplate.kb,
+    },
+    dataStreamAdapter: getDataStreamAdapter({ useDataStreamForAlerts: false }),
+  });
+}
@@ -17,7 +17,7 @@ import { ObservabilityAIAssistantClient } from './client';
 import { KnowledgeBaseService } from './knowledge_base_service';
 import type { RegistrationCallback, RespondFunctionResources } from './types';
 import { ObservabilityAIAssistantConfig } from '../config';
-import { setupConversationAndKbIndexAssets } from './setup_conversation_and_kb_index_assets';
+import { createOrUpdateIndexAssets } from './create_or_update_index_assets';
 
 function getResourceName(resource: string) {
   return `.kibana-observability-ai-assistant-${resource}`;
@@ -40,11 +40,15 @@ export const resourceNames = {
     conversations: getResourceName('index-template-conversations'),
     kb: getResourceName('index-template-kb'),
   },
+  concreteIndexName: {
+    conversations: getResourceName('conversations-000001'),
+    kb: getResourceName('kb-000001'),
+  },
 };
 
 const createIndexAssetsOnce = once(
   (logger: Logger, core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>) =>
-    pRetry(() => setupConversationAndKbIndexAssets({ logger, core }))
+    pRetry(() => createOrUpdateIndexAssets({ logger, core }))
 );
 
 export class ObservabilityAIAssistantService {

@@ -28,6 +28,11 @@ import {
 import { recallFromSearchConnectors } from './recall_from_search_connectors';
 import { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
 import { ObservabilityAIAssistantConfig } from '../../config';
+import {
+  isKnowledgeBaseIndexWriteBlocked,
+  isSemanticTextUnsupportedError,
+} from './reindex_knowledge_base';
+import { scheduleKbSemanticTextMigrationTask } from '../task_manager_definitions/register_kb_semantic_text_migration_task';
 
 interface Dependencies {
   core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
@@ -406,7 +411,9 @@ export class KnowledgeBaseService {
     }
 
     try {
-      await this.dependencies.esClient.asInternalUser.index({
+      await this.dependencies.esClient.asInternalUser.index<
+        Omit<KnowledgeBaseEntry, 'id'> & { namespace: string }
+      >({
         index: resourceNames.aliases.kb,
         id,
         document: {
@@ -418,10 +425,40 @@ export class KnowledgeBaseService {
         },
         refresh: 'wait_for',
       });
+      this.dependencies.logger.debug(`Entry added to knowledge base`);
     } catch (error) {
+      this.dependencies.logger.debug(`Failed to add entry to knowledge base ${error}`);
       if (isInferenceEndpointMissingOrUnavailable(error)) {
         throwKnowledgeBaseNotReady(error.body);
       }
+
+      if (isSemanticTextUnsupportedError(error)) {
+        this.dependencies.core
+          .getStartServices()
+          .then(([_, pluginsStart]) => {
+            return scheduleKbSemanticTextMigrationTask({
+              taskManager: pluginsStart.taskManager,
+              logger: this.dependencies.logger,
+              runSoon: true,
+            });
+          })
+          .catch((e) => {
+            this.dependencies.logger.error(
+              `Failed to schedule knowledge base semantic text migration task: ${e}`
+            );
+          });
+
+        throw serverUnavailable(
+          'The knowledge base is currently being re-indexed. Please try again later'
+        );
+      }
+
+      if (isKnowledgeBaseIndexWriteBlocked(error)) {
+        throw new Error(
+          `Writes to the knowledge base are currently blocked due to an Elasticsearch write index block. This is most likely due to an ongoing re-indexing operation. Please try again later. Error: ${error.message}`
+        );
+      }
+
       throw error;
     }
   };