[8.11] [Security Solution] [Elastic AI Assistant] Throw error if Know…

…ledge Base is enabled but ELSER is unavailable (#169330) (#169455) # Backport This will backport the following commits from `main` to `8.11`: - [[Security Solution] [Elastic AI Assistant] Throw error if Knowledge Base is enabled but ELSER is unavailable (#169330)](#169330)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)
elastic · Oct 19, 2023 · 97c65a5 · 97c65a5
1 parent 107eab4
commit 97c65a5
Show file tree

Hide file tree

Showing 9 changed files with 175 additions and 50 deletions.
diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api.tsx
@@ -95,7 +95,7 @@ export const fetchConnectorExecuteAction = async ({
     };
   } catch (error) {
     return {
-      response: API_ERROR,
+      response: `${API_ERROR}\n\n${error?.body?.message ?? error?.message}`,
       isError: true,
     };
   }

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/context_pills/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/context_pills/index.tsx
@@ -66,21 +66,30 @@ const ContextPillsComponent: React.FC<Props> = ({
 
   return (
     <EuiFlexGroup gutterSize="none" wrap>
-      {sortedPromptContexts.map(({ description, id, getPromptContext, tooltip }) => (
-        <EuiFlexItem grow={false} key={id}>
-          <EuiToolTip content={tooltip}>
-            <PillButton
-              data-test-subj={`pillButton-${id}`}
-              disabled={selectedPromptContexts[id] != null}
-              iconSide="left"
-              iconType="plus"
-              onClick={() => selectPromptContext(id)}
-            >
-              {description}
-            </PillButton>
-          </EuiToolTip>
-        </EuiFlexItem>
-      ))}
+      {sortedPromptContexts.map(({ description, id, getPromptContext, tooltip }) => {
+        // Workaround for known issue where tooltip won't dismiss after button state is changed once clicked
+        // See: https://github.com/elastic/eui/issues/6488#issuecomment-1379656704
+        const button = (
+          <PillButton
+            data-test-subj={`pillButton-${id}`}
+            disabled={selectedPromptContexts[id] != null}
+            iconSide="left"
+            iconType="plus"
+            onClick={() => selectPromptContext(id)}
+          >
+            {description}
+          </PillButton>
+        );
+        return (
+          <EuiFlexItem grow={false} key={id}>
+            {selectedPromptContexts[id] != null ? (
+              button
+            ) : (
+              <EuiToolTip content={tooltip}>{button}</EuiToolTip>
+            )}
+          </EuiFlexItem>
+        );
+      })}
     </EuiFlexGroup>
   );
 };

diff --git a/...elastic-assistant/impl/knowledge_base/knowledge_base_settings/knowledge_base_settings.tsx b/...elastic-assistant/impl/knowledge_base/knowledge_base_settings/knowledge_base_settings.tsx
@@ -20,6 +20,7 @@ import {
   EuiFlexItem,
   EuiHealth,
   EuiButtonEmpty,
+  EuiToolTip,
   EuiSwitch,
 } from '@elastic/eui';
 
@@ -56,18 +57,20 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
     const { mutate: deleteKB, isLoading: isDeletingUpKB } = useDeleteKnowledgeBase({ http });
 
     // Resource enabled state
-    const isKnowledgeBaseEnabled =
-      (kbStatus?.index_exists && kbStatus?.pipeline_exists && kbStatus?.elser_exists) ?? false;
+    const isElserEnabled = kbStatus?.elser_exists ?? false;
+    const isKnowledgeBaseEnabled = (kbStatus?.index_exists && kbStatus?.pipeline_exists) ?? false;
     const isESQLEnabled = kbStatus?.esql_exists ?? false;
 
     // Resource availability state
     const isLoadingKb = isLoading || isFetching || isSettingUpKB || isDeletingUpKB;
     const isKnowledgeBaseAvailable = knowledgeBase.assistantLangChain && kbStatus?.elser_exists;
     const isESQLAvailable =
       knowledgeBase.assistantLangChain && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled;
+    // Prevent enabling if elser doesn't exist, but always allow to disable
+    const isSwitchDisabled = !kbStatus?.elser_exists && !knowledgeBase.assistantLangChain;
 
     // Calculated health state for EuiHealth component
-    const elserHealth = kbStatus?.elser_exists ? 'success' : 'subdued';
+    const elserHealth = isElserEnabled ? 'success' : 'subdued';
     const knowledgeBaseHealth = isKnowledgeBaseEnabled ? 'success' : 'subdued';
     const esqlHealth = isESQLEnabled ? 'success' : 'subdued';
 
@@ -93,15 +96,24 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
       return isLoadingKb ? (
         <EuiLoadingSpinner size="s" />
       ) : (
-        <EuiSwitch
-          showLabel={false}
-          checked={knowledgeBase.assistantLangChain}
-          onChange={onEnableAssistantLangChainChange}
-          label={i18n.KNOWLEDGE_BASE_LABEL}
-          compressed
-        />
+        <EuiToolTip content={isSwitchDisabled && i18n.KNOWLEDGE_BASE_TOOLTIP} position={'right'}>
+          <EuiSwitch
+            showLabel={false}
+            data-test-subj="assistantLangChainSwitch"
+            disabled={isSwitchDisabled}
+            checked={knowledgeBase.assistantLangChain}
+            onChange={onEnableAssistantLangChainChange}
+            label={i18n.KNOWLEDGE_BASE_LABEL}
+            compressed
+          />
+        </EuiToolTip>
       );
-    }, [isLoadingKb, knowledgeBase.assistantLangChain, onEnableAssistantLangChainChange]);
+    }, [
+      isLoadingKb,
+      isSwitchDisabled,
+      knowledgeBase.assistantLangChain,
+      onEnableAssistantLangChainChange,
+    ]);
 
     //////////////////////////////////////////////////////////////////////////////////////////
     // Knowledge Base Resource
