[Obs AI Assistant] Support querying semantic_text fields in search …

…connectors (#200184) This adds support for querying search connectors that have `semantic_text` fields. --------- Co-authored-by: kibanamachine <[email protected]>
elastic · Nov 18, 2024 · 9e4e8a5 · 9e4e8a5
1 parent 5de75c0
commit 9e4e8a5
Show file tree

Hide file tree

Showing 7 changed files with 297 additions and 209 deletions.
diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/plugin.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/plugin.ts
@@ -12,14 +12,13 @@ import {
   Plugin,
   PluginInitializerContext,
 } from '@kbn/core/server';
-import { mapValues, once } from 'lodash';
+import { mapValues } 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 { KibanaFeatureScope } from '@kbn/features-plugin/common';
 import { OBSERVABILITY_AI_ASSISTANT_FEATURE_ID } from '../common/feature';
 import type { ObservabilityAIAssistantConfig } from './config';
@@ -114,47 +113,10 @@ export class ObservabilityAIAssistantPlugin
       };
     }) as ObservabilityAIAssistantRouteHandlerResources['plugins'];
 
-    // Using once to make sure the same model ID is used during service init and Knowledge base setup
-    const getSearchConnectorModelId = once(async () => {
-      const defaultModelId = '.elser_model_2';
-      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
-      const license = await firstValueFrom(pluginsStart.licensing.license$);
-      if (!license.hasAtLeast('enterprise')) {
-        return defaultModelId;
-      }
-
-      try {
-        // Wait for the ML plugin's dependency on the internal saved objects client to be ready
-        const { ml } = await core.plugins.onSetup('ml');
-
-        if (!ml.found) {
-          throw new Error('Could not find ML plugin');
-        }
-
-        const elserModelDefinition = await (
-          ml.contract as {
-            trainedModelsProvider: (
-              request: {},
-              soClient: {}
-            ) => { getELSER: () => Promise<{ model_id: string }> };
-          }
-        )
-          .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}`);
-        return defaultModelId;
-      }
-    });
-
     const service = (this.service = new ObservabilityAIAssistantService({
       logger: this.logger.get('service'),
       core,
-      getSearchConnectorModelId,
-      enableKnowledgeBase: this.config.enableKnowledgeBase,
+      config: this.config,
     }));
 
     registerMigrateKnowledgeBaseEntriesTask({

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts
@@ -20,6 +20,7 @@ import { conversationComponentTemplate } from './conversation_component_template
 import { kbComponentTemplate } from './kb_component_template';
 import { KnowledgeBaseService } from './knowledge_base_service';
 import type { RegistrationCallback, RespondFunctionResources } from './types';
+import { ObservabilityAIAssistantConfig } from '../config';
 
 function getResourceName(resource: string) {
   return `.kibana-observability-ai-assistant-${resource}`;
@@ -47,27 +48,23 @@ export const resourceNames = {
 export class ObservabilityAIAssistantService {
   private readonly core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
   private readonly logger: Logger;
-  private readonly getSearchConnectorModelId: () => Promise<string>;
   private kbService?: KnowledgeBaseService;
-  private enableKnowledgeBase: boolean;
+  private config: ObservabilityAIAssistantConfig;
 
   private readonly registrations: RegistrationCallback[] = [];
 
   constructor({
     logger,
     core,
-    getSearchConnectorModelId,
-    enableKnowledgeBase,
+    config,
   }: {
     logger: Logger;
     core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
-    getSearchConnectorModelId: () => Promise<string>;
-    enableKnowledgeBase: boolean;
+    config: ObservabilityAIAssistantConfig;
   }) {
     this.core = core;
     this.logger = logger;
-    this.getSearchConnectorModelId = getSearchConnectorModelId;
-    this.enableKnowledgeBase = enableKnowledgeBase;
+    this.config = config;
 
     this.resetInit();
   }
@@ -166,12 +163,12 @@ export class ObservabilityAIAssistantService {
       });
 
       this.kbService = new KnowledgeBaseService({
+        core: this.core,
         logger: this.logger.get('kb'),
+        config: this.config,
         esClient: {
           asInternalUser,
         },
-        getSearchConnectorModelId: this.getSearchConnectorModelId,
-        enabled: this.enableKnowledgeBase,
       });
 
       this.logger.info('Successfully set up index assets');

diff --git a/...bility_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/...bility_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts
@@ -6,7 +6,7 @@
  */
 
 import { serverUnavailable } from '@hapi/boom';
-import type { ElasticsearchClient, IUiSettingsClient } from '@kbn/core/server';
+import type { CoreSetup, ElasticsearchClient, IUiSettingsClient } from '@kbn/core/server';
 import type { Logger } from '@kbn/logging';
 import { orderBy } from 'lodash';
 import { encode } from 'gpt-tokenizer';
@@ -26,15 +26,17 @@ import {
   getInferenceEndpoint,
   isInferenceEndpointMissingOrUnavailable,
 } from '../inference_endpoint';
-import { recallFromConnectors } from './recall_from_connectors';
+import { recallFromSearchConnectors } from './recall_from_search_connectors';
+import { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
+import { ObservabilityAIAssistantConfig } from '../../config';
 
 interface Dependencies {
+  core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
   esClient: {
     asInternalUser: ElasticsearchClient;
   };
   logger: Logger;
-  getSearchConnectorModelId: () => Promise<string>;
-  enabled: boolean;
+  config: ObservabilityAIAssistantConfig;
 }
 
 export interface RecalledEntry {
@@ -141,14 +143,13 @@ export class KnowledgeBaseService {
     esClient: { asCurrentUser: ElasticsearchClient; asInternalUser: ElasticsearchClient };
     uiSettingsClient: IUiSettingsClient;
   }): Promise<RecalledEntry[]> => {
-    if (!this.dependencies.enabled) {
+    if (!this.dependencies.config.enableKnowledgeBase) {
       return [];
     }
 
     this.dependencies.logger.debug(
       () => `Recalling entries from KB for queries: "${JSON.stringify(queries)}"`
     );
-    const modelId = await this.dependencies.getSearchConnectorModelId();
 
     const [documentsFromKb, documentsFromConnectors] = await Promise.all([
       this.recallFromKnowledgeBase({
@@ -162,11 +163,11 @@ export class KnowledgeBaseService {
         }
         throw error;
       }),
-      recallFromConnectors({
+      recallFromSearchConnectors({
         esClient,
         uiSettingsClient,
         queries,
-        modelId,
+        core: this.dependencies.core,
         logger: this.dependencies.logger,
       }).catch((error) => {
         this.dependencies.logger.debug('Error getting data from search indices');
@@ -214,7 +215,7 @@ export class KnowledgeBaseService {
     namespace: string,
     user?: { name: string }
   ): Promise<Array<Instruction & { public?: boolean }>> => {
-    if (!this.dependencies.enabled) {
+    if (!this.dependencies.config.enableKnowledgeBase) {
       return [];
     }
     try {
@@ -257,7 +258,7 @@ export class KnowledgeBaseService {
     sortBy?: string;
     sortDirection?: 'asc' | 'desc';
   }): Promise<{ entries: KnowledgeBaseEntry[] }> => {
-    if (!this.dependencies.enabled) {
+    if (!this.dependencies.config.enableKnowledgeBase) {
       return { entries: [] };
     }
     try {
@@ -330,7 +331,7 @@ export class KnowledgeBaseService {
     user?: { name: string; id?: string };
     namespace?: string;
   }) => {
-    if (!this.dependencies.enabled) {
+    if (!this.dependencies.config.enableKnowledgeBase) {
       return null;
     }
     const res = await this.dependencies.esClient.asInternalUser.search<KnowledgeBaseEntry>({
@@ -393,7 +394,7 @@ export class KnowledgeBaseService {
     user?: { name: string; id?: string };
     namespace?: string;
   }): Promise<void> => {
-    if (!this.dependencies.enabled) {
+    if (!this.dependencies.config.enableKnowledgeBase) {
       return;
     }
 
@@ -448,7 +449,7 @@ export class KnowledgeBaseService {
       errorMessage = error.message;
     });
 
-    const enabled = this.dependencies.enabled;
+    const enabled = this.dependencies.config.enableKnowledgeBase;
     if (!endpoint) {
       return { ready: false, enabled, errorMessage };
     }

diff --git a/...bservability_ai_assistant/server/service/knowledge_base_service/recall_from_connectors.ts b/...bservability_ai_assistant/server/service/knowledge_base_service/recall_from_connectors.ts