[Security Assistant] Fixes Knowledge Base setup when ML API's are una…

…vailable or return an error (elastic#189137) ## Summary This PR resolves two issues: 1. If ML API's are unavailable, we still show the 'Setup Knowledge Base' button. 2. If an error occurs during KB setup, we don't show an error toast. To test scenario `1.`, start Elasticsearch without ML, ala `yarn es snapshot -E xpack.enabled.ml=false`, and observe the following disabled 'setup' buttons with tooltip directing users to the docs: <p align="center"> <img width="200" src="https://github.com/user-attachments/assets/cd4575fe-2d74-4e2c-8c6a-d5e458a00f6c" /> <img width="200" src="https://github.com/user-attachments/assets/b79a31d2-5d8d-42ed-9270-f646daa1402c" /> <img width="200" src="https://github.com/user-attachments/assets/a043c3b8-987a-4d07-afb8-b5f1ce6d7d6c" /> </p> To test scenario `2.`, start Elasticsearch with insufficient ML memory, ala `yarn es snapshot -E xpack.ml.max_machine_memory_percent=5`, and observe the following error toasts when setting up the KB: <p align="center"> <img width="200" src="https://github.com/user-attachments/assets/6ef592ce-b4dc-4bfb-a8ec-8e16b7557557" /> <img width="200" src="https://github.com/user-attachments/assets/9e5165a0-66a9-432d-9608-85b0680b3249" /> <img width="200" src="https://github.com/user-attachments/assets/e85d4c7c-80ba-4ea3-be4a-1addd3d2520f" /> </p> ### Checklist Delete any items that are not applicable to this PR. - [X] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md) - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <[email protected]>
lcawl · Jul 24, 2024 · f5459ba · f5459ba
1 parent cbb91f1
commit f5459ba
Show file tree

Hide file tree

Showing 13 changed files with 121 additions and 61 deletions.
diff --git a/...ck/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/...ck/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts
@@ -69,6 +69,7 @@ export const ReadKnowledgeBaseResponse = z.object({
   elser_exists: z.boolean().optional(),
   esql_exists: z.boolean().optional(),
   index_exists: z.boolean().optional(),
+  is_setup_available: z.boolean().optional(),
   is_setup_in_progress: z.boolean().optional(),
   pipeline_exists: z.boolean().optional(),
 });
diff --git a/...ckages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml b/...ckages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml
@@ -66,6 +66,8 @@ paths:
                     type: boolean
                   index_exists:
                     type: boolean
+                  is_setup_available:
+                    type: boolean
                   is_setup_in_progress:
                     type: boolean
                   pipeline_exists:

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/api.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/api.test.tsx
@@ -18,6 +18,7 @@ const mockHttp = {
 describe('API tests', () => {
   beforeEach(() => {
     jest.clearAllMocks();
+    (mockHttp.fetch as jest.Mock).mockImplementation(() => Promise.resolve({}));
   });
 
   const knowledgeBaseArgs = {
@@ -68,7 +69,7 @@ describe('API tests', () => {
         throw new Error(error);
       });
 
-      await expect(postKnowledgeBase(knowledgeBaseArgs)).resolves.toThrowError('simulated error');
+      await expect(postKnowledgeBase(knowledgeBaseArgs)).rejects.toThrowError('simulated error');
     });
   });
 

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/api.tsx
@@ -57,7 +57,7 @@ export const getKnowledgeBaseStatus = async ({
  * @param {string} [options.resource] - Resource to be added to the KB, otherwise sets up the base KB
  * @param {AbortSignal} [options.signal] - AbortSignal
  *
- * @returns {Promise<CreateKnowledgeBaseResponse | IHttpFetchError>}
+ * @returns {Promise<CreateKnowledgeBaseResponse>}
  */
 export const postKnowledgeBase = async ({
   http,
@@ -66,19 +66,15 @@ export const postKnowledgeBase = async ({
 }: CreateKnowledgeBaseRequestParams & {
   http: HttpSetup;
   signal?: AbortSignal | undefined;
-}): Promise<CreateKnowledgeBaseResponse | IHttpFetchError> => {
-  try {
-    const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
-    const response = await http.fetch(path, {
-      method: 'POST',
-      signal,
-      version: API_VERSIONS.internal.v1,
-    });
+}): Promise<CreateKnowledgeBaseResponse> => {
+  const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
+  const response = await http.fetch(path, {
+    method: 'POST',
+    signal,
+    version: API_VERSIONS.internal.v1,
+  });
 
-    return response as CreateKnowledgeBaseResponse;
-  } catch (error) {
-    return error as IHttpFetchError;
-  }
+  return response as CreateKnowledgeBaseResponse;
 };
 
 /**

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx
@@ -101,6 +101,7 @@ describe('Knowledge base settings', () => {
           esql_exists: false,
           index_exists: false,
           pipeline_exists: false,
+          is_setup_available: true,
         },
         isLoading: false,
         isFetching: false,

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx
@@ -18,6 +18,7 @@ import {
   EuiHealth,
   EuiButtonEmpty,
   EuiLink,
+  EuiToolTip,
 } from '@elastic/eui';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { css } from '@emotion/react';
@@ -28,6 +29,7 @@ import type { KnowledgeBaseConfig } from '../assistant/types';
 import * as i18n from './translations';
 import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
 import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base';
+import { SETUP_KNOWLEDGE_BASE_BUTTON_TOOLTIP } from './translations';
 
 const ESQL_RESOURCE = 'esql';
 const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-knowledge-base-(SPACE)';
@@ -42,13 +44,13 @@ interface Props {
  */
 export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
   ({ knowledgeBase, setUpdatedKnowledgeBaseSettings }) => {
-    const { http } = useAssistantContext();
+    const { http, toasts } = useAssistantContext();
     const {
       data: kbStatus,
       isLoading,
       isFetching,
     } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
-    const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http });
+    const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
 
     // Resource enabled state
     const isElserEnabled = kbStatus?.elser_exists ?? false;
@@ -57,6 +59,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
       (isElserEnabled && isESQLEnabled && kbStatus?.index_exists && kbStatus?.pipeline_exists) ??
       false;
     const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false;
+    const isSetupAvailable = kbStatus?.is_setup_available ?? false;
 
     // Resource availability state
     const isLoadingKb = isLoading || isFetching || isSettingUpKB || isSetupInProgress;
@@ -72,21 +75,32 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
       setupKB(ESQL_RESOURCE);
     }, [setupKB]);
 
+    const toolTipContent = !isSetupAvailable ? SETUP_KNOWLEDGE_BASE_BUTTON_TOOLTIP : undefined;
+
     const setupKnowledgeBaseButton = useMemo(() => {
       return isKnowledgeBaseSetup ? (
         <></>
       ) : (
-        <EuiButtonEmpty
-          color={'primary'}
-          data-test-subj={'setupKnowledgeBaseButton'}
-          onClick={onSetupKnowledgeBaseButtonClick}
-          size="xs"
-          isLoading={isLoadingKb}
-        >
-          {i18n.SETUP_KNOWLEDGE_BASE_BUTTON}
-        </EuiButtonEmpty>
+        <EuiToolTip position={'bottom'} content={toolTipContent}>
+          <EuiButtonEmpty
+            color={'primary'}
+            data-test-subj={'setupKnowledgeBaseButton'}
+            disabled={!isSetupAvailable}
+            onClick={onSetupKnowledgeBaseButtonClick}
+            size="xs"
+            isLoading={isLoadingKb}
+          >
+            {i18n.SETUP_KNOWLEDGE_BASE_BUTTON}
+          </EuiButtonEmpty>
+        </EuiToolTip>
       );
-    }, [isKnowledgeBaseSetup, isLoadingKb, onSetupKnowledgeBaseButtonClick]);
+    }, [
+      isKnowledgeBaseSetup,
+      isLoadingKb,
+      isSetupAvailable,
+      onSetupKnowledgeBaseButtonClick,
+      toolTipContent,
+    ]);
 
     //////////////////////////////////////////////////////////////////////////////////////////
     // Knowledge Base Resource