@@ -123,6 +135,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
         <EuiButtonEmpty
           color={isKnowledgeBaseEnabled ? 'danger' : 'primary'}
           flush="left"
+          data-test-subj={'knowledgeBaseActionButton'}
           onClick={() => onEnableKB(!isKnowledgeBaseEnabled)}
           size="xs"
         >
@@ -135,14 +148,14 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
 
     const knowledgeBaseDescription = useMemo(() => {
       return isKnowledgeBaseEnabled ? (
-        <>
+        <span data-test-subj="kb-installed">
           {i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(KNOWLEDGE_BASE_INDEX_PATTERN)}{' '}
           {knowledgeBaseActionButton}
-        </>
+        </span>
       ) : (
-        <>
+        <span data-test-subj="install-kb">
           {i18n.KNOWLEDGE_BASE_DESCRIPTION} {knowledgeBaseActionButton}
-        </>
+        </span>
       );
     }, [isKnowledgeBaseEnabled, knowledgeBaseActionButton]);
 
@@ -166,6 +179,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
         <EuiButtonEmpty
           color={isESQLEnabled ? 'danger' : 'primary'}
           flush="left"
+          data-test-subj="esqlEnableButton"
           onClick={() => onEnableESQL(!isESQLEnabled)}
           size="xs"
         >
@@ -176,13 +190,13 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
 
     const esqlDescription = useMemo(() => {
       return isESQLEnabled ? (
-        <>
+        <span data-test-subj="esql-installed">
           {i18n.ESQL_DESCRIPTION_INSTALLED} {esqlActionButton}
-        </>
+        </span>
       ) : (
-        <>
+        <span data-test-subj="install-esql">
           {i18n.ESQL_DESCRIPTION} {esqlActionButton}
-        </>
+        </span>
       );
     }, [esqlActionButton, isESQLEnabled]);
 
