[Obs AI Assistant] Only show ELSER in EIS if available (elastic#220096)

Author: viduni94

x-pack/platform/packages/shared/kbn-ai-assistant/src/knowledge_base/select_model_and_install_knowledge_base.tsx

@@ -49,22 +49,22 @@ export function SelectModelAndInstallKnowledgeBase({
 
   const { inferenceEndpoints, isLoading: isLoadingEndpoints, error } = useInferenceEndpoints();
 
+  const modelOptions: ModelOptionsData[] = getModelOptionsForInferenceEndpoints({
+    endpoints: inferenceEndpoints,
+  });
+
   useEffect(() => {
-    if (!selectedInferenceId && inferenceEndpoints.length) {
-      setSelectedInferenceId(inferenceEndpoints[0].inference_id);
+    if (!selectedInferenceId && modelOptions?.length) {
+      setSelectedInferenceId(modelOptions[0].key);
     }
-  }, [inferenceEndpoints, selectedInferenceId]);
+  }, [modelOptions, selectedInferenceId]);
 
   const handleInstall = () => {
     if (selectedInferenceId) {
       onInstall(selectedInferenceId);
     }
   };
 
-  const modelOptions: ModelOptionsData[] = getModelOptionsForInferenceEndpoints({
-    endpoints: inferenceEndpoints,
-  });
-
   const superSelectOptions = modelOptions.map((option: ModelOptionsData) => ({
     value: option.key,
     inputDisplay: option.label,

x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/get_model_options_for_inference_endpoints.test.ts

@@ -14,27 +14,100 @@ import {
   ModelOptionsData,
 } from './get_model_options_for_inference_endpoints';
 import type { InferenceAPIConfigResponse } from '@kbn/ml-trained-models-utils';
+import {
+  ELSER_ON_ML_NODE_INFERENCE_ID,
+  E5_LARGE_IN_EIS_INFERENCE_ID,
+  E5_SMALL_INFERENCE_ID,
+  ELSER_IN_EIS_INFERENCE_ID,
+} from '@kbn/observability-ai-assistant-plugin/public';
 
 describe('getModelOptionsForInferenceEndpoints', () => {
   it('maps known inference endpoints to user-friendly titles and descriptions', () => {
     const endpoints = [
-      { inference_id: '.elser-2-elasticsearch' },
-      { inference_id: '.multilingual-e5-small-elasticsearch' },
+      { inference_id: ELSER_ON_ML_NODE_INFERENCE_ID },
+      { inference_id: E5_SMALL_INFERENCE_ID },
     ] as InferenceAPIConfigResponse[];
 
-    const options: ModelOptionsData[] = getModelOptionsForInferenceEndpoints({ endpoints });
+    const options: ModelOptionsData[] = getModelOptionsForInferenceEndpoints({
+      endpoints,
+    });
 
     expect(options).toEqual([
       {
-        key: '.elser-2-elasticsearch',
+        key: ELSER_ON_ML_NODE_INFERENCE_ID,
         label: elserTitle,
         description: elserDescription,
       },
       {
-        key: '.multilingual-e5-small-elasticsearch',
+        key: E5_SMALL_INFERENCE_ID,
         label: e5SmallTitle,
         description: e5SmallDescription,
       },
     ]);
   });
+
+  it('shows only ELSER in EIS when both ELSER models are available', () => {
+    const endpoints = [
+      { inference_id: ELSER_ON_ML_NODE_INFERENCE_ID },
+      { inference_id: ELSER_IN_EIS_INFERENCE_ID },
+      { inference_id: E5_SMALL_INFERENCE_ID },
+    ] as InferenceAPIConfigResponse[];
+
+    const options = getModelOptionsForInferenceEndpoints({
+      endpoints,
+    });
+
+    expect(options.map((o) => o.key)).toEqual([ELSER_IN_EIS_INFERENCE_ID, E5_SMALL_INFERENCE_ID]);
+  });
+
+  it('shows only E5-large in EIS when both E5 (small and large) models are available', () => {
+    const endpoints = [
+      { inference_id: ELSER_ON_ML_NODE_INFERENCE_ID },
+      { inference_id: E5_SMALL_INFERENCE_ID },
+      { inference_id: E5_LARGE_IN_EIS_INFERENCE_ID },
+    ] as InferenceAPIConfigResponse[];
+
+    const options = getModelOptionsForInferenceEndpoints({
+      endpoints,
+    });
+
+    expect(options.map((o) => o.key)).toEqual([
+      ELSER_ON_ML_NODE_INFERENCE_ID,
+      E5_LARGE_IN_EIS_INFERENCE_ID,
+    ]);
+  });
+
+  it('shows only EIS models when both ELSER and E5-large pre-configured endpoints are available in EIS', () => {
+    const endpoints = [
+      { inference_id: ELSER_ON_ML_NODE_INFERENCE_ID },
+      { inference_id: ELSER_IN_EIS_INFERENCE_ID },
+      { inference_id: E5_SMALL_INFERENCE_ID },
+      { inference_id: E5_LARGE_IN_EIS_INFERENCE_ID },
+    ] as InferenceAPIConfigResponse[];
+
+    const options = getModelOptionsForInferenceEndpoints({
+      endpoints,
+    });
+
+    expect(options.map((o) => o.key)).toEqual([
+      ELSER_IN_EIS_INFERENCE_ID,
+      E5_LARGE_IN_EIS_INFERENCE_ID,
+    ]);
+  });
+
+  it('does not show any EIS models if EIS pre-configured endpoints are not available', () => {
+    const endpoints = [
+      { inference_id: ELSER_ON_ML_NODE_INFERENCE_ID },
+      { inference_id: E5_SMALL_INFERENCE_ID },
+    ] as InferenceAPIConfigResponse[];
+
+    const options = getModelOptionsForInferenceEndpoints({
+      endpoints,
+    });
+
+    expect(options.map((o) => o.key)).toEqual([
+      ELSER_ON_ML_NODE_INFERENCE_ID,
+      E5_SMALL_INFERENCE_ID,
+    ]);
+  });
 });

x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/get_model_options_for_inference_endpoints.ts

@@ -7,6 +7,12 @@
 
 import { i18n } from '@kbn/i18n';
 import { InferenceAPIConfigResponse } from '@kbn/ml-trained-models-utils';
+import {
+  ELSER_ON_ML_NODE_INFERENCE_ID,
+  E5_LARGE_IN_EIS_INFERENCE_ID,
+  E5_SMALL_INFERENCE_ID,
+  ELSER_IN_EIS_INFERENCE_ID,
+} from '@kbn/observability-ai-assistant-plugin/public';
 
 export interface ModelOptionsData {
   key: string;
@@ -63,19 +69,19 @@ const PRECONFIGURED_INFERENCE_ENDPOINT_METADATA: Record<
   string,
   { title: string; description: string }
 > = {
-  '.elser-2-elasticsearch': {
+  [ELSER_ON_ML_NODE_INFERENCE_ID]: {
     title: elserTitle,
     description: elserDescription,
   },
-  '.elser-v2-elastic': {
+  [ELSER_IN_EIS_INFERENCE_ID]: {
     title: elserTitle,
     description: elserDescription,
   },
-  '.multilingual-e5-small-elasticsearch': {
+  [E5_SMALL_INFERENCE_ID]: {
     title: e5SmallTitle,
     description: e5SmallDescription,
   },
-  '.multilingual-e5-large-elasticsearch': {
+  [E5_LARGE_IN_EIS_INFERENCE_ID]: {
     title: e5LargeTitle,
     description: e5LargeDescription,
   },
@@ -86,22 +92,30 @@ export const getModelOptionsForInferenceEndpoints = ({
 }: {
   endpoints: InferenceAPIConfigResponse[];
 }): ModelOptionsData[] => {
-  // TODO: add logic to show the EIS models if EIS is enabled, if not show the other models
-  const preConfiguredEndpoints = endpoints
-    .map((endpoint) => {
-      const meta = PRECONFIGURED_INFERENCE_ENDPOINT_METADATA[endpoint.inference_id];
+  const hasElserEIS = endpoints.some((ep) => ep.inference_id === ELSER_IN_EIS_INFERENCE_ID);
+  const hasE5EIS = endpoints.some((ep) => ep.inference_id === E5_LARGE_IN_EIS_INFERENCE_ID);
+
+  return endpoints
+    .filter((endpoint) => {
+      // if ELSER exists in EIS, skip the other ELSER model
+      if (endpoint.inference_id === ELSER_ON_ML_NODE_INFERENCE_ID && hasElserEIS) {
+        return false;
+      }
 
-      if (!meta) {
-        return undefined;
+      if (endpoint.inference_id === E5_SMALL_INFERENCE_ID && hasE5EIS) {
+        return false;
       }
 
+      // Only include preconfigured endpoints and skip custom endpoints
+      return Boolean(PRECONFIGURED_INFERENCE_ENDPOINT_METADATA[endpoint.inference_id]);
+    })
+    .map((endpoint) => {
+      const meta = PRECONFIGURED_INFERENCE_ENDPOINT_METADATA[endpoint.inference_id]!;
+
       return {
         key: endpoint.inference_id,
         label: meta.title,
         description: meta.description,
       };
-    })
-    .filter(Boolean) as ModelOptionsData[];
-
-  return preConfiguredEndpoints;
+    });
 };

x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts

@@ -52,3 +52,10 @@ export {
 export { concatenateChatCompletionChunks } from './utils/concatenate_chat_completion_chunks';
 
 export { ShortIdTable } from './utils/short_id_table';
+
+export {
+  ELSER_ON_ML_NODE_INFERENCE_ID,
+  ELSER_IN_EIS_INFERENCE_ID,
+  E5_SMALL_INFERENCE_ID,
+  E5_LARGE_IN_EIS_INFERENCE_ID,
+} from './preconfigured_inference_ids';

x-pack/platform/plugins/shared/observability_ai_assistant/common/preconfigured_inference_ids.ts

@@ -0,0 +1,12 @@
+/*
+ * 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.
+ */
+
+export const ELSER_ON_ML_NODE_INFERENCE_ID = '.elser-2-elasticsearch';
+export const ELSER_IN_EIS_INFERENCE_ID = '.elser-v2-elastic';
+
+export const E5_SMALL_INFERENCE_ID = '.multilingual-e5-small-elasticsearch';
+export const E5_LARGE_IN_EIS_INFERENCE_ID = '.multilingual-e5-large-elastic'; // TODO: verify the inference ID once it's created in EIS

x-pack/platform/plugins/shared/observability_ai_assistant/public/index.ts

@@ -69,6 +69,10 @@ export {
   ConversationAccess,
   KnowledgeBaseType,
   KnowledgeBaseState,
+  ELSER_ON_ML_NODE_INFERENCE_ID,
+  ELSER_IN_EIS_INFERENCE_ID,
+  E5_SMALL_INFERENCE_ID,
+  E5_LARGE_IN_EIS_INFERENCE_ID,
 } from '../common';
 
 export type {

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/index_assets/update_existing_index_assets.ts

@@ -6,15 +6,14 @@
  */
 
 import type { CoreSetup, Logger } from '@kbn/core/server';
+import { ELSER_ON_ML_NODE_INFERENCE_ID } from '../../../common';
 import type { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
 import { createOrUpdateConversationIndexAssets } from './create_or_update_conversation_index_assets';
 import { createOrUpdateKnowledgeBaseIndexAssets } from './create_or_update_knowledge_base_index_assets';
 import { hasKbWriteIndex } from '../knowledge_base_service/has_kb_index';
 import { getInferenceIdFromWriteIndex } from '../knowledge_base_service/get_inference_id_from_write_index';
 import { resourceNames } from '..';
 
-export const DEFAULT_INFERENCE_ENDPOINT = '.elser-2-elasticsearch';
-
 export async function updateExistingIndexAssets({
   logger,
   core,
@@ -48,7 +47,7 @@ export async function updateExistingIndexAssets({
       logger.debug(
         `Current KB write index does not have an inference_id. This is to be expected for indices created before 8.16`
       );
-      return DEFAULT_INFERENCE_ENDPOINT;
+      return ELSER_ON_ML_NODE_INFERENCE_ID;
     });
 
     await createOrUpdateKnowledgeBaseIndexAssets({ logger, core, inferenceId: currentInferenceId });