diff --git a/...packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx b/...packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx
@@ -17,6 +17,7 @@ import {
   EuiHealth,
   EuiButtonEmpty,
   EuiPanel,
+  EuiToolTip,
 } from '@elastic/eui';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { css } from '@emotion/react';
@@ -33,6 +34,7 @@ import {
 } from '../assistant/settings/use_settings_updater/use_settings_updater';
 import { AssistantSettingsBottomBar } from '../assistant/settings/assistant_settings_bottom_bar';
 import { SETTINGS_UPDATED_TOAST_TITLE } from '../assistant/settings/translations';
+import { SETUP_KNOWLEDGE_BASE_BUTTON_TOOLTIP } from './translations';
 
 const ESQL_RESOURCE = 'esql';
 const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-knowledge-base-(SPACE)';
@@ -87,7 +89,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
     isLoading,
     isFetching,
   } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
-  const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http });
+  const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
 
   // Resource enabled state
   const isElserEnabled = kbStatus?.elser_exists ?? false;
@@ -96,6 +98,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
     (isElserEnabled && isESQLEnabled && kbStatus?.index_exists && kbStatus?.pipeline_exists) ??
     false;
   const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false;
+  const isSetupAvailable = kbStatus?.is_setup_available ?? false;
 
   // Resource availability state
   const isLoadingKb = isLoading || isFetching || isSettingUpKB || isSetupInProgress;
