From 60bb1f8da1f50e0cca42d4c73fc9c730e5c76a0a Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 19 Oct 2023 14:34:39 -0600 Subject: [PATCH] [Security Solution] [Elastic AI Assistant] Throw error if Knowledge Base is enabled but ELSER is unavailable (#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.

> [!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: https://github.com/elastic/kibana/issues/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 <42973632+kibanamachine@users.noreply.github.com> --- .../impl/assistant/api.tsx | 2 +- .../impl/assistant/context_pills/index.tsx | 39 +++++++---- .../knowledge_base_settings.tsx | 37 ++++++---- .../impl/knowledge_base/translations.ts | 7 ++ .../elasticsearch_store.test.ts | 68 ++++++++++++++++--- .../elasticsearch_store.ts | 15 ++-- .../execute_custom_llm_chain/index.test.ts | 28 ++++++++ .../execute_custom_llm_chain/index.ts | 8 +++ .../routes/post_actions_connector_execute.ts | 4 ++ 9 files changed, 165 insertions(+), 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 index 8ccb2e72cfee9..d6b61ee70b1ae 100644 --- 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 index be5374e37bd67..89aef125b5d0e 100644 --- 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 = ({ return ( - {sortedPromptContexts.map(({ description, id, getPromptContext, tooltip }) => ( - - - selectPromptContext(id)} - > - {description} - - - - ))} + {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 = ( + selectPromptContext(id)} + > + {description} + + ); + return ( + + {selectedPromptContexts[id] != null ? ( + button + ) : ( + {button} + )} + + ); + })} ); }; 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 index bed527b6434e3..e49602a52411a 100644 --- 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,8 +57,8 @@ export const KnowledgeBaseSettings: React.FC = 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 @@ -65,9 +66,11 @@ export const KnowledgeBaseSettings: React.FC = React.memo( 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 = React.memo( return isLoadingKb ? ( ) : ( - + + + ); - }, [isLoadingKb, knowledgeBase.assistantLangChain, onEnableAssistantLangChainChange]); + }, [ + isLoadingKb, + isSwitchDisabled, + knowledgeBase.assistantLangChain, + onEnableAssistantLangChainChange, + ]); ////////////////////////////////////////////////////////////////////////////////////////// // Knowledge Base Resource @@ -205,7 +216,7 @@ export const KnowledgeBaseSettings: React.FC = 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 index 95417ddf6a889..1aa295e311e68 100644 --- 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/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.test.ts index 1de907c3ddc9c..0dc1cd10499cc 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.test.ts +++ b/x-pack/plugins/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/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts index 52f3fe87275db..b6c8f37384862 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts +++ b/x-pack/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 indicating whether the model is installed */ - async isModelInstalled(modelId: string): Promise { + async isModelInstalled(modelId?: string): Promise { 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 index e63da9257aa36..7f06b4b456194 100644 --- 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 index 694bd44bfd471..713a330b46c48 100644 --- 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 index 8da820288ae1b..299d8ade24a3f 100644 --- 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;