[Security Assistant] Fixes Knowledge Base setup when ML API's are unavailable or return an error (#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 <42973632+kibanamachine@users.noreply.github.com>
This commit is contained in:
Garrett Spong 2024-07-24 17:47:51 -06:00 committed by GitHub
parent cbb91f1f6b
commit f5459ba835
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 121 additions and 61 deletions

View file

@ -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(),
});

View file

@ -66,6 +66,8 @@ paths:
type: boolean
index_exists:
type: boolean
is_setup_available:
type: boolean
is_setup_in_progress:
type: boolean
pipeline_exists:

View file

@ -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');
});
});

View file

@ -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;
};
/**

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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>
);
});

View file

@ -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',
{

View file

@ -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);
};

View file

@ -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);

View file

@ -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,
};

View file

@ -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,
});
}
}