[8.12] [Obs AI Assistant] Make KB use recommended ELSER model (#173543)…

… (#173937) # Backport This will backport the following commits from `main` to `8.12`: - [[Obs AI Assistant] Make KB use recommended ELSER model (#173543)](#173543)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Milton Hultgren <[email protected]>
elastic · Dec 22, 2023 · d7417f9 · d7417f9
1 parent 37634be
commit d7417f9
Show file tree

Hide file tree

Showing 10 changed files with 80 additions and 23 deletions.
diff --git a/x-pack/plugins/ml/public/application/services/get_shared_ml_services.ts b/x-pack/plugins/ml/public/application/services/get_shared_ml_services.ts
@@ -21,5 +21,6 @@ export function getMlSharedServices(httpStart: HttpStart) {
 
   return {
     elasticModels: new ElasticModels(mlApiServices.trainedModels),
+    mlApiServices,
   };
 }
diff --git a/x-pack/plugins/ml/public/plugin.ts b/x-pack/plugins/ml/public/plugin.ts
@@ -69,6 +69,7 @@ import {
 } from '../common/constants/app';
 import type { MlCapabilities } from './shared';
 import { ElasticModels } from './application/services/elastic_models_service';
+import type { MlApiServices } from './application/services/ml_api_service';
 
 export interface MlStartDependencies {
   cases?: CasesUiStart;
@@ -283,7 +284,11 @@ export class MlPlugin implements Plugin<MlPluginSetup, MlPluginStart> {
   start(
     core: CoreStart,
     deps: MlStartDependencies
-  ): { locator?: LocatorPublic<MlLocatorParams>; elasticModels?: ElasticModels } {
+  ): {
+    locator?: LocatorPublic<MlLocatorParams>;
+    elasticModels?: ElasticModels;
+    mlApi?: MlApiServices;
+  } {
     setDependencyCache({
       docLinks: core.docLinks!,
       basePath: core.http.basePath,
@@ -295,6 +300,7 @@ export class MlPlugin implements Plugin<MlPluginSetup, MlPluginStart> {
     return {
       locator: this.locator,
       elasticModels: this.sharedMlServices?.elasticModels,
+      mlApi: this.sharedMlServices?.mlApiServices,
     };
   }
 

diff --git a/x-pack/plugins/observability_ai_assistant/kibana.jsonc b/x-pack/plugins/observability_ai_assistant/kibana.jsonc
@@ -20,7 +20,8 @@
       "share",
       "taskManager",
       "triggersActionsUi",
-      "dataViews"
+      "dataViews",
+      "ml"
     ],
     "requiredBundles": [ "kibanaReact", "kibanaUtils"],
     "optionalPlugins": [],

diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_knowledge_base.tsx b/x-pack/plugins/observability_ai_assistant/public/hooks/use_knowledge_base.tsx
@@ -30,6 +30,9 @@ export interface UseKnowledgeBaseResult {
 export function useKnowledgeBase(): UseKnowledgeBaseResult {
   const {
     notifications: { toasts },
+    plugins: {
+      start: { ml },
+    },
   } = useKibana().services;
   const service = useObservabilityAIAssistant();
 
@@ -53,6 +56,7 @@ export function useKnowledgeBase(): UseKnowledgeBaseResult {
         .callApi('POST /internal/observability_ai_assistant/kb/setup', {
           signal: null,
         })
+        .then(() => ml.mlApi?.savedObjects.syncSavedObjects())
         .then(() => {
           status.refresh();
         })
@@ -82,5 +86,5 @@ export function useKnowledgeBase(): UseKnowledgeBaseResult {
       isInstalling,
       installError,
     };
-  }, [status, isInstalling, installError, service, toasts]);
+  }, [status, isInstalling, installError, service, ml.mlApi?.savedObjects, toasts]);
 }
diff --git a/x-pack/plugins/observability_ai_assistant/public/types.ts b/x-pack/plugins/observability_ai_assistant/public/types.ts
@@ -30,6 +30,7 @@ import type { LicensingPluginStart, ILicense } from '@kbn/licensing-plugin/publi
 import type { SharePluginStart } from '@kbn/share-plugin/public';
 import { ForwardRefExoticComponent, RefAttributes } from 'react';
 import { WithSuspenseExtendedDeps } from '@kbn/shared-ux-utility';
+import { MlPluginSetup, MlPluginStart } from '@kbn/ml-plugin/public';
 import type {
   ContextDefinition,
   FunctionDefinition,
@@ -106,6 +107,7 @@ export interface ObservabilityAIAssistantPluginSetupDependencies {
   observabilityShared: ObservabilitySharedPluginSetup;
   security: SecurityPluginSetup;
   triggersActionsUi: TriggersAndActionsUIPublicPluginSetup;
+  ml: MlPluginSetup;
 }
 
 export interface ObservabilityAIAssistantPluginStartDependencies {
@@ -117,6 +119,7 @@ export interface ObservabilityAIAssistantPluginStartDependencies {
   security: SecurityPluginStart;
   share: SharePluginStart;
   triggersActionsUi: TriggersAndActionsUIPublicPluginStart;
+  ml: MlPluginStart;
 }
 
 export interface ObservabilityAIAssistantPluginSetup {}

diff --git a/x-pack/plugins/observability_ai_assistant/server/plugin.ts b/x-pack/plugins/observability_ai_assistant/server/plugin.ts
@@ -12,13 +12,14 @@ import {
   Plugin,
   PluginInitializerContext,
 } from '@kbn/core/server';
-import { mapValues } from 'lodash';
+import { mapValues, once } from 'lodash';
 import { i18n } from '@kbn/i18n';
 import {
   CONNECTOR_TOKEN_SAVED_OBJECT_TYPE,
   ACTION_SAVED_OBJECT_TYPE,
   ACTION_TASK_PARAMS_SAVED_OBJECT_TYPE,
 } from '@kbn/actions-plugin/server/constants/saved_objects';
+import { firstValueFrom } from 'rxjs';
 import { OBSERVABILITY_AI_ASSISTANT_FEATURE_ID } from '../common/feature';
 import type { ObservabilityAIAssistantConfig } from './config';
 import { registerServerRoutes } from './routes/register_routes';
@@ -105,10 +106,34 @@ export class ObservabilityAIAssistantPlugin
       };
     }) as ObservabilityAIAssistantRouteHandlerResources['plugins'];
 
+    const getModelId = once(async () => {
+      // Using once to make sure the same model ID is used during service init and Knowledge base setup
+
+      try {
+        // Wait for the ML plugin's dependency on the internal saved objects client to be ready
+        const [_, pluginsStart] = await core.getStartServices();
+
+        // Wait for the license to be available so the ML plugin's guards pass once we ask for ELSER stats
+        await firstValueFrom(pluginsStart.licensing.license$);
+
+        const elserModelDefinition = await plugins.ml
+          .trainedModelsProvider({} as any, {} as any) // request, savedObjectsClient (but we fake it to use the internal user)
+          .getELSER();
+
+        return elserModelDefinition.model_id;
+      } catch (error) {
+        this.logger.error(`Failed to resolve ELSER model definition: ${error}`);
+
+        // Fallback to ELSER v2
+        return '.elser_model_2';
+      }
+    });
+
     const service = (this.service = new ObservabilityAIAssistantService({
       logger: this.logger.get('service'),
       core,
       taskManager: plugins.taskManager,
+      getModelId,
     }));
 
     service.register(registerFunctions);

diff --git a/x-pack/plugins/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
@@ -66,8 +66,6 @@ export function createResourceNamesMap() {
   };
 }
 
-export const ELSER_MODEL_ID = '.elser_model_2';
-
 export const INDEX_QUEUED_DOCUMENTS_TASK_ID = 'observabilityAIAssistant:indexQueuedDocumentsTask';
 
 export const INDEX_QUEUED_DOCUMENTS_TASK_TYPE = INDEX_QUEUED_DOCUMENTS_TASK_ID + 'Type';
@@ -84,6 +82,7 @@ type KnowledgeBaseEntryRequest = { id: string; labels?: Record<string, string> }
 export class ObservabilityAIAssistantService {
   private readonly core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
   private readonly logger: Logger;
+  private readonly getModelId: () => Promise<string>;
   private kbService?: KnowledgeBaseService;
 
   private readonly resourceNames: ObservabilityAIAssistantResourceNames = createResourceNamesMap();
@@ -94,13 +93,16 @@ export class ObservabilityAIAssistantService {
     logger,
     core,
     taskManager,
+    getModelId,
   }: {
     logger: Logger;
     core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
     taskManager: TaskManagerSetupContract;
+    getModelId: () => Promise<string>;
   }) {
     this.core = core;
     this.logger = logger;
+    this.getModelId = getModelId;
 
     taskManager.registerTaskDefinitions({
       [INDEX_QUEUED_DOCUMENTS_TASK_TYPE]: {
@@ -132,6 +134,8 @@ export class ObservabilityAIAssistantService {
     try {
       const [coreStart, pluginsStart] = await this.core.getStartServices();
 
+      const elserModelId = await this.getModelId();
+
       const esClient = coreStart.elasticsearch.client.asInternalUser;
 
       await esClient.cluster.putComponentTemplate({
@@ -153,7 +157,7 @@ export class ObservabilityAIAssistantService {
           },
           mappings: {
             _meta: {
-              model: ELSER_MODEL_ID,
+              model: elserModelId,
             },
           },
         },
@@ -186,7 +190,7 @@ export class ObservabilityAIAssistantService {
         processors: [
           {
             inference: {
-              model_id: ELSER_MODEL_ID,
+              model_id: elserModelId,
               target_field: 'ml',
               field_map: {
                 text: 'text_field',
@@ -237,6 +241,7 @@ export class ObservabilityAIAssistantService {
         esClient,
         resources: this.resourceNames,
         taskManagerStart: pluginsStart.taskManager,
+        getModelId: this.getModelId,
       });
 
       this.logger.info('Successfully set up index assets');

diff --git a/x-pack/plugins/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/knowledge_base_service/index.ts
@@ -14,11 +14,7 @@ import pLimit from 'p-limit';
 import pRetry from 'p-retry';
 import { map, orderBy } from 'lodash';
 import { encode } from 'gpt-tokenizer';
-import {
-  ELSER_MODEL_ID,
-  INDEX_QUEUED_DOCUMENTS_TASK_ID,
-  INDEX_QUEUED_DOCUMENTS_TASK_TYPE,
-} from '..';
+import { INDEX_QUEUED_DOCUMENTS_TASK_ID, INDEX_QUEUED_DOCUMENTS_TASK_TYPE } from '..';
 import { KnowledgeBaseEntry, KnowledgeBaseEntryRole } from '../../../common/types';
 import type { ObservabilityAIAssistantResourceNames } from '../types';
 import { getAccessQuery } from '../util/get_access_query';
@@ -29,6 +25,7 @@ interface Dependencies {
   resources: ObservabilityAIAssistantResourceNames;
   logger: Logger;
   taskManagerStart: TaskManagerStartContract;
+  getModelId: () => Promise<string>;
 }
 
 export interface RecalledEntry {
@@ -81,13 +78,15 @@ export class KnowledgeBaseService {
   }
 
   setup = async () => {
+    const elserModelId = await this.dependencies.getModelId();
+
     const retryOptions = { factor: 1, minTimeout: 10000, retries: 12 };
 
     const installModel = async () => {
       this.dependencies.logger.info('Installing ELSER model');
       await this.dependencies.esClient.ml.putTrainedModel(
         {
-          model_id: ELSER_MODEL_ID,
+          model_id: elserModelId,
           input: {
             field_names: ['text_field'],
           },
@@ -101,7 +100,7 @@ export class KnowledgeBaseService {
 
     const getIsModelInstalled = async () => {
       const getResponse = await this.dependencies.esClient.ml.getTrainedModels({
-        model_id: ELSER_MODEL_ID,
+        model_id: elserModelId,
         include: 'definition_status',
       });
 
@@ -132,7 +131,7 @@ export class KnowledgeBaseService {
 
     try {
       await this.dependencies.esClient.ml.startTrainedModelDeployment({
-        model_id: ELSER_MODEL_ID,
+        model_id: elserModelId,
         wait_for: 'fully_allocated',
       });
     } catch (error) {
@@ -145,7 +144,7 @@ export class KnowledgeBaseService {
 
     await pRetry(async () => {
       const response = await this.dependencies.esClient.ml.getTrainedModelsStats({
-        model_id: ELSER_MODEL_ID,
+        model_id: elserModelId,
       });
 
       if (
@@ -269,9 +268,11 @@ export class KnowledgeBaseService {
   }
 
   status = async () => {
+    const elserModelId = await this.dependencies.getModelId();
+
     try {
       const modelStats = await this.dependencies.esClient.ml.getTrainedModelsStats({
-        model_id: ELSER_MODEL_ID,
+        model_id: elserModelId,
       });
       const elserModelStats = modelStats.trained_model_stats[0];
       const deploymentState = elserModelStats.deployment_stats?.state;
@@ -281,13 +282,13 @@ export class KnowledgeBaseService {
         ready: deploymentState === 'started' && allocationState === 'fully_allocated',
         deployment_state: deploymentState,
         allocation_state: allocationState,
-        model_name: ELSER_MODEL_ID,
+        model_name: elserModelId,
       };
     } catch (error) {
       return {
         error: error instanceof errors.ResponseError ? error.body.error : String(error),
         ready: false,
-        model_name: ELSER_MODEL_ID,
+        model_name: elserModelId,
       };
     }
   };
@@ -435,7 +436,7 @@ export class KnowledgeBaseService {
   }): Promise<{
     entries: RecalledEntry[];
   }> => {
-    const modelId = ELSER_MODEL_ID;
+    const modelId = await this.dependencies.getModelId();
 
     const [documentsFromKb, documentsFromConnectors] = await Promise.all([
       this.recallFromKnowledgeBase({

diff --git a/x-pack/plugins/observability_ai_assistant/server/types.ts b/x-pack/plugins/observability_ai_assistant/server/types.ts
@@ -17,7 +17,12 @@ import type {
   TaskManagerSetupContract,
   TaskManagerStartContract,
 } from '@kbn/task-manager-plugin/server';
-import { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
+import type {
+  DataViewsServerPluginSetup,
+  DataViewsServerPluginStart,
+} from '@kbn/data-views-plugin/server';
+import type { MlPluginSetup, MlPluginStart } from '@kbn/ml-plugin/server';
+import type { LicensingPluginSetup, LicensingPluginStart } from '@kbn/licensing-plugin/server';
 import { ObservabilityAIAssistantService } from './service';
 
 export interface ObservabilityAIAssistantPluginSetup {
@@ -39,11 +44,16 @@ export interface ObservabilityAIAssistantPluginSetupDependencies {
   security: SecurityPluginSetup;
   features: FeaturesPluginSetup;
   taskManager: TaskManagerSetupContract;
+  dataViews: DataViewsServerPluginSetup;
+  ml: MlPluginSetup;
+  licensing: LicensingPluginSetup;
 }
 export interface ObservabilityAIAssistantPluginStartDependencies {
   actions: ActionsPluginStart;
   security: SecurityPluginStart;
   features: FeaturesPluginStart;
   taskManager: TaskManagerStartContract;
   dataViews: DataViewsServerPluginStart;
+  ml: MlPluginStart;
+  licensing: LicensingPluginStart;
 }
diff --git a/x-pack/plugins/observability_ai_assistant/tsconfig.json b/x-pack/plugins/observability_ai_assistant/tsconfig.json
@@ -55,7 +55,8 @@
     "@kbn/dev-cli-runner",
     "@kbn/core-analytics-browser",
     "@kbn/core-http-browser",
-    "@kbn/security-plugin-types-common"
+    "@kbn/security-plugin-types-common",
+    "@kbn/ml-plugin"
   ],
   "exclude": ["target/**/*"]
 }