From 97c65a5cbd6c498043f7a34e0478f5483b040f62 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 19 Oct 2023 16:53:34 -0600 Subject: [PATCH] [8.11] [Security Solution] [Elastic AI Assistant] Throw error if Knowledge 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)](https://github.com/elastic/kibana/pull/169330) ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport) --- .../impl/assistant/api.tsx | 2 +- .../impl/assistant/context_pills/index.tsx | 39 +++++++---- .../knowledge_base_settings.tsx | 54 +++++++++------ .../knowledge_base_settings/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, 175 insertions(+), 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 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/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings/knowledge_base_settings.tsx index 2c667a5739ef6..e8b454b068ca7 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings/knowledge_base_settings.tsx +++ b/x-pack/packages/kbn-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,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,15 +96,24 @@ export const KnowledgeBaseSettings: React.FC = React.memo( return isLoadingKb ? ( ) : ( - + + + ); - }, [isLoadingKb, knowledgeBase.assistantLangChain, onEnableAssistantLangChainChange]); + }, [ + isLoadingKb, + isSwitchDisabled, + knowledgeBase.assistantLangChain, + onEnableAssistantLangChainChange, + ]); ////////////////////////////////////////////////////////////////////////////////////////// // Knowledge Base Resource @@ -123,6 +135,7 @@ export const KnowledgeBaseSettings: React.FC = React.memo( onEnableKB(!isKnowledgeBaseEnabled)} size="xs" > @@ -135,14 +148,14 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const knowledgeBaseDescription = useMemo(() => { return isKnowledgeBaseEnabled ? ( - <> + {i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(KNOWLEDGE_BASE_INDEX_PATTERN)}{' '} {knowledgeBaseActionButton} - + ) : ( - <> + {i18n.KNOWLEDGE_BASE_DESCRIPTION} {knowledgeBaseActionButton} - + ); }, [isKnowledgeBaseEnabled, knowledgeBaseActionButton]); @@ -166,6 +179,7 @@ export const KnowledgeBaseSettings: React.FC = React.memo( onEnableESQL(!isESQLEnabled)} size="xs" > @@ -176,13 +190,13 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const esqlDescription = useMemo(() => { return isESQLEnabled ? ( - <> + {i18n.ESQL_DESCRIPTION_INSTALLED} {esqlActionButton} - + ) : ( - <> + {i18n.ESQL_DESCRIPTION} {esqlActionButton} - + ); }, [esqlActionButton, isESQLEnabled]); @@ -202,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/knowledge_base_settings/translations.ts b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings/translations.ts index 95417ddf6a889..1aa295e311e68 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings/translations.ts +++ b/x-pack/packages/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/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;