@@ -202,7 +216,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
           display="columnCompressedSwitch"
           label={i18n.KNOWLEDGE_BASE_LABEL}
           css={css`
-            div {
+            .euiFormRow__labelWrapper {
               min-width: 95px !important;
             }
           `}

diff --git a/...ackages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings/translations.ts b/...ackages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings/translations.ts
@@ -36,6 +36,13 @@ export const KNOWLEDGE_BASE_LABEL = i18n.translate(
   }
 );
 
+export const KNOWLEDGE_BASE_TOOLTIP = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.knowledgeBaseTooltip',
+  {
+    defaultMessage: 'ELSER must be configured to enable the Knowledge Base',
+  }
+);
+
 export const KNOWLEDGE_BASE_DESCRIPTION = i18n.translate(
   'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.knowledgeBaseDescription',
   {

diff --git a/...ns/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.test.ts b/...ns/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.test.ts
@@ -9,7 +9,7 @@ import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
 import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
 import {
   IndicesCreateResponse,
-  MlGetTrainedModelsResponse,
+  MlGetTrainedModelsStatsResponse,
 } from '@elastic/elasticsearch/lib/api/types';
 import { Document } from 'langchain/document';
 
@@ -142,17 +142,69 @@ describe('ElasticsearchStore', () => {
     });
   });
 