@@ -111,21 +114,32 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
     setupKB(ESQL_RESOURCE);
   }, [setupKB]);
 
+  const toolTipContent = !isSetupAvailable ? SETUP_KNOWLEDGE_BASE_BUTTON_TOOLTIP : undefined;
+
   const setupKnowledgeBaseButton = useMemo(() => {
     return isKnowledgeBaseSetup ? (
       <></>
     ) : (
-      <EuiButtonEmpty
-        color={'primary'}
-        data-test-subj={'setupKnowledgeBaseButton'}
-        onClick={onSetupKnowledgeBaseButtonClick}
-        size="xs"
-        isLoading={isLoadingKb}
-      >
-        {i18n.SETUP_KNOWLEDGE_BASE_BUTTON}
-      </EuiButtonEmpty>
+      <EuiToolTip position={'bottom'} content={toolTipContent}>
+        <EuiButtonEmpty
+          color={'primary'}
+          data-test-subj={'setupKnowledgeBaseButton'}
+          disabled={!isSetupAvailable}
+          onClick={onSetupKnowledgeBaseButtonClick}
+          size="xs"
+          isLoading={isLoadingKb}
+        >
+          {i18n.SETUP_KNOWLEDGE_BASE_BUTTON}
+        </EuiButtonEmpty>
+      </EuiToolTip>
     );
-  }, [isKnowledgeBaseSetup, isLoadingKb, onSetupKnowledgeBaseButtonClick]);
+  }, [
+    isKnowledgeBaseSetup,
+    isLoadingKb,
+    isSetupAvailable,
+    onSetupKnowledgeBaseButtonClick,
+    toolTipContent,
+  ]);
 
   //////////////////////////////////////////////////////////////////////////////////////////
   // Knowledge Base Resource

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx
@@ -6,7 +6,7 @@
  */
 
 import React, { useCallback } from 'react';
-import { EuiButton } from '@elastic/eui';
+import { EuiButton, EuiToolTip } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 
 import { useAssistantContext } from '../..';
@@ -40,19 +40,28 @@ export const SetupKnowledgeBaseButton: React.FC = React.memo(() => {
     return null;
   }
 
+  const toolTipContent = !kbStatus?.is_setup_available
+    ? i18n.translate('xpack.elasticAssistant.knowledgeBase.installKnowledgeBaseButtonToolTip', {
+        defaultMessage: 'Knowledge Base unavailable, please see documentation for more details.',
+      })
+    : undefined;
+
   return (
-    <EuiButton
-      color="primary"
-      data-test-subj="setup-knowledge-base-button"
-      fill
-      isLoading={isSetupInProgress}
-      iconType="importAction"
-      onClick={onInstallKnowledgeBase}
-    >
-      {i18n.translate('xpack.elasticAssistant.knowledgeBase.installKnowledgeBaseButton', {
-        defaultMessage: 'Setup Knowledge Base',
-      })}
-    </EuiButton>
+    <EuiToolTip position={'bottom'} content={toolTipContent}>
+      <EuiButton
+        color="primary"
+        data-test-subj="setup-knowledge-base-button"
+        fill
+        disabled={!kbStatus?.is_setup_available}
+        isLoading={isSetupInProgress}
+        iconType="importAction"
+        onClick={onInstallKnowledgeBase}
+      >
+        {i18n.translate('xpack.elasticAssistant.knowledgeBase.installKnowledgeBaseButton', {
+          defaultMessage: 'Setup Knowledge Base',
+        })}
+      </EuiButton>
+    </EuiToolTip>
   );
 });
 

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/translations.ts b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/translations.ts
@@ -87,6 +87,13 @@ export const SETUP_KNOWLEDGE_BASE_BUTTON = i18n.translate(
   }
 );
 
