Skip to content

Commit

Permalink
[Security Assistant] Fixes Knowledge Base setup when ML API's are una…
Browse files Browse the repository at this point in the history
…vailable or return an error (elastic#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 <[email protected]>
  • Loading branch information
spong and kibanamachine authored Jul 24, 2024
1 parent cbb91f1 commit f5459ba
Show file tree
Hide file tree
Showing 13 changed files with 121 additions and 61 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
});
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ paths:
type: boolean
index_exists:
type: boolean
is_setup_available:
type: boolean
is_setup_in_progress:
type: boolean
pipeline_exists:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ const mockHttp = {
describe('API tests', () => {
beforeEach(() => {
jest.clearAllMocks();
(mockHttp.fetch as jest.Mock).mockImplementation(() => Promise.resolve({}));
});

const knowledgeBaseArgs = {
Expand Down Expand Up @@ -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');
});
});

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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;
};

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import {
EuiHealth,
EuiButtonEmpty,
EuiLink,
EuiToolTip,
} from '@elastic/eui';
import { FormattedMessage } from '@kbn/i18n-react';
import { css } from '@emotion/react';
Expand All @@ -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)';
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import {
EuiHealth,
EuiButtonEmpty,
EuiPanel,
EuiToolTip,
} from '@elastic/eui';
import { FormattedMessage } from '@kbn/i18n-react';
import { css } from '@emotion/react';
Expand All @@ -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)';
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 '../..';
Expand Down Expand Up @@ -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>
);
});

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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',
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
*
Expand Down Expand Up @@ -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}`);
}
};

Expand Down Expand Up @@ -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);
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
* 2.0.
*/

import { transformError } from '@kbn/securitysolution-es-utils';

import {
ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION,
CreateKnowledgeBaseRequestParams,
Expand Down Expand Up @@ -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,
});
}
}
Expand Down

0 comments on commit f5459ba

Please sign in to comment.