-  describe('Model Management', () => {
-    it('Checks if a model is installed', async () => {
-      mockEsClient.ml.getTrainedModels.mockResolvedValue({
-        trained_model_configs: [{ fully_defined: true }],
-      } as MlGetTrainedModelsResponse);
+  describe('isModelInstalled', () => {
+    it('returns true if model is started and fully allocated', async () => {
+      mockEsClient.ml.getTrainedModelsStats.mockResolvedValue({
+        trained_model_stats: [
+          {
+            deployment_stats: {
+              state: 'started',
+              allocation_status: {
+                state: 'fully_allocated',
+              },
+            },
+          },
+        ],
+      } as MlGetTrainedModelsStatsResponse);
 
       const isInstalled = await esStore.isModelInstalled('.elser_model_2');
 
       expect(isInstalled).toBe(true);
-      expect(mockEsClient.ml.getTrainedModels).toHaveBeenCalledWith({
-        include: 'definition_status',
+      expect(mockEsClient.ml.getTrainedModelsStats).toHaveBeenCalledWith({
+        model_id: '.elser_model_2',
+      });
+    });
+
+    it('returns false if model is not started', async () => {
+      mockEsClient.ml.getTrainedModelsStats.mockResolvedValue({
+        trained_model_stats: [
+          {
+            deployment_stats: {
+              state: 'starting',
+              allocation_status: {
+                state: 'fully_allocated',
+              },
+            },
+          },
+        ],
+      } as MlGetTrainedModelsStatsResponse);
+
+      const isInstalled = await esStore.isModelInstalled('.elser_model_2');
+
+      expect(isInstalled).toBe(false);
+      expect(mockEsClient.ml.getTrainedModelsStats).toHaveBeenCalledWith({
+        model_id: '.elser_model_2',
+      });
+    });
+
+    it('returns false if model is not fully allocated', async () => {
+      mockEsClient.ml.getTrainedModelsStats.mockResolvedValue({
+        trained_model_stats: [
+          {
+            deployment_stats: {
+              state: 'started',
+              allocation_status: {
+                state: 'starting',
+              },
+            },
+          },
+        ],
+      } as MlGetTrainedModelsStatsResponse);
+
+      const isInstalled = await esStore.isModelInstalled('.elser_model_2');
+
+      expect(isInstalled).toBe(false);
+      expect(mockEsClient.ml.getTrainedModelsStats).toHaveBeenCalledWith({
         model_id: '.elser_model_2',
       });
     });

diff --git a/...plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/...plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts
@@ -37,7 +37,7 @@ interface CreateIndexParams {
 }
 
 /**
- * A fallback for the the query `size` that determines how many documents to
+ * A fallback for the query `size` that determines how many documents to
  * return from Elasticsearch when performing a similarity search.
  *
  * The size is typically determined by the implementation of LangChain's
@@ -360,14 +360,17 @@ export class ElasticsearchStore extends VectorStore {
    * @param modelId ID of the model to check
    * @returns Promise<boolean> indicating whether the model is installed
    */
-  async isModelInstalled(modelId: string): Promise<boolean> {
+  async isModelInstalled(modelId?: string): Promise<boolean> {
     try {
-      const getResponse = await this.esClient.ml.getTrainedModels({
-        model_id: modelId,
-        include: 'definition_status',
+      const getResponse = await this.esClient.ml.getTrainedModelsStats({
+        model_id: modelId ?? this.model,
       });
 
-      return Boolean(getResponse.trained_model_configs[0]?.fully_defined);
+      return getResponse.trained_model_stats.some(
+        (stats) =>
+          stats.deployment_stats?.state === 'started' &&
+          stats.deployment_stats?.allocation_status.state === 'fully_allocated'
+      );
     } catch (e) {
       // Returns 404 if it doesn't exist
       return false;

diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts
@@ -16,6 +16,7 @@ import { langChainMessages } from '../../../__mocks__/lang_chain_messages';
 import { ESQL_RESOURCE } from '../../../routes/knowledge_base/constants';
 import { ResponseBody } from '../types';
 import { callAgentExecutor } from '.';
+import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store';
 
 jest.mock('../llm/actions_client_llm');
 
@@ -36,6 +37,13 @@ jest.mock('langchain/agents', () => ({
   })),
 }));
 
+jest.mock('../elasticsearch_store/elasticsearch_store', () => ({
+  ElasticsearchStore: jest.fn().mockImplementation(() => ({
+    asRetriever: jest.fn(),
+    isModelInstalled: jest.fn().mockResolvedValue(true),
+  })),
+}));
+
 const mockConnectorId = 'mock-connector-id';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -129,4 +137,24 @@ describe('callAgentExecutor', () => {
       status: 'ok',
     });
   });
+
+  it('throws an error if ELSER model is not installed', async () => {
+    (ElasticsearchStore as unknown as jest.Mock).mockImplementationOnce(() => ({
+      isModelInstalled: jest.fn().mockResolvedValue(false),
+    }));
+
+    await expect(
+      callAgentExecutor({
+        actions: mockActions,
+        connectorId: mockConnectorId,
+        esClient: esClientMock,
+        langChainMessages,
+        logger: mockLogger,
+        request: mockRequest,
+        kbResource: ESQL_RESOURCE,
+      })
+    ).rejects.toThrow(
+      'Please ensure ELSER is configured to use the Knowledge Base, otherwise disable the Knowledge Base in Advanced Settings to continue.'
+    );
+  });
 });
diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts
@@ -47,6 +47,14 @@ export const callAgentExecutor = async ({
     elserId,
     kbResource
   );
+
+  const modelExists = await esStore.isModelInstalled();
+  if (!modelExists) {
+    throw new Error(
+      'Please ensure ELSER is configured to use the Knowledge Base, otherwise disable the Knowledge Base in Advanced Settings to continue.'
+    );
+  }
+
   const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever());
 
   const tools: Tool[] = [

diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts
@@ -44,12 +44,16 @@ export const postActionsConnectorExecuteRoute = (
 
         // if not langchain, call execute action directly and return the response:
         if (!request.body.assistantLangChain) {
+          logger.debug('Executing via actions framework directly, assistantLangChain: false');
           const result = await executeAction({ actions, request, connectorId });
           return response.ok({
             body: result,
           });
         }
 
+        // TODO: Add `traceId` to actions request when calling via langchain
+        logger.debug('Executing via langchain, assistantLangChain: true');
+
         // get a scoped esClient for assistant memory
         const esClient = (await context.core).elasticsearch.client.asCurrentUser;