+export const SETUP_KNOWLEDGE_BASE_BUTTON_TOOLTIP = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.setupKnowledgeBaseButtonToolTip',
+  {
+    defaultMessage: 'Knowledge Base unavailable, please see documentation for more details.',
+  }
+);
+
 export const KNOWLEDGE_BASE_TOOLTIP = i18n.translate(
   'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.knowledgeBaseTooltip',
   {

diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts
@@ -46,6 +46,23 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
     return this.options.getIsKBSetupInProgress();
   }
 
+  /**
+   * Returns whether setup of the Knowledge Base can be performed (essentially an ML features check)
+   *
+   */
+  public isSetupAvailable = async () => {
+    // ML plugin requires request to retrieve capabilities, which are in turn scoped to the user from the request,
+    // so we just test the API for a 404 instead to determine if ML is 'available'
+    // TODO: expand to include memory check, see https://github.com/elastic/ml-team/issues/1208#issuecomment-2115770318
+    try {
+      const esClient = await this.options.elasticsearchClientPromise;
+      await esClient.ml.getMemoryStats({ human: true });
+    } catch (error) {
+      return false;
+    }
+    return true;
+  };
+
   /**
    * Downloads and installs ELSER model if not already installed
    *
@@ -104,10 +121,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
         wait_for: 'fully_allocated',
       });
     } catch (error) {
-      if (!isModelAlreadyExistsError(error)) {
-        this.options.logger.error(`Error deploying ELSER model '${elserId}':\n${error}`);
-      }
-      this.options.logger.debug(`Error deploying ELSER model '${elserId}', model already deployed`);
+      this.options.logger.error(`Error deploying ELSER model '${elserId}':\n${error}`);
+      throw new Error(`Error deploying ELSER model '${elserId}':\n${error}`);
     }
   };
 
@@ -214,7 +229,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
         this.options.logger.debug(`Knowledge Base docs already loaded!`);
       }
     } catch (e) {
+      this.options.setIsKBSetupInProgress(false);
       this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`);
+      throw new Error(`Error setting up Knowledge Base: ${e.message}`);
     }
     this.options.setIsKBSetupInProgress(false);
   };

diff --git a/.../plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts b/.../plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts
@@ -38,6 +38,7 @@ describe('Get Knowledge Base Status Route', () => {
         alias: 'knowledge-base-alias',
       },
       isModelInstalled: jest.fn().mockResolvedValue(true),
+      isSetupAvailable: jest.fn().mockResolvedValue(true),
     });
 
     getKnowledgeBaseStatusRoute(server.router, mockGetElser);

diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts
@@ -79,11 +79,13 @@ export const getKnowledgeBaseStatusRoute = (
           const indexExists = await esStore.indexExists();
           const pipelineExists = await esStore.pipelineExists();
           const modelExists = await esStore.isModelInstalled(elserId);
+          const setupAvailable = await kbDataClient.isSetupAvailable();
 
           const body: ReadKnowledgeBaseResponse = {
             elser_exists: modelExists,
             index_exists: indexExists,
             is_setup_in_progress: kbDataClient.isSetupInProgress,
+            is_setup_available: setupAvailable,
             pipeline_exists: pipelineExists,
           };
 

diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts
@@ -5,8 +5,6 @@
  * 2.0.
  */
 
-import { transformError } from '@kbn/securitysolution-es-utils';
-
 import {
   ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION,
   CreateKnowledgeBaseRequestParams,
@@ -88,13 +86,10 @@ export const postKnowledgeBaseRoute = (
           await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, soClient });
 
           return response.ok({ body: { success: true } });
-        } catch (err) {
-          logger.log(err);
-          const error = transformError(err);
-
+        } catch (error) {
           return resp.error({
             body: error.message,
-            statusCode: error.statusCode,
+            statusCode: 500,
           });
         }
       }