[Security Solution] [Elastic AI Assistant] Throw error if Knowledge B…

…ase is enabled but ELSER is unavailable (elastic#169330) ## Summary This fixes the Knowledge Base UX a bit by throwing an error if somehow ELSER has been disabled in the background, and instructs the user on how to resolve or to disable the Knowledge Base to continue. Additionally, if ELSER is not available, we prevent the enabling of the Knowledge Base as to not provide a degraded experience when ELSER and the ES|QL documentation is not available. <p align="center"> <img width="500" src="https://github.com/elastic/kibana/assets/2946766/e4d326fa-c996-43ad-9d1c-d76f7d16f916" /> </p> > [!NOTE] > `isModelInstalled` logic has been updated to not just check the model `definition_status`, but to actually ensure that it's deployed by checking to see that it is `started` and `fully_allocated`. This better guards ELSER availability as the previous check would return true if the model was just downloaded and not actually deployed. Also resolves: elastic#169403 ## Test Instructions After enabling the KB, disable the ELSER deployment in the `Trained Models` ML UI and then try using the assistant. --------- Co-authored-by: kibanamachine <[email protected]>
davismcphee · Oct 19, 2023 · 60bb1f8 · 60bb1f8
1 parent 44f3910
commit 60bb1f8
Show file tree

Hide file tree

Showing 9 changed files with 165 additions and 43 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/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
@@ -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,16 +96,24 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
       return isLoadingKb ? (
         <EuiLoadingSpinner size="s" />
       ) : (
-        <EuiSwitch
-          showLabel={false}
-          data-test-subj="assistantLangChainSwitch"
-          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
@@ -205,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/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/translations.ts b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/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;