github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-24 01:38:56 -04:00
Code Issues Projects Releases Packages Wiki Activity

[Security Assistant] Enables automatic setup of Knowledge Base and LangGraph code paths for 8.15 (#188168)

Browse source 
## Summary

This PR enables the automatic setup of the Knowledge Base and LangGraph
code paths for the `8.15` release. These features were behind the
`assistantKnowledgeBaseByDefault` feature flag, which will remain as a
gate for upcoming Knowledge Base features that were not ready for this
release.

As part of these changes, we now only support the new LangGraph code
path, and so were able to clean up the non-kb and non-RAGonAlerts code
paths. All paths within the `post_actions_executor` route funnel to the
LangGraph implementation.

> [!NOTE]
> We were planning to do the switch to the new
[`chat/completions`](https://github.com/elastic/kibana/pull/184485/files)
public API, however this would've required additional refactoring since
the API's slightly differ. We will make this change and delete the
`post_actions_executor` route for the next release.






### 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)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- Working with docs team to ensure updates before merging, cc
@benironside
- [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: Steph Milovic <stephanie.milovic@elastic.co>
This commit is contained in:
Garrett Spong 2024-07-17 15:44:24 -06:00 • committed by GitHub
parent 1a1c7d6101
commit 661c25133d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
55 changed files with 351 additions and 1968 deletions
Download patch file Download diff file Expand all files Collapse all files

2
.buildkite/ftr_configs.yml
View file

@ -555,8 +555,6 @@ enabled:
- x-pack/test/security_solution_api_integration/test_suites/detections_response/telemetry/trial_license_complete_tier/configs/serverless.config.ts
- x-pack/test/security_solution_api_integration/test_suites/detections_response/user_roles/trial_license_complete_tier/configs/ess.config.ts
- x-pack/test/security_solution_api_integration/test_suites/detections_response/user_roles/trial_license_complete_tier/configs/serverless.config.ts
- x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/trial_license_complete_tier/configs/ess.config.ts
- x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/trial_license_complete_tier/configs/serverless.config.ts
- x-pack/test/security_solution_api_integration/test_suites/genai/nlp_cleanup_task/trial_license_complete_tier/configs/serverless.config.ts
- x-pack/test/security_solution_api_integration/test_suites/genai/nlp_cleanup_task/basic_license_essentials_tier/configs/serverless.config.ts
- x-pack/test/security_solution_api_integration/test_suites/entity_analytics/risk_engine/trial_license_complete_tier/configs/ess.config.ts

2
x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts
View file

@ -38,8 +38,6 @@ export const ExecuteConnectorRequestBody = z.object({
alertsIndexPattern: z.string().optional(),
allow: z.array(z.string()).optional(),
allowReplacement: z.array(z.string()).optional(),
isEnabledKnowledgeBase: z.boolean().optional(),
isEnabledRAGAlerts: z.boolean().optional(),
replacements: Replacements,
size: z.number().optional(),
langSmithProject: z.string().optional(),

4
x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml
View file

@ -53,10 +53,6 @@ paths:
type: array
items:
type: string
isEnabledKnowledgeBase:
type: boolean
isEnabledRAGAlerts:
type: boolean
replacements:
$ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements'
size:

43
x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.test.tsx
View file

@ -17,31 +17,6 @@ describe('AlertsSettings', () => {
jest.clearAllMocks();
});
it('updates the knowledgeBase settings when the switch is toggled', () => {
const knowledgeBase: KnowledgeBaseConfig = {
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};
const setUpdatedKnowledgeBaseSettings = jest.fn();
render(
<AlertsSettings
knowledgeBase={knowledgeBase}
setUpdatedKnowledgeBaseSettings={setUpdatedKnowledgeBaseSettings}
/>
);
const alertsSwitch = screen.getByTestId('alertsSwitch');
fireEvent.click(alertsSwitch);
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
});
});
it('updates the knowledgeBase settings when the alerts range slider is changed', () => {
const setUpdatedKnowledgeBaseSettings = jest.fn();
const knowledgeBase: KnowledgeBaseConfig = {
@ -66,22 +41,4 @@ describe('AlertsSettings', () => {
latestAlerts: 10,
});
});
it('enables the alerts range slider when knowledgeBase.isEnabledRAGAlerts is true', () => {
const setUpdatedKnowledgeBaseSettings = jest.fn();
const knowledgeBase: KnowledgeBaseConfig = {
isEnabledRAGAlerts: true, // <-- true
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};
render(
<AlertsSettings
knowledgeBase={knowledgeBase}
setUpdatedKnowledgeBaseSettings={setUpdatedKnowledgeBaseSettings}
/>
);
expect(screen.getByTestId('alertsRange')).not.toBeDisabled();
});
});

23
x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.tsx
View file

@ -11,13 +11,11 @@ import {
EuiFormRow,
EuiRange,
EuiSpacer,
EuiSwitch,
EuiSwitchEvent,
EuiText,
useGeneratedHtmlId,
} from '@elastic/eui';
import { css } from '@emotion/react';
import React, { useCallback } from 'react';
import React from 'react';
import { KnowledgeBaseConfig } from '../../assistant/types';
import * as i18n from '../../knowledge_base/translations';
@ -36,16 +34,6 @@ interface Props {
const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSettings }: Props) => {
const inputRangeSliderId = useGeneratedHtmlId({ prefix: 'inputRangeSlider' });
const onEnableAlertsChange = useCallback(
(event: EuiSwitchEvent) => {
setUpdatedKnowledgeBaseSettings({
...knowledgeBase,
isEnabledRAGAlerts: event.target.checked,
});
},
[knowledgeBase, setUpdatedKnowledgeBaseSettings]
);
return (
<>
<EuiFormRow
@ -57,14 +45,7 @@ const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSetting
}
`}
>
<EuiSwitch
checked={knowledgeBase.isEnabledRAGAlerts}
compressed
data-test-subj="alertsSwitch"
label={i18n.ALERTS_LABEL}
onChange={onEnableAlertsChange}
showLabel={false}
/>
<></>
</EuiFormRow>
<EuiSpacer size="xs" />

64
x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.test.tsx
View file

@ -36,9 +36,7 @@ const apiConfig: Record<'openai' | 'bedrock' | 'gemini', ApiConfig> = {
};
const fetchConnectorArgs: FetchConnectorExecuteAction = {
isEnabledRAGAlerts: false,
apiConfig: apiConfig.openai,
isEnabledKnowledgeBase: true,
assistantStreamingEnabled: true,
http: mockHttp,
message: 'This is a test',
@ -75,7 +73,7 @@ describe('API tests', () => {
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...staticDefaults,
body: '{"model":"gpt-4","message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gen-ai","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
body: '{"model":"gpt-4","message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gen-ai","replacements":{}}',
}
);
});
@ -87,12 +85,12 @@ describe('API tests', () => {
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...streamingDefaults,
body: '{"model":"gpt-4","message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gen-ai","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
body: '{"model":"gpt-4","message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gen-ai","replacements":{}}',
}
);
});
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is bedrock and isEnabledKnowledgeBase is true', async () => {
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is bedrock', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
apiConfig: apiConfig.bedrock,
@ -104,31 +102,12 @@ describe('API tests', () => {
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...streamingDefaults,
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{}}',
}
);
});
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is bedrock and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
apiConfig: apiConfig.bedrock,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: true,
};
await fetchConnectorExecuteAction(testProps);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...streamingDefaults,
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
}
);
});
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is true', async () => {
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is gemini', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
apiConfig: apiConfig.gemini,
@ -140,35 +119,15 @@ describe('API tests', () => {
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...streamingDefaults,
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gemini","replacements":{}}',
}
);
});
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
apiConfig: apiConfig.gemini,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: true,
};
await fetchConnectorExecuteAction(testProps);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...streamingDefaults,
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
}
);
});
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is .bedrock and isEnabledKnowledgeBase is false', async () => {
it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is .bedrock', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
apiConfig: apiConfig.bedrock,
isEnabledKnowledgeBase: false,
};
await fetchConnectorExecuteAction(testProps);
@ -177,7 +136,7 @@ describe('API tests', () => {
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...streamingDefaults,
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":false}',
body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{}}',
}
);
});
@ -185,7 +144,6 @@ describe('API tests', () => {
it('calls the api with the expected optional request parameters', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
isEnabledRAGAlerts: true,
assistantStreamingEnabled: false,
alertsIndexPattern: '.alerts-security.alerts-default',
replacements: { auuid: 'real.hostname' },
@ -198,7 +156,7 @@ describe('API tests', () => {
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
...staticDefaults,
body: '{"model":"gpt-4","message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gen-ai","replacements":{"auuid":"real.hostname"},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":true,"alertsIndexPattern":".alerts-security.alerts-default","size":30}',
body: '{"model":"gpt-4","message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gen-ai","replacements":{"auuid":"real.hostname"},"alertsIndexPattern":".alerts-security.alerts-default","size":30}',
}
);
});
@ -239,7 +197,6 @@ describe('API tests', () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
isEnabledKnowledgeBase: false,
};
const result = await fetchConnectorExecuteAction(testProps);
@ -258,7 +215,6 @@ describe('API tests', () => {
});
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
isEnabledKnowledgeBase: false,
};
const result = await fetchConnectorExecuteAction(testProps);
@ -281,7 +237,7 @@ describe('API tests', () => {
expect(result).toEqual({ response: API_ERROR, isStream: false, isError: true });
});
it('returns the original when isEnabledKnowledgeBase is true, and `content` is not JSON', async () => {
it('returns the original when `content` is not JSON', async () => {
const response = 'plain text content';
(mockHttp.fetch as jest.Mock).mockResolvedValue({

7
x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx
View file

@ -15,9 +15,7 @@ export * from './prompts';
export interface FetchConnectorExecuteAction {
conversationId: string;
isEnabledRAGAlerts: boolean;
alertsIndexPattern?: string;
isEnabledKnowledgeBase: boolean;
assistantStreamingEnabled: boolean;
apiConfig: ApiConfig;
http: HttpSetup;
@ -40,9 +38,7 @@ export interface FetchConnectorExecuteResponse {
export const fetchConnectorExecuteAction = async ({
conversationId,
isEnabledRAGAlerts,
alertsIndexPattern,
isEnabledKnowledgeBase,
assistantStreamingEnabled,
http,
message,
@ -56,7 +52,6 @@ export const fetchConnectorExecuteAction = async ({
const isStream = assistantStreamingEnabled;
const optionalRequestParams = getOptionalRequestParams({
isEnabledRAGAlerts,
alertsIndexPattern,
size,
});
@ -68,8 +63,6 @@ export const fetchConnectorExecuteAction = async ({
conversationId,
actionTypeId: apiConfig.actionTypeId,
replacements,
isEnabledKnowledgeBase,
isEnabledRAGAlerts,
langSmithProject:
traceOptions?.langSmithProject === '' ? undefined : traceOptions?.langSmithProject,
langSmithApiKey:

13
x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.test.ts
View file

@ -163,21 +163,8 @@ describe('helpers', () => {
});
describe('getOptionalRequestParams', () => {
it('should return an empty object when alerts is false', () => {
const params = {
isEnabledRAGAlerts: false, // <-- false
alertsIndexPattern: 'indexPattern',
size: 10,
};
const result = getOptionalRequestParams(params);
expect(result).toEqual({});
});
it('should return the optional request params when alerts is true', () => {
const params = {
isEnabledRAGAlerts: true,
alertsIndexPattern: 'indexPattern',
size: 10,
};

7
x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.ts
View file

@ -100,22 +100,15 @@ interface OptionalRequestParams {
}
export const getOptionalRequestParams = ({
isEnabledRAGAlerts,
alertsIndexPattern,
size,
}: {
isEnabledRAGAlerts: boolean;
alertsIndexPattern?: string;
size?: number;
}): OptionalRequestParams => {
const optionalAlertsIndexPattern = alertsIndexPattern ? { alertsIndexPattern } : undefined;
const optionalSize = size ? { size } : undefined;
// the settings toggle must be enabled:
if (!isEnabledRAGAlerts) {
return {}; // don't send any optional params
}
return {
...optionalAlertsIndexPattern,
...optionalSize,

4
x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
View file

@ -80,7 +80,7 @@ import { Conversation } from '../assistant_context/types';
import { getGenAiConfig } from '../connectorland/helpers';
import { AssistantAnimatedIcon } from './assistant_animated_icon';
import { useFetchAnonymizationFields } from './api/anonymization_fields/use_fetch_anonymization_fields';
import { InstallKnowledgeBaseButton } from '../knowledge_base/install_knowledge_base_button';
import { SetupKnowledgeBaseButton } from '../knowledge_base/setup_knowledge_base_button';
import { useFetchPrompts } from './api/prompts/use_fetch_prompts';
export interface Props {
@ -798,7 +798,7 @@ const AssistantComponent: React.FC<Props> = ({
/>
</EuiFlexItem>
<EuiFlexItem grow={false}>
<InstallKnowledgeBaseButton />
<SetupKnowledgeBaseButton />
</EuiFlexItem>
</EuiFlexGroup>
</EuiPanel>

11
x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_message/index.tsx
View file

@ -41,10 +41,8 @@ export const useSendMessage = (): UseSendMessage => {
try {
return await fetchConnectorExecuteAction({
conversationId,
isEnabledRAGAlerts: knowledgeBase.isEnabledRAGAlerts, // settings toggle
alertsIndexPattern,
apiConfig,
isEnabledKnowledgeBase: knowledgeBase.isEnabledKnowledgeBase,
assistantStreamingEnabled,
http,
message,
@ -57,14 +55,7 @@ export const useSendMessage = (): UseSendMessage => {
setIsLoading(false);
}
},
[
alertsIndexPattern,
assistantStreamingEnabled,
knowledgeBase.isEnabledRAGAlerts,
knowledgeBase.isEnabledKnowledgeBase,
knowledgeBase.latestAlerts,
traceOptions,
]
[alertsIndexPattern, assistantStreamingEnabled, knowledgeBase.latestAlerts, traceOptions]
);
const cancelRequest = useCallback(() => {

86
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx
View file

@ -94,78 +94,6 @@ describe('Knowledge base settings', () => {
expect(getByTestId('esql-installed')).toBeInTheDocument();
expect(queryByTestId('install-esql')).not.toBeInTheDocument();
expect(getByTestId('esqlEnableButton')).toBeInTheDocument();
});
it('On click enable esql button, esql is enabled', () => {
(useKnowledgeBaseStatus as jest.Mock).mockImplementation(() => {
return {
data: {
elser_exists: true,
esql_exists: false,
index_exists: true,
pipeline_exists: true,
},
isLoading: false,
isFetching: false,
};
});
const { getByTestId, queryByTestId } = render(
<TestProviders>
<KnowledgeBaseSettings {...defaultProps} />
</TestProviders>
);
expect(queryByTestId('esql-installed')).not.toBeInTheDocument();
expect(getByTestId('install-esql')).toBeInTheDocument();
fireEvent.click(getByTestId('esqlEnableButton'));
expect(mockSetup).toHaveBeenCalledWith('esql');
});
it('On disable lang chain, set isEnabledKnowledgeBase to false', () => {
const { getByTestId } = render(
<TestProviders>
<KnowledgeBaseSettings {...defaultProps} />
</TestProviders>
);
fireEvent.click(getByTestId('isEnabledKnowledgeBaseSwitch'));
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
});
expect(mockSetup).not.toHaveBeenCalled();
});
it('On enable lang chain, set up with esql by default if ELSER exists', () => {
const { getByTestId } = render(
<TestProviders>
<KnowledgeBaseSettings
{...defaultProps}
knowledgeBase={{
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
}}
/>
</TestProviders>
);
fireEvent.click(getByTestId('isEnabledKnowledgeBaseSwitch'));
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
});
expect(mockSetup).toHaveBeenCalledWith('esql');
});
it('On disable knowledge base, call delete knowledge base setup', () => {
const { getByTestId, queryByTestId } = render(
<TestProviders>
<KnowledgeBaseSettings {...defaultProps} />
</TestProviders>
);
expect(queryByTestId('install-kb')).not.toBeInTheDocument();
expect(getByTestId('kb-installed')).toBeInTheDocument();
fireEvent.click(getByTestId('knowledgeBaseActionButton'));
expect(mockDelete).toHaveBeenCalledWith();
});
it('On enable knowledge base, call setup knowledge base setup', () => {
(useKnowledgeBaseStatus as jest.Mock).mockImplementation(() => {
@ -187,8 +115,8 @@ describe('Knowledge base settings', () => {
);
expect(queryByTestId('kb-installed')).not.toBeInTheDocument();
expect(getByTestId('install-kb')).toBeInTheDocument();
fireEvent.click(getByTestId('knowledgeBaseActionButton'));
expect(mockSetup).toHaveBeenCalledWith();
fireEvent.click(getByTestId('setupKnowledgeBaseButton'));
expect(mockSetup).toHaveBeenCalledWith('esql');
});
it('If elser does not exist, do not offer knowledge base', () => {
(useKnowledgeBaseStatus as jest.Mock).mockImplementation(() => {
@ -210,14 +138,4 @@ describe('Knowledge base settings', () => {
);
expect(queryByTestId('knowledgeBaseActionButton')).not.toBeInTheDocument();
});
it('renders the alerts settings', () => {
const { getByTestId } = render(
<TestProviders>
<KnowledgeBaseSettings {...defaultProps} />
</TestProviders>
);
expect(getByTestId('alertsSwitch')).toBeInTheDocument();
});
});

198
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx
View file

@ -11,17 +11,13 @@ import {
EuiTitle,
EuiText,
EuiHorizontalRule,
EuiLoadingSpinner,
EuiSpacer,
EuiSwitchEvent,
EuiLink,
EuiBetaBadge,
EuiFlexGroup,
EuiFlexItem,
EuiHealth,
EuiButtonEmpty,
EuiToolTip,
EuiSwitch,
EuiLink,
} from '@elastic/eui';
import { FormattedMessage } from '@kbn/i18n-react';
import { css } from '@emotion/react';
@ -30,12 +26,10 @@ import { AlertsSettings } from '../alerts/settings/alerts_settings';
import { useAssistantContext } from '../assistant_context';
import type { KnowledgeBaseConfig } from '../assistant/types';
import * as i18n from './translations';
import { useDeleteKnowledgeBase } from '../assistant/api/knowledge_base/use_delete_knowledge_base';
import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base';
const ESQL_RESOURCE = 'esql';
const KNOWLEDGE_BASE_INDEX_PATTERN_OLD = '.kibana-elastic-ai-assistant-kb';
const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-knowledge-base-(SPACE)';
interface Props {
@ -44,179 +38,77 @@ interface Props {
}
/**
* Knowledge Base Settings -- enable and disable LangChain integration, Knowledge Base, and ESQL KB Documents
* Knowledge Base Settings -- set up the Knowledge Base and configure RAG on alerts
*/
export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
({ knowledgeBase, setUpdatedKnowledgeBaseSettings }) => {
const {
assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault },
http,
} = useAssistantContext();
const { http } = useAssistantContext();
const {
data: kbStatus,
isLoading,
isFetching,
} = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http });
const { mutate: deleteKB, isLoading: isDeletingUpKB } = useDeleteKnowledgeBase({ http });
// Resource enabled state
const isElserEnabled = kbStatus?.elser_exists ?? false;
const isKnowledgeBaseEnabled = (kbStatus?.index_exists && kbStatus?.pipeline_exists) ?? false;
const isESQLEnabled = kbStatus?.esql_exists ?? false;
const isKnowledgeBaseSetup =
(isElserEnabled && isESQLEnabled && kbStatus?.index_exists && kbStatus?.pipeline_exists) ??
false;
const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false;
// Resource availability state
const isLoadingKb =
isLoading || isFetching || isSettingUpKB || isDeletingUpKB || isSetupInProgress;
const isKnowledgeBaseAvailable = knowledgeBase.isEnabledKnowledgeBase && kbStatus?.elser_exists;
const isESQLAvailable =
knowledgeBase.isEnabledKnowledgeBase && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled;
// Prevent enabling if elser doesn't exist, but always allow to disable
const isSwitchDisabled = enableKnowledgeBaseByDefault
? false
: !kbStatus?.elser_exists && !knowledgeBase.isEnabledKnowledgeBase;
const isLoadingKb = isLoading || isFetching || isSettingUpKB || isSetupInProgress;
// Calculated health state for EuiHealth component
const elserHealth = isElserEnabled ? 'success' : 'subdued';
const knowledgeBaseHealth = isKnowledgeBaseEnabled ? 'success' : 'subdued';
const knowledgeBaseHealth = isKnowledgeBaseSetup ? 'success' : 'subdued';
const esqlHealth = isESQLEnabled ? 'success' : 'subdued';
//////////////////////////////////////////////////////////////////////////////////////////
// Main `Knowledge Base` switch, which toggles the `isEnabledKnowledgeBase` UI feature toggle
// setting that is saved to localstorage
const onEnableAssistantLangChainChange = useCallback(
(event: EuiSwitchEvent) => {
setUpdatedKnowledgeBaseSettings({
...knowledgeBase,
isEnabledKnowledgeBase: event.target.checked,
});
// Main `Knowledge Base` setup button
const onSetupKnowledgeBaseButtonClick = useCallback(() => {
setupKB(ESQL_RESOURCE);
}, [setupKB]);
// If enabling and ELSER exists or automatic KB setup FF is enabled, try to set up automatically
if (event.target.checked && (enableKnowledgeBaseByDefault || kbStatus?.elser_exists)) {
setupKB(ESQL_RESOURCE);
}
},
[
enableKnowledgeBaseByDefault,
kbStatus?.elser_exists,
knowledgeBase,
setUpdatedKnowledgeBaseSettings,
setupKB,
]
);
const isEnabledKnowledgeBaseSwitch = useMemo(() => {
return isLoadingKb ? (
<EuiLoadingSpinner size="s" />
const setupKnowledgeBaseButton = useMemo(() => {
return isKnowledgeBaseSetup ? (
<></>
) : (
<EuiToolTip content={isSwitchDisabled && i18n.KNOWLEDGE_BASE_TOOLTIP} position={'right'}>
<EuiSwitch
showLabel={false}
data-test-subj="isEnabledKnowledgeBaseSwitch"
disabled={isSwitchDisabled}
checked={knowledgeBase.isEnabledKnowledgeBase}
onChange={onEnableAssistantLangChainChange}
label={i18n.KNOWLEDGE_BASE_LABEL}
compressed
/>
</EuiToolTip>
<EuiButtonEmpty
color={'primary'}
data-test-subj={'setupKnowledgeBaseButton'}
onClick={onSetupKnowledgeBaseButtonClick}
size="xs"
isLoading={isLoadingKb}
>
{i18n.SETUP_KNOWLEDGE_BASE_BUTTON}
</EuiButtonEmpty>
);
}, [
isLoadingKb,
isSwitchDisabled,
knowledgeBase.isEnabledKnowledgeBase,
onEnableAssistantLangChainChange,
]);
}, [isKnowledgeBaseSetup, isLoadingKb, onSetupKnowledgeBaseButtonClick]);
//////////////////////////////////////////////////////////////////////////////////////////
// Knowledge Base Resource
const onEnableKB = useCallback(
(enabled: boolean) => {
if (enabled) {
setupKB();
} else {
deleteKB();
}
},
[deleteKB, setupKB]
);
const knowledgeBaseActionButton = useMemo(() => {
return isLoadingKb || !isKnowledgeBaseAvailable ? (
<></>
) : (
<EuiButtonEmpty
color={isKnowledgeBaseEnabled ? 'danger' : 'primary'}
flush="left"
data-test-subj={'knowledgeBaseActionButton'}
onClick={() => onEnableKB(!isKnowledgeBaseEnabled)}
size="xs"
>
{isKnowledgeBaseEnabled
? i18n.KNOWLEDGE_BASE_DELETE_BUTTON
: i18n.KNOWLEDGE_BASE_INIT_BUTTON}
</EuiButtonEmpty>
);
}, [isKnowledgeBaseAvailable, isKnowledgeBaseEnabled, isLoadingKb, onEnableKB]);
const knowledgeBaseDescription = useMemo(() => {
return isKnowledgeBaseEnabled ? (
return isKnowledgeBaseSetup ? (
<span data-test-subj="kb-installed">
{i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(
enableKnowledgeBaseByDefault
? KNOWLEDGE_BASE_INDEX_PATTERN
: KNOWLEDGE_BASE_INDEX_PATTERN_OLD
)}{' '}
{knowledgeBaseActionButton}
{i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(KNOWLEDGE_BASE_INDEX_PATTERN)}
</span>
) : (
<span data-test-subj="install-kb">
{i18n.KNOWLEDGE_BASE_DESCRIPTION} {knowledgeBaseActionButton}
</span>
<span data-test-subj="install-kb">{i18n.KNOWLEDGE_BASE_DESCRIPTION}</span>
);
}, [enableKnowledgeBaseByDefault, isKnowledgeBaseEnabled, knowledgeBaseActionButton]);
}, [isKnowledgeBaseSetup]);
//////////////////////////////////////////////////////////////////////////////////////////
// ESQL Resource
const onEnableESQL = useCallback(
(enabled: boolean) => {
if (enabled) {
setupKB(ESQL_RESOURCE);
} else {
deleteKB(ESQL_RESOURCE);
}
},
[deleteKB, setupKB]
);
const esqlActionButton = useMemo(() => {
return isLoadingKb || !isESQLAvailable ? (
<></>
) : (
<EuiButtonEmpty
color={isESQLEnabled ? 'danger' : 'primary'}
flush="left"
data-test-subj="esqlEnableButton"
onClick={() => onEnableESQL(!isESQLEnabled)}
size="xs"
>
{isESQLEnabled ? i18n.KNOWLEDGE_BASE_DELETE_BUTTON : i18n.KNOWLEDGE_BASE_INIT_BUTTON}
</EuiButtonEmpty>
);
}, [isLoadingKb, isESQLAvailable, isESQLEnabled, onEnableESQL]);
const esqlDescription = useMemo(() => {
return isESQLEnabled ? (
<span data-test-subj="esql-installed">
{i18n.ESQL_DESCRIPTION_INSTALLED} {esqlActionButton}
</span>
<span data-test-subj="esql-installed">{i18n.ESQL_DESCRIPTION_INSTALLED}</span>
) : (
<span data-test-subj="install-esql">
{i18n.ESQL_DESCRIPTION} {esqlActionButton}
</span>
<span data-test-subj="install-esql">{i18n.ESQL_DESCRIPTION}</span>
);
}, [esqlActionButton, isESQLEnabled]);
}, [isESQLEnabled]);
return (
<>
@ -255,7 +147,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
}
`}
>
{isEnabledKnowledgeBaseSwitch}
{setupKnowledgeBaseButton}
</EuiFormRow>
<EuiSpacer size="s" />
@ -277,30 +169,8 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
`}
>
<FormattedMessage
defaultMessage="Configure ELSER within {machineLearning} to get started. {seeDocs}"
defaultMessage="Elastic Learned Sparse EncodeR - or ELSER - is a retrieval model trained by Elastic for performing semantic search."
id="xpack.elasticAssistant.assistant.settings.knowledgeBasedSettings.knowledgeBaseDescription"
values={{
machineLearning: (
<EuiLink
external
href={http.basePath.prepend('/app/ml/trained_models')}
target="_blank"
>
{i18n.KNOWLEDGE_BASE_ELSER_MACHINE_LEARNING}
</EuiLink>
),
seeDocs: (
<EuiLink
external
href={
'https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-elser.html#download-deploy-elser'
}
target="_blank"
>
{i18n.KNOWLEDGE_BASE_ELSER_SEE_DOCS}
</EuiLink>
),
}}
/>
</EuiText>
</div>

198
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx
View file

@ -10,16 +10,12 @@ import {
EuiFormRow,
EuiText,
EuiHorizontalRule,
EuiLoadingSpinner,
EuiSpacer,
EuiSwitchEvent,
EuiLink,
EuiFlexGroup,
EuiFlexItem,
EuiHealth,
EuiButtonEmpty,
EuiToolTip,
EuiSwitch,
EuiPanel,
} from '@elastic/eui';
import { FormattedMessage } from '@kbn/i18n-react';
@ -28,7 +24,6 @@ import { css } from '@emotion/react';
import { AlertsSettings } from '../alerts/settings/alerts_settings';
import { useAssistantContext } from '../assistant_context';
import * as i18n from './translations';
import { useDeleteKnowledgeBase } from '../assistant/api/knowledge_base/use_delete_knowledge_base';
import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base';
import {
@ -40,18 +35,13 @@ import { AssistantSettingsBottomBar } from '../assistant/settings/assistant_sett
import { SETTINGS_UPDATED_TOAST_TITLE } from '../assistant/settings/translations';
const ESQL_RESOURCE = 'esql';
const KNOWLEDGE_BASE_INDEX_PATTERN_OLD = '.kibana-elastic-ai-assistant-kb';
const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-knowledge-base-(SPACE)';
/**
* Knowledge Base Settings -- enable and disable LangChain integration, Knowledge Base, and ESQL KB Documents
* Knowledge Base Settings -- set up the Knowledge Base and configure RAG on alerts
*/
export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
const {
assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault },
http,
toasts,
} = useAssistantContext();
const { http, toasts } = useAssistantContext();
const [hasPendingChanges, setHasPendingChanges] = useState(false);
const { knowledgeBase, setUpdatedKnowledgeBaseSettings, resetSettings, saveSettings } =
@ -98,165 +88,67 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
isFetching,
} = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http });
const { mutate: deleteKB, isLoading: isDeletingUpKB } = useDeleteKnowledgeBase({ http });
// Resource enabled state
const isElserEnabled = kbStatus?.elser_exists ?? false;
const isKnowledgeBaseEnabled = (kbStatus?.index_exists && kbStatus?.pipeline_exists) ?? false;
const isESQLEnabled = kbStatus?.esql_exists ?? false;
const isKnowledgeBaseSetup =
(isElserEnabled && isESQLEnabled && kbStatus?.index_exists && kbStatus?.pipeline_exists) ??
false;
const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false;
// Resource availability state
const isLoadingKb =
isLoading || isFetching || isSettingUpKB || isDeletingUpKB || isSetupInProgress;
const isKnowledgeBaseAvailable = knowledgeBase.isEnabledKnowledgeBase && kbStatus?.elser_exists;
const isESQLAvailable =
knowledgeBase.isEnabledKnowledgeBase && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled;
// Prevent enabling if elser doesn't exist, but always allow to disable
const isSwitchDisabled = enableKnowledgeBaseByDefault
? false
: !kbStatus?.elser_exists && !knowledgeBase.isEnabledKnowledgeBase;
const isLoadingKb = isLoading || isFetching || isSettingUpKB || isSetupInProgress;
// Calculated health state for EuiHealth component
const elserHealth = isElserEnabled ? 'success' : 'subdued';
const knowledgeBaseHealth = isKnowledgeBaseEnabled ? 'success' : 'subdued';
const knowledgeBaseHealth = isKnowledgeBaseSetup ? 'success' : 'subdued';
const esqlHealth = isESQLEnabled ? 'success' : 'subdued';
//////////////////////////////////////////////////////////////////////////////////////////
// Main `Knowledge Base` switch, which toggles the `isEnabledKnowledgeBase` UI feature toggle
// setting that is saved to localstorage
const onEnableAssistantLangChainChange = useCallback(
(event: EuiSwitchEvent) => {
handleUpdateKnowledgeBaseSettings({
...knowledgeBase,
isEnabledKnowledgeBase: event.target.checked,
});
// Main `Knowledge Base` setup button
const onSetupKnowledgeBaseButtonClick = useCallback(() => {
setupKB(ESQL_RESOURCE);
}, [setupKB]);
// If enabling and ELSER exists or automatic KB setup FF is enabled, try to set up automatically
if (event.target.checked && (enableKnowledgeBaseByDefault || kbStatus?.elser_exists)) {
setupKB(ESQL_RESOURCE);
}
},
[
enableKnowledgeBaseByDefault,
handleUpdateKnowledgeBaseSettings,
kbStatus?.elser_exists,
knowledgeBase,
setupKB,
]
);
const isEnabledKnowledgeBaseSwitch = useMemo(() => {
return isLoadingKb ? (
<EuiLoadingSpinner size="s" />
const setupKnowledgeBaseButton = useMemo(() => {
return isKnowledgeBaseSetup ? (
<></>
) : (
<EuiToolTip content={isSwitchDisabled && i18n.KNOWLEDGE_BASE_TOOLTIP} position={'right'}>
<EuiSwitch
showLabel={false}
data-test-subj="isEnabledKnowledgeBaseSwitch"
disabled={isSwitchDisabled}
checked={knowledgeBase.isEnabledKnowledgeBase}
onChange={onEnableAssistantLangChainChange}
label={i18n.KNOWLEDGE_BASE_LABEL}
compressed
/>
</EuiToolTip>
<EuiButtonEmpty
color={'primary'}
data-test-subj={'setupKnowledgeBaseButton'}
onClick={onSetupKnowledgeBaseButtonClick}
size="xs"
isLoading={isLoadingKb}
>
{i18n.SETUP_KNOWLEDGE_BASE_BUTTON}
</EuiButtonEmpty>
);
}, [
isLoadingKb,
isSwitchDisabled,
knowledgeBase.isEnabledKnowledgeBase,
onEnableAssistantLangChainChange,
]);
}, [isKnowledgeBaseSetup, isLoadingKb, onSetupKnowledgeBaseButtonClick]);
//////////////////////////////////////////////////////////////////////////////////////////
// Knowledge Base Resource
const onEnableKB = useCallback(
(enabled: boolean) => {
if (enabled) {
setupKB();
} else {
deleteKB();
}
},
[deleteKB, setupKB]
);
const knowledgeBaseActionButton = useMemo(() => {
return isLoadingKb || !isKnowledgeBaseAvailable ? (
<></>
) : (
<EuiButtonEmpty
color={isKnowledgeBaseEnabled ? 'danger' : 'primary'}
flush="left"
data-test-subj={'knowledgeBaseActionButton'}
onClick={() => onEnableKB(!isKnowledgeBaseEnabled)}
size="xs"
>
{isKnowledgeBaseEnabled
? i18n.KNOWLEDGE_BASE_DELETE_BUTTON
: i18n.KNOWLEDGE_BASE_INIT_BUTTON}
</EuiButtonEmpty>
);
}, [isKnowledgeBaseAvailable, isKnowledgeBaseEnabled, isLoadingKb, onEnableKB]);
const knowledgeBaseDescription = useMemo(() => {
return isKnowledgeBaseEnabled ? (
return isKnowledgeBaseSetup ? (
<span data-test-subj="kb-installed">
{i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(
enableKnowledgeBaseByDefault
? KNOWLEDGE_BASE_INDEX_PATTERN
: KNOWLEDGE_BASE_INDEX_PATTERN_OLD
)}{' '}
{knowledgeBaseActionButton}
{i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(KNOWLEDGE_BASE_INDEX_PATTERN)}
</span>
) : (
<span data-test-subj="install-kb">
{i18n.KNOWLEDGE_BASE_DESCRIPTION} {knowledgeBaseActionButton}
</span>
<span data-test-subj="install-kb">{i18n.KNOWLEDGE_BASE_DESCRIPTION}</span>
);
}, [enableKnowledgeBaseByDefault, isKnowledgeBaseEnabled, knowledgeBaseActionButton]);
}, [isKnowledgeBaseSetup]);
//////////////////////////////////////////////////////////////////////////////////////////
// ESQL Resource
const onEnableESQL = useCallback(
(enabled: boolean) => {
if (enabled) {
setupKB(ESQL_RESOURCE);
} else {
deleteKB(ESQL_RESOURCE);
}
},
[deleteKB, setupKB]
);
const esqlActionButton = useMemo(() => {
return isLoadingKb || !isESQLAvailable ? (
<></>
) : (
<EuiButtonEmpty
color={isESQLEnabled ? 'danger' : 'primary'}
flush="left"
data-test-subj="esqlEnableButton"
onClick={() => onEnableESQL(!isESQLEnabled)}
size="xs"
>
{isESQLEnabled ? i18n.KNOWLEDGE_BASE_DELETE_BUTTON : i18n.KNOWLEDGE_BASE_INIT_BUTTON}
</EuiButtonEmpty>
);
}, [isLoadingKb, isESQLAvailable, isESQLEnabled, onEnableESQL]);
const esqlDescription = useMemo(() => {
return isESQLEnabled ? (
<span data-test-subj="esql-installed">
{i18n.ESQL_DESCRIPTION_INSTALLED} {esqlActionButton}
</span>
<span data-test-subj="esql-installed">{i18n.ESQL_DESCRIPTION_INSTALLED}</span>
) : (
<span data-test-subj="install-esql">
{i18n.ESQL_DESCRIPTION} {esqlActionButton}
</span>
<span data-test-subj="install-esql">{i18n.ESQL_DESCRIPTION}</span>
);
}, [esqlActionButton, isESQLEnabled]);
}, [isESQLEnabled]);
return (
<EuiPanel hasShadow={false} hasBorder paddingSize="l">
@ -288,7 +180,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
}
`}
>
{isEnabledKnowledgeBaseSwitch}
{setupKnowledgeBaseButton}
</EuiFormRow>
<EuiSpacer size="s" />
@ -310,30 +202,8 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => {
`}
>
<FormattedMessage
defaultMessage="Configure ELSER within {machineLearning} to get started. {seeDocs}"
id="xpack.elasticAssistant.assistant.settings.knowledgeBasedSettings.knowledgeBaseDescription"
values={{
machineLearning: (
<EuiLink
external
href={http.basePath.prepend('/app/ml/trained_models')}
target="_blank"
>
{i18n.KNOWLEDGE_BASE_ELSER_MACHINE_LEARNING}
</EuiLink>
),
seeDocs: (
<EuiLink
external
href={
'https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-elser.html#download-deploy-elser'
}
target="_blank"
>
{i18n.KNOWLEDGE_BASE_ELSER_SEE_DOCS}
</EuiLink>
),
}}
defaultMessage="Elastic Learned Sparse EncodeR - or ELSER - is a retrieval model trained by Elastic for performing semantic search."
id="xpack.elasticAssistant.assistant.settings.knowledgeBasedSettingsManagement.knowledgeBaseDescription"
/>
</EuiText>
</div>

19
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx → x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx
View file

@ -16,16 +16,11 @@ import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_know
const ESQL_RESOURCE = 'esql';
/**
* Self-contained component that renders a button to install the knowledge base.
* Self-contained component that renders a button to set up the knowledge base.
*
* Only renders if `assistantKnowledgeBaseByDefault` feature flag is enabled.
*/
export const InstallKnowledgeBaseButton: React.FC = React.memo(() => {
const {
assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault },
http,
toasts,
} = useAssistantContext();
export const SetupKnowledgeBaseButton: React.FC = React.memo(() => {
const { http, toasts } = useAssistantContext();
const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
@ -41,24 +36,24 @@ export const InstallKnowledgeBaseButton: React.FC = React.memo(() => {
setupKB(ESQL_RESOURCE);
}, [setupKB]);
if (!enableKnowledgeBaseByDefault || isSetupComplete) {
if (isSetupComplete) {
return null;
}
return (
<EuiButton
color="primary"
data-test-subj="install-knowledge-base-button"
data-test-subj="setup-knowledge-base-button"
fill
isLoading={isSetupInProgress}
iconType="importAction"
onClick={onInstallKnowledgeBase}
>
{i18n.translate('xpack.elasticAssistant.knowledgeBase.installKnowledgeBaseButton', {
defaultMessage: 'Install Knowledge Base',
defaultMessage: 'Setup Knowledge Base',
})}
</EuiButton>
);
});
InstallKnowledgeBaseButton.displayName = 'InstallKnowledgeBaseButton';
SetupKnowledgeBaseButton.displayName = 'SetupKnowledgeBaseButton';

7
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/translations.ts
View file

@ -80,6 +80,13 @@ export const KNOWLEDGE_BASE_LABEL = i18n.translate(
}
);
export const SETUP_KNOWLEDGE_BASE_BUTTON = i18n.translate(
'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.setupKnowledgeBaseButton',
{
defaultMessage: 'Setup',
}
);
export const KNOWLEDGE_BASE_TOOLTIP = i18n.translate(
'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettings.knowledgeBaseTooltip',
{

4
x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts
View file

@ -123,10 +123,10 @@ const createElasticAssistantRequestContextMock = (
() => clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient
) as unknown as jest.MockInstance<
Promise<AIAssistantKnowledgeBaseDataClient | null>,
[boolean],
[],
unknown
> &
((initializeKnowledgeBase: boolean) => Promise<AIAssistantKnowledgeBaseDataClient | null>),
(() => Promise<AIAssistantKnowledgeBaseDataClient | null>),
getCurrentUser: jest.fn(),
getServerBasePath: jest.fn(),
getSpaceId: jest.fn(),

22
x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts
View file

@ -98,6 +98,7 @@ const mockUser1 = {
describe('AI Assistant Service', () => {
let pluginStop$: Subject<void>;
let assistantServiceOpts: AIAssistantServiceOpts;
let ml: MlPluginSetup;
beforeEach(() => {
jest.resetAllMocks();
@ -110,12 +111,16 @@ describe('AI Assistant Service', () => {
);
clusterClient.indices.getAlias.mockImplementation(async () => GetAliasResponse);
clusterClient.indices.getDataStream.mockImplementation(async () => GetDataStreamResponse);
ml = mlPluginMock.createSetupContract() as unknown as MlPluginSetup; // Missing SharedServices mock, so manually mocking trainedModelsProvider
ml.trainedModelsProvider = jest.fn().mockImplementation(() => ({
getELSER: jest.fn().mockImplementation(() => '.elser_model_2'),
}));
assistantServiceOpts = {
logger,
elasticsearchClientPromise: Promise.resolve(clusterClient),
pluginStop$,
kibanaVersion: '8.8.0',
ml: mlPluginMock.createSetupContract() as unknown as MlPluginSetup, // Missing SharedServices mock
ml,
taskManager: taskManagerMock.createSetup(),
};
});
@ -136,10 +141,11 @@ describe('AI Assistant Service', () => {
expect(assistantService.isInitialized()).toEqual(true);
expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(4);
expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(5);
const expectedTemplates = [
'.kibana-elastic-ai-assistant-component-template-conversations',
'.kibana-elastic-ai-assistant-component-template-knowledge-base',
'.kibana-elastic-ai-assistant-component-template-prompts',
'.kibana-elastic-ai-assistant-component-template-anonymization-fields',
'.kibana-elastic-ai-assistant-component-template-attack-discovery',
@ -634,12 +640,13 @@ describe('AI Assistant Service', () => {
'AI Assistant service initialized',
async () => assistantService.isInitialized() === true
);
expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(6);
expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(7);
const expectedTemplates = [
'.kibana-elastic-ai-assistant-component-template-conversations',
'.kibana-elastic-ai-assistant-component-template-conversations',
'.kibana-elastic-ai-assistant-component-template-conversations',
'.kibana-elastic-ai-assistant-component-template-knowledge-base',
'.kibana-elastic-ai-assistant-component-template-prompts',
'.kibana-elastic-ai-assistant-component-template-anonymization-fields',
'.kibana-elastic-ai-assistant-component-template-attack-discovery',
@ -667,11 +674,12 @@ describe('AI Assistant Service', () => {
async () => (await getSpaceResourcesInitialized(assistantService)) === true
);
expect(clusterClient.indices.putIndexTemplate).toHaveBeenCalledTimes(6);
expect(clusterClient.indices.putIndexTemplate).toHaveBeenCalledTimes(7);
const expectedTemplates = [
'.kibana-elastic-ai-assistant-index-template-conversations',
'.kibana-elastic-ai-assistant-index-template-conversations',
'.kibana-elastic-ai-assistant-index-template-conversations',
'.kibana-elastic-ai-assistant-index-template-knowledge-base',
'.kibana-elastic-ai-assistant-index-template-prompts',
'.kibana-elastic-ai-assistant-index-template-anonymization-fields',
'.kibana-elastic-ai-assistant-index-template-attack-discovery',
@ -698,7 +706,7 @@ describe('AI Assistant Service', () => {
async () => (await getSpaceResourcesInitialized(assistantService)) === true
);
expect(clusterClient.indices.putSettings).toHaveBeenCalledTimes(6);
expect(clusterClient.indices.putSettings).toHaveBeenCalledTimes(7);
});
test('should retry updating index mappings for existing indices for transient ES errors', async () => {
@ -718,7 +726,7 @@ describe('AI Assistant Service', () => {
async () => (await getSpaceResourcesInitialized(assistantService)) === true
);
expect(clusterClient.indices.putMapping).toHaveBeenCalledTimes(6);
expect(clusterClient.indices.putMapping).toHaveBeenCalledTimes(7);
});
test('should retry creating concrete index for transient ES errors', async () => {
@ -751,7 +759,7 @@ describe('AI Assistant Service', () => {
async () => (await getSpaceResourcesInitialized(assistantService)) === true
);
expect(clusterClient.indices.createDataStream).toHaveBeenCalledTimes(5);
expect(clusterClient.indices.createDataStream).toHaveBeenCalledTimes(6);
});
});
});

68
x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts
View file

@ -66,8 +66,6 @@ export type CreateDataStream = (params: {
export class AIAssistantService {
private initialized: boolean;
// Temporary 'feature flag' to determine if we should initialize the knowledge base, toggled when accessing data client
private initializeKnowledgeBase: boolean = false;
private isInitializing: boolean = false;
private getElserId: GetElser;
private conversationsDataStream: DataStreamSpacesAdapter;
@ -124,6 +122,7 @@ export class AIAssistantService {
public getIsKBSetupInProgress() {
return this.isKBSetupInProgress;
}
public setIsKBSetupInProgress(isInProgress: boolean) {
this.isKBSetupInProgress = isInProgress;
}
@ -172,34 +171,32 @@ export class AIAssistantService {
pluginStop$: this.options.pluginStop$,
});
if (this.initializeKnowledgeBase) {
await this.knowledgeBaseDataStream.install({
esClient,
logger: this.options.logger,
pluginStop$: this.options.pluginStop$,
});
await this.knowledgeBaseDataStream.install({
esClient,
logger: this.options.logger,
pluginStop$: this.options.pluginStop$,
});
// TODO: Pipeline creation is temporary as we'll be moving to semantic_text field once available in ES
const pipelineCreated = await pipelineExists({
// TODO: Pipeline creation is temporary as we'll be moving to semantic_text field once available in ES
const pipelineCreated = await pipelineExists({
esClient,
id: this.resourceNames.pipelines.knowledgeBase,
});
if (!pipelineCreated) {
this.options.logger.debug(
`Installing ingest pipeline - ${this.resourceNames.pipelines.knowledgeBase}`
);
const response = await createPipeline({
esClient,
id: this.resourceNames.pipelines.knowledgeBase,
modelId: await this.getElserId(),
});
if (!pipelineCreated) {
this.options.logger.debug(
`Installing ingest pipeline - ${this.resourceNames.pipelines.knowledgeBase}`
);
const response = await createPipeline({
esClient,
id: this.resourceNames.pipelines.knowledgeBase,
modelId: await this.getElserId(),
});
this.options.logger.debug(`Installed ingest pipeline: ${response}`);
} else {
this.options.logger.debug(
`Ingest pipeline already exists - ${this.resourceNames.pipelines.knowledgeBase}`
);
}
this.options.logger.debug(`Installed ingest pipeline: ${response}`);
} else {
this.options.logger.debug(
`Ingest pipeline already exists - ${this.resourceNames.pipelines.knowledgeBase}`
);
}
await this.promptsDataStream.install({
@ -330,15 +327,8 @@ export class AIAssistantService {
}
public async createAIAssistantKnowledgeBaseDataClient(
opts: CreateAIAssistantClientParams & { initializeKnowledgeBase: boolean }
opts: CreateAIAssistantClientParams
): Promise<AIAssistantKnowledgeBaseDataClient | null> {
// Note: Due to plugin lifecycle and feature flag registration timing, we need to pass in the feature flag here
// Remove this param and initialization when the `assistantKnowledgeBaseByDefault` feature flag is removed
if (opts.initializeKnowledgeBase) {
this.initializeKnowledgeBase = true;
await this.initializeResources();
}
const res = await this.checkResourcesInstallation(opts);
if (res === null) {
@ -446,13 +436,11 @@ export class AIAssistantService {
await this.conversationsDataStream.installSpace(spaceId);
}
if (this.initializeKnowledgeBase) {
const knowledgeBaseIndexName = await this.knowledgeBaseDataStream.getInstalledSpaceName(
spaceId
);
if (!knowledgeBaseIndexName) {
await this.knowledgeBaseDataStream.installSpace(spaceId);
}
const knowledgeBaseIndexName = await this.knowledgeBaseDataStream.getInstalledSpaceName(
spaceId
);
if (!knowledgeBaseIndexName) {
await this.knowledgeBaseDataStream.installSpace(spaceId);
}
const promptsIndexName = await this.promptsDataStream.getInstalledSpaceName(spaceId);

14
x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts
View file

@ -67,10 +67,16 @@ export class DocumentsDataWriter implements DocumentsDataWriter {
return { errors: [], docs_created: [], docs_deleted: [], docs_updated: [], took: 0 };
}
const { errors, items, took } = await this.options.esClient.bulk({
refresh: 'wait_for',
body: await this.buildBulkOperations(params),
});
const { errors, items, took } = await this.options.esClient.bulk(
{
refresh: 'wait_for',
body: await this.buildBulkOperations(params),
},
{
// Increasing timout to 2min as KB docs were failing to load after 30s
requestTimeout: 120000,
}
);
return {
errors: errors ? this.formatErrorsResponse(items) : [],

1
x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts
View file

@ -97,7 +97,6 @@ const esStoreMock = new ElasticsearchStore(
);
const defaultProps: AgentExecutorParams<true> = {
actionsClient,
isEnabledKnowledgeBase: true,
connectorId: mockConnectorId,
esClient: esClientMock,
esStore: esStoreMock,

3
x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts
View file

@ -35,7 +35,6 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
abortSignal,
actionsClient,
alertsIndexPattern,
isEnabledKnowledgeBase,
assistantTools = [],
connectorId,
esClient,
@ -105,7 +104,7 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
anonymizationFields,
chain,
esClient,
isEnabledKnowledgeBase,
isEnabledKnowledgeBase: true,
llm,
logger,
modelExists,

1
x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts
View file

@ -37,7 +37,6 @@ export interface AgentExecutorParams<T extends boolean> {
abortSignal?: AbortSignal;
alertsIndexPattern?: string;
actionsClient: PublicMethodsOf<ActionsClient>;
isEnabledKnowledgeBase: boolean;
assistantTools?: AssistantTool[];
connectorId: string;
conversationId?: string;

14
x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
View file

@ -29,7 +29,6 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
abortSignal,
actionsClient,
alertsIndexPattern,
isEnabledKnowledgeBase,
assistantTools = [],
connectorId,
conversationId,
@ -87,6 +86,9 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
// Create a chain that uses the ELSER backed ElasticsearchStore, override k=10 for esql query generation for now
const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever(10));
// Check if KB is available
const isEnabledKnowledgeBase = (await dataClients?.kbDataClient?.isModelDeployed()) ?? false;
// Fetch any applicable tools that the source plugin may have registered
const assistantToolParams: AssistantToolParams = {
alertsIndexPattern,
@ -137,7 +139,15 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
const inputs = { input: latestMessage[0]?.content as string };
if (isStream) {
return streamGraph({ apmTracer, assistantGraph, inputs, logger, onLlmResponse, request });
return streamGraph({
apmTracer,
assistantGraph,
inputs,
logger,
onLlmResponse,
request,
traceOptions,
});
}
const graphResponse = await invokeGraph({

8
x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts
View file

@ -21,7 +21,6 @@ export const AGENT_NODE = 'agent';
export const AGENT_NODE_TAG = 'agent_run';
const NO_HISTORY = '[No existing knowledge history]';
/**
* Node to run the agent
*
@ -40,17 +39,10 @@ export const runAgent = async ({
}: RunAgentParams) => {
logger.debug(() => `Node state:\n${JSON.stringify(state, null, 2)}`);
const knowledgeHistory = await dataClients?.kbDataClient?.getKnowledgeBaseDocuments({
kbResource: 'user',
required: true,
query: '',
});
const agentOutcome = await agentRunnable.withConfig({ tags: [AGENT_NODE_TAG] }).invoke(
{
...state,
chat_history: state.messages, // TODO: Message de-dupe with ...state spread
knowledge_history: JSON.stringify(knowledgeHistory?.length ? knowledgeHistory : NO_HISTORY),
},
config
);

10
x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts
View file

@ -8,10 +8,7 @@
import { ChatPromptTemplate } from '@langchain/core/prompts';
export const openAIFunctionAgentPrompt = ChatPromptTemplate.fromMessages([
[
'system',
'You are a helpful assistant\n\nUse the below context as a sample of information about the user from their knowledge base:\n\n```{knowledge_history}```',
],
['system', 'You are a helpful assistant'],
['placeholder', '{chat_history}'],
['human', '{input}'],
['placeholder', '{agent_scratchpad}'],
@ -52,8 +49,5 @@ export const structuredChatAgentPrompt = ChatPromptTemplate.fromMessages([
'Begin! Reminder to ALWAYS respond with a valid json blob of a single action with no additional output. When using tools, ALWAYS input the expected JSON schema args. Your answer will be parsed as JSON, so never use double quotes within the output and instead use backticks. Single quotes may be used, such as apostrophes. Response format is Action:```$JSON_BLOB```then Observation',
],
['placeholder', '{chat_history}'],
[
'human',
'Use the below context as a sample of information about the user from their knowledge base:\n\n```\n{knowledge_history}\n```\n\n{input}\n\n{agent_scratchpad}\n(reminder to respond in a JSON blob with no additional output no matter what)',
],
['human', '{input}\n\n{agent_scratchpad}\n\n(reminder to respond in a JSON blob no matter what)'],
]);

28
x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts
View file

@ -71,26 +71,12 @@ export const KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT: EventTypeOpts<{
};
export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{
isEnabledKnowledgeBase: boolean;
isEnabledRAGAlerts: boolean;
assistantStreamingEnabled: boolean;
actionTypeId: string;
model?: string;
}> = {
eventType: 'invoke_assistant_success',
schema: {
isEnabledKnowledgeBase: {
type: 'boolean',
_meta: {
description: 'Is Knowledge Base enabled',
},
},
isEnabledRAGAlerts: {
type: 'boolean',
_meta: {
description: 'Is RAG Alerts enabled',
},
},
assistantStreamingEnabled: {
type: 'boolean',
_meta: {
@ -115,8 +101,6 @@ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{
export const INVOKE_ASSISTANT_ERROR_EVENT: EventTypeOpts<{
errorMessage: string;
isEnabledKnowledgeBase: boolean;
isEnabledRAGAlerts: boolean;
assistantStreamingEnabled: boolean;
actionTypeId: string;
model?: string;
@ -129,18 +113,6 @@ export const INVOKE_ASSISTANT_ERROR_EVENT: EventTypeOpts<{
description: 'Error message from Elasticsearch',
},
},
isEnabledKnowledgeBase: {
type: 'boolean',
_meta: {
description: 'Is Knowledge Base enabled',
},
},
isEnabledRAGAlerts: {
type: 'boolean',
_meta: {
description: 'Is RAG Alerts enabled',
},
},
assistantStreamingEnabled: {
type: 'boolean',
_meta: {

170
x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts
View file

@ -5,27 +5,26 @@
* 2.0.
*/
import { ElasticsearchClient, IRouter, KibanaRequest, Logger } from '@kbn/core/server';
import type { ActionsClient } from '@kbn/actions-plugin/server';
import { BaseMessage } from '@langchain/core/messages';
import { IRouter } from '@kbn/core/server';
import { NEVER } from 'rxjs';
import { mockActionResponse } from '../../__mocks__/action_result_data';
import { ElasticAssistantRequestHandlerContext } from '../../types';
import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
import { coreMock } from '@kbn/core/server/mocks';
import {
INVOKE_ASSISTANT_ERROR_EVENT,
INVOKE_ASSISTANT_SUCCESS_EVENT,
} from '../../lib/telemetry/event_based_telemetry';
import { INVOKE_ASSISTANT_ERROR_EVENT } from '../../lib/telemetry/event_based_telemetry';
import { PassThrough } from 'stream';
import { getConversationResponseMock } from '../../ai_assistant_data_clients/conversations/update_conversation.test';
import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock';
import { getFindAnonymizationFieldsResultWithSingleHit } from '../../__mocks__/response';
import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
import { chatCompleteRoute } from './chat_complete_route';
import { PublicMethodsOf } from '@kbn/utility-types';
import { licensingMock } from '@kbn/licensing-plugin/server/mocks';
import {
appendAssistantMessageToConversation,
createOrUpdateConversationWithUserInput,
langChainExecute,
} from '../helpers';
const license = licensingMock.createLicenseMock();
@ -33,43 +32,12 @@ const actionsClient = actionsClientMock.create();
jest.mock('../../lib/build_response', () => ({
buildResponse: jest.fn().mockImplementation((x) => x),
}));
jest.mock('../helpers');
const mockAppendAssistantMessageToConversation = appendAssistantMessageToConversation as jest.Mock;
const mockLangChainExecute = langChainExecute as jest.Mock;
const mockStream = jest.fn().mockImplementation(() => new PassThrough());
jest.mock('../../lib/langchain/execute_custom_llm_chain', () => ({
callAgentExecutor: jest.fn().mockImplementation(
async ({
connectorId,
isStream,
onLlmResponse,
}: {
onLlmResponse: (
content: string,
replacements: Record<string, string>,
isError: boolean
) => Promise<void>;
actionsClient: PublicMethodsOf<ActionsClient>;
connectorId: string;
esClient: ElasticsearchClient;
langChainMessages: BaseMessage[];
logger: Logger;
isStream: boolean;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
request: KibanaRequest<unknown, unknown, any, any>;
}) => {
if (!isStream && connectorId === 'mock-connector-id') {
return {
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
};
} else if (isStream && connectorId === 'mock-connector-id') {
return mockStream;
} else {
onLlmResponse('simulated error', {}, true).catch(() => {});
throw new Error('simulated error');
}
}
),
}));
const existingConversation = getConversationResponseMock();
const reportEvent = jest.fn();
const appendConversationMessages = jest.fn();
@ -103,6 +71,13 @@ const mockContext = {
appendConversationMessages:
appendConversationMessages.mockResolvedValue(existingConversation),
}),
getAIAssistantKnowledgeBaseDataClient: jest.fn().mockResolvedValue({
getKnowledgeBaseDocuments: jest.fn().mockResolvedValue([]),
indexTemplateAndPattern: {
alias: 'knowledge-base-alias',
},
isModelDeployed: jest.fn().mockResolvedValue(true),
}),
getAIAssistantAnonymizationFieldsDataClient: jest.fn().mockResolvedValue({
findDocuments: jest.fn().mockResolvedValue(getFindAnonymizationFieldsResultWithSingleHit()),
}),
@ -125,8 +100,6 @@ const mockRequest = {
conversationId: 'mock-conversation-id',
connectorId: 'mock-connector-id',
persist: true,
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
model: 'gpt-4',
messages: [
{
@ -164,7 +137,37 @@ describe('chatCompleteRoute', () => {
beforeEach(() => {
jest.clearAllMocks();
mockAppendAssistantMessageToConversation.mockResolvedValue(true);
license.hasAtLeast.mockReturnValue(true);
(createOrUpdateConversationWithUserInput as jest.Mock).mockResolvedValue({ id: 'something' });
mockLangChainExecute.mockImplementation(
async ({
connectorId,
isStream,
onLlmResponse,
}: {
connectorId: string;
isStream: boolean;
onLlmResponse: (
content: string,
replacements: Record<string, string>,
isError: boolean
) => Promise<void>;
}) => {
if (!isStream && connectorId === 'mock-connector-id') {
return {
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
};
} else if (isStream && connectorId === 'mock-connector-id') {
return mockStream;
} else {
onLlmResponse('simulated error', {}, true).catch(() => {});
throw new Error('simulated error');
}
}
);
actionsClient.execute.mockImplementation(
jest.fn().mockResolvedValue(() => ({
data: 'mockChatCompletion',
@ -256,72 +259,6 @@ describe('chatCompleteRoute', () => {
);
});
it('reports success events to telemetry - kb on, RAG alerts off', async () => {
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, mockRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await chatCompleteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports success events to telemetry - kb on, RAG alerts on', async () => {
const ragRequest = {
...mockRequest,
body: {
...mockRequest.body,
isEnabledRAGAlerts: true,
anonymizationFields: [
{ id: '@timestamp', field: '@timestamp', allowed: true, anonymized: false },
{ id: 'host.name', field: 'host.name', allowed: true, anonymized: true },
],
},
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, ragRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await chatCompleteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports error events to telemetry - kb on, RAG alerts off', async () => {
const requestWithBadConnectorId = {
...mockRequest,
@ -340,8 +277,6 @@ describe('chatCompleteRoute', () => {
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
errorMessage: 'simulated error',
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
@ -363,7 +298,7 @@ describe('chatCompleteRoute', () => {
...mockRequest,
body: {
...mockRequest.body,
conversationId: '99999',
conversationId: undefined,
connectorId: 'bad-connector-id',
},
};
@ -374,11 +309,10 @@ describe('chatCompleteRoute', () => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, badRequest, mockResponse);
expect(appendConversationMessages.mock.calls[1][0].messages[0]).toEqual(
expect(mockAppendAssistantMessageToConversation).toHaveBeenCalledWith(
expect.objectContaining({
content: 'simulated error',
messageContent: 'simulated error',
isError: true,
role: 'assistant',
})
);
}),

22
x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts
View file

@ -23,10 +23,8 @@ import { INVOKE_ASSISTANT_ERROR_EVENT } from '../../lib/telemetry/event_based_te
import { ElasticAssistantPluginRouter, GetElser } from '../../types';
import { buildResponse } from '../../lib/build_response';
import {
DEFAULT_PLUGIN_NAME,
appendAssistantMessageToConversation,
createOrUpdateConversationWithUserInput,
getPluginNameFromRequest,
langChainExecute,
performChecks,
} from '../helpers';
@ -151,20 +149,9 @@ export const chatCompleteRoute = (
});
let updatedConversation: ConversationResponse | undefined | null;
// Fetch any tools registered by the request's originating plugin
const pluginName = getPluginNameFromRequest({
request,
defaultPluginName: DEFAULT_PLUGIN_NAME,
logger,
});
const enableKnowledgeBaseByDefault =
ctx.elasticAssistant.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
// TODO: remove non-graph persistance when KB will be enabled by default
if (
(!enableKnowledgeBaseByDefault || (enableKnowledgeBaseByDefault && !conversationId)) &&
request.body.persist &&
conversationsDataClient
) {
// TODO: Remove non-graph persistence now that KB is enabled by default
if (!conversationId && request.body.persist && conversationsDataClient) {
updatedConversation = await createOrUpdateConversationWithUserInput({
actionsClient,
actionTypeId,
@ -209,7 +196,6 @@ export const chatCompleteRoute = (
return await langChainExecute({
abortSignal,
isEnabledKnowledgeBase: true,
isStream: request.body.isStream ?? false,
actionsClient,
actionTypeId,
@ -231,8 +217,6 @@ export const chatCompleteRoute = (
const error = transformError(err as Error);
telemetry?.reportEvent(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
actionTypeId: actionTypeId ?? '',
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
model: request.body.model,
errorMessage: error.message,
// TODO rm actionTypeId check when llmClass for bedrock streaming is implemented

5
x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
View file

@ -153,8 +153,6 @@ export const postEvaluateRoute = (
actionTypeId: '.gen-ai',
replacements: {},
size: DEFAULT_SIZE,
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
conversationId: '',
},
};
@ -164,7 +162,7 @@ export const postEvaluateRoute = (
const enableKnowledgeBaseByDefault =
assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
const kbDataClient = enableKnowledgeBaseByDefault
? (await assistantContext.getAIAssistantKnowledgeBaseDataClient(false)) ?? undefined
? (await assistantContext.getAIAssistantKnowledgeBaseDataClient()) ?? undefined
: undefined;
const kbIndex =
enableKnowledgeBaseByDefault && kbDataClient != null
@ -194,7 +192,6 @@ export const postEvaluateRoute = (
agentEvaluator: async (langChainMessages, exampleId) => {
const evalResult = await AGENT_EXECUTOR_MAP[agentName]({
actionsClient,
isEnabledKnowledgeBase: true,
assistantTools,
connectorId,
esClient,

82
x-pack/plugins/elastic_assistant/server/routes/helpers.ts
View file

@ -29,8 +29,7 @@ import { ActionsClient } from '@kbn/actions-plugin/server';
import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities';
import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
import { MINIMUM_AI_ASSISTANT_LICENSE } from '../../common/constants';
import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './knowledge_base/constants';
import { callAgentExecutor } from '../lib/langchain/execute_custom_llm_chain';
import { ESQL_RESOURCE } from './knowledge_base/constants';
import { buildResponse, getLlmType } from './utils';
import {
AgentExecutorParams,
@ -275,57 +274,10 @@ export interface NonLangChainExecuteParams {
response: KibanaResponseFactory;
telemetry: AnalyticsServiceSetup;
}
export const nonLangChainExecute = async ({
messages,
abortSignal,
actionTypeId,
connectorId,
logger,
actionsClient,
onLlmResponse,
response,
request,
telemetry,
}: NonLangChainExecuteParams) => {
logger.debug('Executing via actions framework directly');
const result = await executeAction({
abortSignal,
onLlmResponse,
actionsClient,
connectorId,
actionTypeId,
params: {
subAction: request.body.subAction,
subActionParams: {
model: request.body.model,
messages,
...(actionTypeId === '.gen-ai'
? { n: 1, stop: null, temperature: 0.2 }
: { temperature: 0, stopSequences: [] }),
},
},
logger,
});
telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
actionTypeId,
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase,
isEnabledRAGAlerts: request.body.isEnabledRAGAlerts,
model: request.body.model,
assistantStreamingEnabled: request.body.subAction !== 'invokeAI',
});
return response.ok({
body: result,
...(request.body.subAction === 'invokeAI'
? { headers: { 'content-type': 'application/json' } }
: {}),
});
};
export interface LangChainExecuteParams {
messages: Array<Pick<Message, 'content' | 'role'>>;
replacements: Replacements;
isEnabledKnowledgeBase: boolean;
isStream?: boolean;
onNewReplacements: (newReplacements: Replacements) => void;
abortSignal: AbortSignal;
@ -353,7 +305,6 @@ export const langChainExecute = async ({
messages,
replacements,
onNewReplacements,
isEnabledKnowledgeBase,
abortSignal,
telemetry,
actionTypeId,
@ -369,10 +320,6 @@ export const langChainExecute = async ({
responseLanguage,
isStream = true,
}: LangChainExecuteParams) => {
// TODO: Add `traceId` to actions request when calling via langchain
logger.debug(
`Executing via langchain, isEnabledKnowledgeBase: ${isEnabledKnowledgeBase}, isEnabledRAGAlerts: ${request.body.isEnabledRAGAlerts}`
);
// Fetch any tools registered by the request's originating plugin
const pluginName = getPluginNameFromRequest({
request,
@ -397,19 +344,11 @@ export const langChainExecute = async ({
const conversationsDataClient = await assistantContext.getAIAssistantConversationsDataClient();
// Create an ElasticsearchStore for KB interactions
// Setup with kbDataClient if `assistantKnowledgeBaseByDefault` FF is enabled
const enableKnowledgeBaseByDefault =
assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
const kbDataClient = enableKnowledgeBaseByDefault
? (await assistantContext.getAIAssistantKnowledgeBaseDataClient(false)) ?? undefined
: undefined;
const kbIndex =
enableKnowledgeBaseByDefault && kbDataClient != null
? kbDataClient.indexTemplateAndPattern.alias
: KNOWLEDGE_BASE_INDEX_PATTERN;
const kbDataClient =
(await assistantContext.getAIAssistantKnowledgeBaseDataClient()) ?? undefined;
const esStore = new ElasticsearchStore(
esClient,
kbIndex,
kbDataClient?.indexTemplateAndPattern?.alias ?? '',
logger,
telemetry,
elserId,
@ -429,7 +368,6 @@ export const langChainExecute = async ({
dataClients,
alertsIndexPattern: request.body.alertsIndexPattern,
actionsClient,
isEnabledKnowledgeBase,
assistantTools,
conversationId,
connectorId,
@ -455,18 +393,12 @@ export const langChainExecute = async ({
},
};
// New code path for LangGraph implementation, behind `assistantKnowledgeBaseByDefault` FF
let result: StreamResponseWithHeaders | StaticReturnType;
if (enableKnowledgeBaseByDefault && request.body.isEnabledKnowledgeBase) {
result = await callAssistantGraph(executorParams);
} else {
result = await callAgentExecutor(executorParams);
}
const result: StreamResponseWithHeaders | StaticReturnType = await callAssistantGraph(
executorParams
);
telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
actionTypeId,
isEnabledKnowledgeBase,
isEnabledRAGAlerts: request.body.isEnabledRAGAlerts ?? true,
model: request.body.model,
// TODO rm actionTypeId check when llmClass for bedrock streaming is implemented
// tracked here: https://github.com/elastic/security-team/issues/7363

50
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.test.ts
View file

@ -1,50 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { deleteKnowledgeBaseRoute } from './delete_knowledge_base';
import { serverMock } from '../../__mocks__/server';
import { requestContextMock } from '../../__mocks__/request_context';
import { getDeleteKnowledgeBaseRequest } from '../../__mocks__/request';
import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
describe('Delete Knowledge Base Route', () => {
let server: ReturnType<typeof serverMock.create>;
// eslint-disable-next-line prefer-const
let { clients, context } = requestContextMock.createTools();
clients.core.elasticsearch.client = elasticsearchServiceMock.createScopedClusterClient();
beforeEach(() => {
server = serverMock.create();
({ context } = requestContextMock.createTools());
deleteKnowledgeBaseRoute(server.router);
});
describe('Status codes', () => {
test('returns 200 if base resources are deleted', async () => {
const response = await server.inject(
getDeleteKnowledgeBaseRequest(),
requestContextMock.convertContext(context)
);
expect(response.status).toEqual(200);
});
test('returns 500 if error is thrown when deleting resources', async () => {
context.core.elasticsearch.client.asInternalUser.indices.delete.mockRejectedValue(
new Error('Test error')
);
const response = await server.inject(
getDeleteKnowledgeBaseRequest('esql'),
requestContextMock.convertContext(context)
);
expect(response.status).toEqual(500);
});
});
});

46
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts
View file

@ -20,8 +20,7 @@ import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/
import { buildResponse } from '../../lib/build_response';
import { ElasticAssistantRequestHandlerContext } from '../../types';
import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store';
import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants';
import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
import { ESQL_RESOURCE } from './constants';
import { getKbResource } from './get_kb_resource';
/**
@ -53,42 +52,25 @@ export const deleteKnowledgeBaseRoute = (
const assistantContext = await context.elasticAssistant;
const logger = assistantContext.logger;
const telemetry = assistantContext.telemetry;
const pluginName = getPluginNameFromRequest({
request,
defaultPluginName: DEFAULT_PLUGIN_NAME,
logger,
});
const enableKnowledgeBaseByDefault =
assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
try {
const kbResource = getKbResource(request);
const esClient = (await context.core).elasticsearch.client.asInternalUser;
let esStore = new ElasticsearchStore(
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
logger,
telemetry
);
// Code path for when `assistantKnowledgeBaseByDefault` FF is enabled, only need an esStore w/ kbDataClient
if (enableKnowledgeBaseByDefault) {
const knowledgeBaseDataClient =
await assistantContext.getAIAssistantKnowledgeBaseDataClient(false);
if (!knowledgeBaseDataClient) {
return response.custom({ body: { success: false }, statusCode: 500 });
}
esStore = new ElasticsearchStore(
esClient,
knowledgeBaseDataClient.indexTemplateAndPattern.alias,
logger,
telemetry,
'elserId', // Not needed for delete ops
kbResource,
knowledgeBaseDataClient
);
const knowledgeBaseDataClient =
await assistantContext.getAIAssistantKnowledgeBaseDataClient();
if (!knowledgeBaseDataClient) {
return response.custom({ body: { success: false }, statusCode: 500 });
}
const esStore = new ElasticsearchStore(
esClient,
knowledgeBaseDataClient.indexTemplateAndPattern.alias,
logger,
telemetry,
'elserId', // Not needed for delete ops
kbResource,
knowledgeBaseDataClient
);
if (kbResource === ESQL_RESOURCE) {
// For now, tearing down the Knowledge Base is fine, but will want to support removing specific assets based

4
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts
View file

@ -166,9 +166,7 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug
// subscribing to completed$, because it handles both cases when request was completed and aborted.
// when route is finished by timeout, aborted$ is not getting fired
request.events.completed$.subscribe(() => abortController.abort());
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(
false
);
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient();
const spaceId = ctx.elasticAssistant.getSpaceId();
// Authenticated user null check completed in `performChecks()` above
const authenticatedUser = ctx.elasticAssistant.getCurrentUser() as AuthenticatedUser;

4
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts
View file

@ -67,9 +67,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
pageContent: request.body.text,
},
];
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(
false
);
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient();
const createResponse = await kbDataClient?.addKnowledgeBaseDocuments({ documents });
if (createResponse == null) {

4
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts
View file

@ -63,9 +63,7 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout
return checkResponse;
}
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(
false
);
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient();
const currentUser = ctx.elasticAssistant.getCurrentUser();
const additionalFilter = query.filter ? ` AND ${query.filter}` : '';

7
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts
View file

@ -32,6 +32,13 @@ describe('Get Knowledge Base Status Route', () => {
server = serverMock.create();
({ context } = requestContextMock.createTools());
context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser);
context.elasticAssistant.getAIAssistantKnowledgeBaseDataClient = jest.fn().mockResolvedValue({
getKnowledgeBaseDocuments: jest.fn().mockResolvedValue([]),
indexTemplateAndPattern: {
alias: 'knowledge-base-alias',
},
isModelInstalled: jest.fn().mockResolvedValue(true),
});
getKnowledgeBaseStatusRoute(server.router, mockGetElser);
});

51
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts
View file

@ -19,8 +19,7 @@ import { getKbResource } from './get_kb_resource';
import { buildResponse } from '../../lib/build_response';
import { ElasticAssistantPluginRouter, GetElser } from '../../types';
import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store';
import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants';
import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE } from './constants';
/**
* Get the status of the Knowledge Base index, pipeline, and resources (collection of documents)
@ -56,49 +55,27 @@ export const getKnowledgeBaseStatusRoute = (
const telemetry = assistantContext.telemetry;
try {
// Use asInternalUser
const esClient = (await context.core).elasticsearch.client.asInternalUser;
const elserId = await getElser();
const kbResource = getKbResource(request);
let esStore = new ElasticsearchStore(
const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient();
if (!kbDataClient) {
return response.custom({ body: { success: false }, statusCode: 500 });
}
// Use old status checks by overriding esStore to use kbDataClient
const esStore = new ElasticsearchStore(
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
kbDataClient.indexTemplateAndPattern.alias,
logger,
telemetry,
elserId,
kbResource
kbResource,
kbDataClient
);
const pluginName = getPluginNameFromRequest({
request,
defaultPluginName: DEFAULT_PLUGIN_NAME,
logger,
});
const enableKnowledgeBaseByDefault =
assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
// Code path for when `assistantKnowledgeBaseByDefault` FF is enabled
let isSetupInProgress = false;
if (enableKnowledgeBaseByDefault) {
const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient(
false
);
if (!kbDataClient) {
return response.custom({ body: { success: false }, statusCode: 500 });
}
// Use old status checks by overriding esStore to use kbDataClient
esStore = new ElasticsearchStore(
esClient,
kbDataClient.indexTemplateAndPattern.alias,
logger,
telemetry,
elserId,
kbResource,
kbDataClient
);
isSetupInProgress = kbDataClient.isSetupInProgress;
}
const indexExists = await esStore.indexExists();
const pipelineExists = await esStore.pipelineExists();
const modelExists = await esStore.isModelInstalled(elserId);
@ -106,7 +83,7 @@ export const getKnowledgeBaseStatusRoute = (
const body: ReadKnowledgeBaseResponse = {
elser_exists: modelExists,
index_exists: indexExists,
is_setup_in_progress: isSetupInProgress,
is_setup_in_progress: kbDataClient.isSetupInProgress,
pipeline_exists: pipelineExists,
};

19
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts
View file

@ -32,6 +32,13 @@ describe('Post Knowledge Base Route', () => {
server = serverMock.create();
({ context } = requestContextMock.createTools());
context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser);
context.elasticAssistant.getAIAssistantKnowledgeBaseDataClient = jest.fn().mockResolvedValue({
setupKnowledgeBase: jest.fn(),
indexTemplateAndPattern: {
alias: 'knowledge-base-alias',
},
isModelInstalled: jest.fn().mockResolvedValue(true),
});
postKnowledgeBaseRoute(server.router, mockGetElser);
});
@ -45,17 +52,5 @@ describe('Post Knowledge Base Route', () => {
expect(response.status).toEqual(200);
});
test('returns 500 if error is thrown when creating resources', async () => {
context.core.elasticsearch.client.asInternalUser.indices.exists.mockRejectedValue(
new Error('Test error')
);
const response = await server.inject(
getPostKnowledgeBaseRequest('esql'),
requestContextMock.convertContext(context)
);
expect(response.status).toEqual(500);
});
});
});

75
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts
View file

@ -18,10 +18,7 @@ import { IKibanaResponse, KibanaRequest } from '@kbn/core/server';
import { buildResponse } from '../../lib/build_response';
import { ElasticAssistantPluginRouter, GetElser } from '../../types';
import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store';
import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants';
import { getKbResource } from './get_kb_resource';
import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader';
import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
// Since we're awaiting on ELSER setup, this could take a bit (especially if ML needs to autoscale)
// Consider just returning if attempt was successful, and switch to client polling
@ -70,79 +67,27 @@ export const postKnowledgeBaseRoute = (
const esClient = core.elasticsearch.client.asInternalUser;
const soClient = core.savedObjects.getClient();
const pluginName = getPluginNameFromRequest({
request,
defaultPluginName: DEFAULT_PLUGIN_NAME,
logger,
});
const enableKnowledgeBaseByDefault =
assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
try {
// Code path for when `assistantKnowledgeBaseByDefault` FF is enabled
if (enableKnowledgeBaseByDefault) {
const knowledgeBaseDataClient =
await assistantContext.getAIAssistantKnowledgeBaseDataClient(true);
if (!knowledgeBaseDataClient) {
return response.custom({ body: { success: false }, statusCode: 500 });
}
// Continue to use esStore for loading esql docs until `semantic_text` is available and we can test the new chunking strategy
const esStore = new ElasticsearchStore(
esClient,
knowledgeBaseDataClient.indexTemplateAndPattern.alias,
logger,
telemetry,
elserId,
getKbResource(request),
knowledgeBaseDataClient
);
await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, soClient });
return response.ok({ body: { success: true } });
const knowledgeBaseDataClient =
await assistantContext.getAIAssistantKnowledgeBaseDataClient();
if (!knowledgeBaseDataClient) {
return response.custom({ body: { success: false }, statusCode: 500 });
}
const kbResource = getKbResource(request);
// Continue to use esStore for loading esql docs until `semantic_text` is available and we can test the new chunking strategy
const esStore = new ElasticsearchStore(
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
knowledgeBaseDataClient.indexTemplateAndPattern.alias,
logger,
telemetry,
elserId,
kbResource
getKbResource(request),
knowledgeBaseDataClient
);
// Pre-check on index/pipeline
let indexExists = await esStore.indexExists();
let pipelineExists = await esStore.pipelineExists();
await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, soClient });
// Load if not exists
if (!pipelineExists) {
pipelineExists = await esStore.createPipeline();
}
if (!indexExists) {
indexExists = await esStore.createIndex();
}
// If specific resource is requested, load it
if (kbResource === ESQL_RESOURCE) {
const esqlExists = (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0;
if (!esqlExists) {
const loadedKnowledgeBase = await loadESQL(esStore, logger);
return response.custom({ body: { success: loadedKnowledgeBase }, statusCode: 201 });
} else {
return response.ok({ body: { success: true } });
}
}
const wasSuccessful = indexExists && pipelineExists;
if (wasSuccessful) {
return response.ok({ body: { success: true } });
} else {
return response.custom({ body: { success: false }, statusCode: 500 });
}
return response.ok({ body: { success: true } });
} catch (err) {
logger.log(err);
const error = transformError(err);

412
x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts
View file

@ -5,9 +5,7 @@
* 2.0.
*/
import { ElasticsearchClient, IRouter, KibanaRequest, Logger } from '@kbn/core/server';
import type { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server';
import { BaseMessage } from '@langchain/core/messages';
import { IRouter, KibanaRequest } from '@kbn/core/server';
import { NEVER } from 'rxjs';
import { mockActionResponse } from '../__mocks__/action_result_data';
import { postActionsConnectorExecuteRoute } from './post_actions_connector_execute';
@ -15,16 +13,17 @@ import { ElasticAssistantRequestHandlerContext } from '../types';
import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
import { coreMock } from '@kbn/core/server/mocks';
import {
INVOKE_ASSISTANT_ERROR_EVENT,
INVOKE_ASSISTANT_SUCCESS_EVENT,
} from '../lib/telemetry/event_based_telemetry';
import { INVOKE_ASSISTANT_ERROR_EVENT } from '../lib/telemetry/event_based_telemetry';
import { PassThrough } from 'stream';
import { getConversationResponseMock } from '../ai_assistant_data_clients/conversations/update_conversation.test';
import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock';
import { getFindAnonymizationFieldsResultWithSingleHit } from '../__mocks__/response';
import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
import {
defaultAssistantFeatures,
ExecuteConnectorRequestBody,
} from '@kbn/elastic-assistant-common';
import { licensingMock } from '@kbn/licensing-plugin/server/mocks';
import { appendAssistantMessageToConversation, langChainExecute } from './helpers';
const license = licensingMock.createLicenseMock();
const actionsClient = actionsClientMock.create();
@ -45,47 +44,9 @@ jest.mock('../lib/executor', () => ({
}),
}));
const mockStream = jest.fn().mockImplementation(() => new PassThrough());
jest.mock('../lib/langchain/execute_custom_llm_chain', () => ({
callAgentExecutor: jest.fn().mockImplementation(
async ({
connectorId,
isStream,
}: {
actions: ActionsPluginStart;
connectorId: string;
esClient: ElasticsearchClient;
langChainMessages: BaseMessage[];
logger: Logger;
isStream: boolean;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
request: KibanaRequest<unknown, unknown, any, any>;
}) => {
if (!isStream && connectorId === 'mock-connector-id') {
return {
body: {
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
},
headers: { 'content-type': 'application/json' },
};
} else if (isStream && connectorId === 'mock-connector-id') {
return {
body: mockStream,
headers: {
'Cache-Control': 'no-cache',
Connection: 'keep-alive',
'Transfer-Encoding': 'chunked',
'X-Accel-Buffering': 'no',
'X-Content-Type-Options': 'nosniff',
},
};
} else {
throw new Error('simulated error');
}
}
),
}));
const mockLangChainExecute = langChainExecute as jest.Mock;
const mockAppendAssistantMessageToConversation = appendAssistantMessageToConversation as jest.Mock;
jest.mock('./helpers');
const existingConversation = getConversationResponseMock();
const reportEvent = jest.fn();
const appendConversationMessages = jest.fn();
@ -121,6 +82,12 @@ const mockContext = {
getAIAssistantAnonymizationFieldsDataClient: jest.fn().mockResolvedValue({
findDocuments: jest.fn().mockResolvedValue(getFindAnonymizationFieldsResultWithSingleHit()),
}),
getAIAssistantKnowledgeBaseDataClient: jest.fn().mockResolvedValue({
getKnowledgeBaseDocuments: jest.fn().mockResolvedValue([]),
indexTemplateAndPattern: {
alias: 'knowledge-base-alias',
},
}),
},
core: {
elasticsearch: {
@ -141,8 +108,6 @@ const mockRequest = {
subAction: 'invokeAI',
message: 'Do you know my name?',
actionTypeId: '.gen-ai',
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
replacements: {},
model: 'gpt-4',
},
@ -162,6 +127,40 @@ describe('postActionsConnectorExecuteRoute', () => {
beforeEach(() => {
jest.clearAllMocks();
license.hasAtLeast.mockReturnValue(true);
mockAppendAssistantMessageToConversation.mockResolvedValue(true);
mockLangChainExecute.mockImplementation(
async ({
connectorId,
request,
}: {
connectorId: string;
request: KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody>;
}) => {
if (request.body.subAction === 'invokeAI' && connectorId === 'mock-connector-id') {
return {
body: {
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
},
headers: { 'content-type': 'application/json' },
};
} else if (request.body.subAction !== 'invokeAI' && connectorId === 'mock-connector-id') {
return {
body: mockStream,
headers: {
'Cache-Control': 'no-cache',
Connection: 'keep-alive',
'Transfer-Encoding': 'chunked',
'X-Accel-Buffering': 'no',
'X-Content-Type-Options': 'nosniff',
},
};
} else {
throw new Error('simulated error');
}
}
);
actionsClient.getBulk.mockResolvedValue([
{
id: '1',
@ -180,45 +179,7 @@ describe('postActionsConnectorExecuteRoute', () => {
]);
});
it('returns the expected response when isEnabledKnowledgeBase=false', async () => {
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
const result = await handler(
mockContext,
{
...mockRequest,
body: {
...mockRequest.body,
isEnabledKnowledgeBase: false,
},
},
mockResponse
);
expect(result).toEqual({
body: {
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
},
headers: { 'content-type': 'application/json' },
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('returns the expected response when isEnabledKnowledgeBase=true', async () => {
it('returns the expected response', async () => {
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
@ -275,189 +236,17 @@ describe('postActionsConnectorExecuteRoute', () => {
);
});
it('reports success events to telemetry - kb on, RAG alerts off', async () => {
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, mockRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports success events to telemetry - kb on, RAG alerts on', async () => {
const ragRequest = {
...mockRequest,
body: {
...mockRequest.body,
anonymizationFields: [
{ id: '@timestamp', field: '@timestamp', allowed: true, anonymized: false },
{ id: 'host.name', field: 'host.name', allowed: true, anonymized: true },
],
replacements: [],
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, ragRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports success events to telemetry - kb off, RAG alerts on', async () => {
const req = {
...mockRequest,
body: {
...mockRequest.body,
isEnabledKnowledgeBase: false,
anonymizationFields: [
{ id: '@timestamp', field: '@timestamp', allowed: true, anonymized: false },
{ id: 'host.name', field: 'host.name', allowed: true, anonymized: true },
],
replacements: [],
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, req, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: true,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports success events to telemetry - kb off, RAG alerts off', async () => {
const req = {
...mockRequest,
body: {
...mockRequest.body,
isEnabledKnowledgeBase: false,
},
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, req, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports error events to telemetry - kb on, RAG alerts off', async () => {
const requestWithBadConnectorId = {
...mockRequest,
params: { connectorId: 'bad-connector-id' },
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, requestWithBadConnectorId, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
errorMessage: 'simulated error',
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports error events to telemetry - kb on, RAG alerts on', async () => {
it('reports error events to telemetry', async () => {
const badRequest = {
...mockRequest,
params: { connectorId: 'bad-connector-id' },
body: {
...mockRequest.body,
isEnabledRAGAlerts: true,
anonymizationFields: [
{ id: '@timestamp', field: '@timestamp', allowed: true, anonymized: false },
{ id: 'host.name', field: 'host.name', allowed: true, anonymized: true },
],
replacements: [],
},
};
@ -470,51 +259,6 @@ describe('postActionsConnectorExecuteRoute', () => {
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
errorMessage: 'simulated error',
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports error events to telemetry - kb off, RAG alerts on', async () => {
const badRequest = {
...mockRequest,
params: { connectorId: 'bad-connector-id' },
body: {
...mockRequest.body,
isEnabledKnowledgeBase: false,
anonymizationFields: [
{ id: '@timestamp', field: '@timestamp', allowed: true, anonymized: false },
{ id: 'host.name', field: 'host.name', allowed: true, anonymized: true },
],
replacements: [],
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, badRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
errorMessage: 'simulated error',
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: true,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
@ -547,11 +291,10 @@ describe('postActionsConnectorExecuteRoute', () => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, badRequest, mockResponse);
expect(appendConversationMessages.mock.calls[1][0].messages[0]).toEqual(
expect(mockAppendAssistantMessageToConversation).toHaveBeenCalledWith(
expect.objectContaining({
content: 'simulated error',
messageContent: 'simulated error',
isError: true,
role: 'assistant',
})
);
}),
@ -566,43 +309,6 @@ describe('postActionsConnectorExecuteRoute', () => {
);
});
it('reports error events to telemetry - kb off, RAG alerts off', async () => {
const badRequest = {
...mockRequest,
params: { connectorId: 'bad-connector-id' },
body: {
...mockRequest.body,
isEnabledKnowledgeBase: false,
},
};
const mockRouter = {
versioned: {
post: jest.fn().mockImplementation(() => {
return {
addVersion: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, badRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
errorMessage: 'simulated error',
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
actionTypeId: '.gen-ai',
model: 'gpt-4',
assistantStreamingEnabled: false,
});
}),
};
}),
},
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('returns the expected response when subAction=invokeStream and actionTypeId=.gen-ai', async () => {
const mockRouter = {
versioned: {

65
x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts
View file

@ -21,14 +21,7 @@ import { INVOKE_ASSISTANT_ERROR_EVENT } from '../lib/telemetry/event_based_telem
import { POST_ACTIONS_CONNECTOR_EXECUTE } from '../../common/constants';
import { buildResponse } from '../lib/build_response';
import { ElasticAssistantRequestHandlerContext, GetElser } from '../types';
import {
DEFAULT_PLUGIN_NAME,
appendAssistantMessageToConversation,
getPluginNameFromRequest,
langChainExecute,
nonLangChainExecute,
updateConversationWithUserInput,
} from './helpers';
import { appendAssistantMessageToConversation, langChainExecute } from './helpers';
export const postActionsConnectorExecuteRoute = (
router: IRouter<ElasticAssistantRequestHandlerContext>,
@ -97,41 +90,6 @@ export const postActionsConnectorExecuteRoute = (
const conversationsDataClient =
await assistantContext.getAIAssistantConversationsDataClient();
// Fetch any tools registered by the request's originating plugin
const pluginName = getPluginNameFromRequest({
request,
defaultPluginName: DEFAULT_PLUGIN_NAME,
logger,
});
const isGraphAvailable =
assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault &&
request.body.isEnabledKnowledgeBase;
// TODO: remove non-graph persistance when KB will be enabled by default
if (!isGraphAvailable && conversationId && conversationsDataClient) {
const updatedConversation = await updateConversationWithUserInput({
actionsClient,
actionTypeId,
connectorId,
conversationId,
conversationsDataClient,
logger,
replacements: latestReplacements,
newMessages: newMessage ? [newMessage] : [],
model: request.body.model,
});
if (updatedConversation == null) {
return response.badRequest({
body: `conversation id: "${conversationId}" not updated`,
});
}
// messages are anonymized by conversationsDataClient
messages = updatedConversation?.messages?.map((c) => ({
role: c.role,
content: c.content,
}));
}
onLlmResponse = async (
content: string,
traceData: Message['traceData'] = {},
@ -149,26 +107,9 @@ export const postActionsConnectorExecuteRoute = (
}
};
if (!request.body.isEnabledKnowledgeBase && !request.body.isEnabledRAGAlerts) {
// if not langchain, call execute action directly and return the response:
return await nonLangChainExecute({
abortSignal,
actionsClient,
actionTypeId,
connectorId,
logger,
messages: messages ?? [],
onLlmResponse,
request,
response,
telemetry,
});
}
return await langChainExecute({
abortSignal,
isStream: request.body.subAction !== 'invokeAI',
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase ?? false,
actionsClient,
actionTypeId,
connectorId,
@ -176,7 +117,7 @@ export const postActionsConnectorExecuteRoute = (
context: ctx,
getElser,
logger,
messages: (isGraphAvailable && newMessage ? [newMessage] : messages) ?? [],
messages: (newMessage ? [newMessage] : messages) ?? [],
onLlmResponse,
onNewReplacements,
replacements: latestReplacements,
@ -192,8 +133,6 @@ export const postActionsConnectorExecuteRoute = (
}
telemetry.reportEvent(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
actionTypeId: request.body.actionTypeId,
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase,
isEnabledRAGAlerts: request.body.isEnabledRAGAlerts,
model: request.body.model,
errorMessage: error.message,
assistantStreamingEnabled: request.body.subAction !== 'invokeAI',

5
x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts
View file

@ -81,15 +81,12 @@ export class RequestContextFactory implements IRequestContextFactory {
telemetry: core.analytics,
// Note: Due to plugin lifecycle and feature flag registration timing, we need to pass in the feature flag here
// Remove `initializeKnowledgeBase` once 'assistantKnowledgeBaseByDefault' feature flag is removed
getAIAssistantKnowledgeBaseDataClient: memoize((initializeKnowledgeBase = false) => {
getAIAssistantKnowledgeBaseDataClient: memoize(() => {
const currentUser = getCurrentUser();
return this.assistantService.createAIAssistantKnowledgeBaseDataClient({
spaceId: getSpaceId(),
logger: this.logger,
currentUser,
initializeKnowledgeBase,
});
}),

4
x-pack/plugins/elastic_assistant/server/types.ts
View file

@ -113,9 +113,7 @@ export interface ElasticAssistantApiRequestHandlerContext {
getSpaceId: () => string;
getCurrentUser: () => AuthenticatedUser | null;
getAIAssistantConversationsDataClient: () => Promise<AIAssistantConversationsDataClient | null>;
getAIAssistantKnowledgeBaseDataClient: (
initializeKnowledgeBase: boolean
) => Promise<AIAssistantKnowledgeBaseDataClient | null>;
getAIAssistantKnowledgeBaseDataClient: () => Promise<AIAssistantKnowledgeBaseDataClient | null>;
getAttackDiscoveryDataClient: () => Promise<AttackDiscoveryDataClient | null>;
getAIAssistantPromptsDataClient: () => Promise<AIAssistantDataClient | null>;
getAIAssistantAnonymizationFieldsDataClient: () => Promise<AIAssistantDataClient | null>;

2
x-pack/plugins/security_solution/common/experimental_features.ts
View file

@ -124,7 +124,7 @@ export const allowedExperimentalValues = Object.freeze({
assistantModelEvaluation: false,
/**
* Enables the Assistant Knowledge Base by default, introduced in `8.15.0`.
* Enables new Knowledge Base Entries features, introduced in `8.15.0`.
*/
assistantKnowledgeBaseByDefault: false,

82
x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts
View file

@ -12,14 +12,7 @@ import {
bedrockClaude2SuccessResponse,
} from '@kbn/actions-simulators-plugin/server/bedrock_simulation';
import { DEFAULT_TOKEN_LIMIT } from '@kbn/stack-connectors-plugin/common/bedrock/constants';
import { PassThrough } from 'stream';
import { EventStreamCodec } from '@smithy/eventstream-codec';
import { fromUtf8, toUtf8 } from '@smithy/util-utf8';
import { TaskErrorSource } from '@kbn/task-manager-plugin/common';
import {
ELASTIC_HTTP_VERSION_HEADER,
X_ELASTIC_INTERNAL_ORIGIN_REQUEST,
} from '@kbn/core-http-common';
import { FtrProviderContext } from '../../../../../common/ftr_provider_context';
import { getUrlPrefix, ObjectRemover } from '../../../../../common/lib';
@ -432,38 +425,6 @@ export default function bedrockTest({ getService }: FtrProviderContext) {
});
});
it('should invoke stream with assistant AI body argument formatted to bedrock expectations', async () => {
await new Promise<void>((resolve, reject) => {
const passThrough = new PassThrough();
supertest
.post(`/internal/elastic_assistant/actions/connector/${bedrockActionId}/_execute`)
.set('kbn-xsrf', 'foo')
.set(ELASTIC_HTTP_VERSION_HEADER, '1')
.set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana')
.on('error', reject)
.send({
actionTypeId: '.bedrock',
subAction: 'invokeStream',
message: 'Hello world',
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
replacements: {},
})
.pipe(passThrough);
const responseBuffer: Uint8Array[] = [];
passThrough.on('data', (chunk) => {
responseBuffer.push(chunk);
});
passThrough.on('end', () => {
const parsed = parseBedrockBuffer(responseBuffer);
expect(parsed).to.eql('Hello world, what a unique string!');
resolve();
});
});
});
describe('Token tracking dashboard', () => {
const dashboardId = 'specific-dashboard-id-default';
@ -647,46 +608,3 @@ export default function bedrockTest({ getService }: FtrProviderContext) {
});
});
}
const parseBedrockBuffer = (chunks: Uint8Array[]): string => {
let bedrockBuffer: Uint8Array = new Uint8Array(0);
return chunks
.map((chunk) => {
bedrockBuffer = concatChunks(bedrockBuffer, chunk);
let messageLength = getMessageLength(bedrockBuffer);
const buildChunks = [];
while (bedrockBuffer.byteLength > 0 && bedrockBuffer.byteLength >= messageLength) {
const extractedChunk = bedrockBuffer.slice(0, messageLength);
buildChunks.push(extractedChunk);
bedrockBuffer = bedrockBuffer.slice(messageLength);
messageLength = getMessageLength(bedrockBuffer);
}
const awsDecoder = new EventStreamCodec(toUtf8, fromUtf8);
return buildChunks
.map((bChunk) => {
const event = awsDecoder.decode(bChunk);
const body = JSON.parse(
Buffer.from(JSON.parse(new TextDecoder().decode(event.body)).bytes, 'base64').toString()
);
return body.delta.text;
})
.join('');
})
.join('');
};
function concatChunks(a: Uint8Array, b: Uint8Array): Uint8Array {
const newBuffer = new Uint8Array(a.length + b.length);
newBuffer.set(a);
newBuffer.set(b, a.length);
return newBuffer;
}
function getMessageLength(buffer: Uint8Array): number {
if (buffer.byteLength === 0) return 0;
const view = new DataView(buffer.buffer, buffer.byteOffset, buffer.byteLength);
return view.getUint32(0, false);
}

7
x-pack/test/security_solution_api_integration/package.json
View file

@ -45,13 +45,6 @@
"intialize-server:explore": "node scripts/index.js server explore trial_license_complete_tier",
"run-tests:explore": "node scripts/index.js runner explore trial_license_complete_tier",
"genai:server:serverless": "npm run initialize-server:genai:trial_complete invoke_ai serverless",
"genai:runner:serverless": "npm run run-tests:genai:trial_complete invoke_ai serverless serverlessEnv",
"genai:qa:serverless": "npm run run-tests:genai:trial_complete invoke_ai serverless qaPeriodicEnv",
"genai:qa:serverless:release": "npm run run-tests:genai:trial_complete invoke_ai serverless qaEnv",
"genai:server:ess": "npm run initialize-server:genai:trial_complete invoke_ai ess",
"genai:runner:ess": "npm run run-tests:genai:trial_complete invoke_ai ess essEnv",
"nlp_cleanup_task:complete:server:serverless": "npm run initialize-server:genai:trial_complete nlp_cleanup_task serverless",
"nlp_cleanup_task:complete:runner:serverless": "npm run run-tests:genai:trial_complete nlp_cleanup_task serverless serverlessEnv",
"nlp_cleanup_task:complete:qa:serverless": "npm run run-tests:genai:trial_complete nlp_cleanup_task serverless qaPeriodicEnv",

102
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/trial_license_complete_tier/basic.ts
View file

@ -1,102 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import expect from '@kbn/expect';
import { BedrockSimulator } from '@kbn/actions-simulators-plugin/server/bedrock_simulation';
import { OpenAISimulator } from '@kbn/actions-simulators-plugin/server/openai_simulation';
import { FtrProviderContext } from '../../../../ftr_provider_context';
import { postActionsClientExecute } from '../utils/post_actions_client_execute';
import { ObjectRemover } from '../utils/object_remover';
import { createConnector } from '../utils/create_connector';
const mockRequest = {
message: 'Do you know my name?',
subAction: 'invokeAI',
actionTypeId: '.bedrock',
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
replacements: {},
};
export default ({ getService }: FtrProviderContext) => {
const supertest = getService('supertest');
const objectRemover = new ObjectRemover(supertest);
const configService = getService('config');
// @skipInServerlessMKI tag because the simulators do not work in the QA env
describe('@ess @serverless @skipInServerlessMKI Basic Security AI Assistant Invoke AI [non-streaming, non-LangChain]', async () => {
after(() => {
objectRemover.removeAll();
});
describe('With Bedrock connector', () => {
const simulator = new BedrockSimulator({
proxy: {
config: configService.get('kbnTestServer.serverArgs'),
},
});
let apiUrl: string;
let bedrockActionId: string;
before(async () => {
apiUrl = await simulator.start();
bedrockActionId = await createConnector(supertest, objectRemover, apiUrl, 'bedrock');
});
after(() => {
simulator.close();
});
it('should execute a chat completion', async () => {
const response = await postActionsClientExecute(bedrockActionId, mockRequest, supertest);
const expected = {
connector_id: bedrockActionId,
data: 'Hello there! How may I assist you today?',
status: 'ok',
};
expect(response.body).to.eql(expected);
});
});
describe('With OpenAI connector', () => {
const simulator = new OpenAISimulator({
returnError: false,
proxy: {
config: configService.get('kbnTestServer.serverArgs'),
},
});
let apiUrl: string;
let openaiActionId: string;
before(async () => {
apiUrl = await simulator.start();
openaiActionId = await createConnector(supertest, objectRemover, apiUrl, 'openai');
});
after(() => {
simulator.close();
});
it('should execute a chat completion', async () => {
const response = await postActionsClientExecute(
openaiActionId,
{ ...mockRequest, actionTypeId: '.gen-ai' },
supertest
);
const expected = {
connector_id: openaiActionId,
data: 'Hello there! How may I assist you today?',
status: 'ok',
};
expect(response.body).to.eql(expected);
});
});
});
};

34
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/trial_license_complete_tier/configs/ess.config.ts
View file

@ -1,34 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { FtrConfigProviderContext } from '@kbn/test';
import getPort from 'get-port';
export default async function ({ readConfigFile }: FtrConfigProviderContext) {
const functionalConfig = await readConfigFile(
require.resolve('../../../../../config/ess/config.base.trial')
);
const proxyPort = await getPort({ port: getPort.makeRange(6200, 6299) });
return {
...functionalConfig.getAll(),
kbnTestServer: {
...functionalConfig.get('kbnTestServer'),
serverArgs: [
...functionalConfig.get('kbnTestServer.serverArgs'),
// used for connector simulators
`--xpack.actions.proxyUrl=http://localhost:${proxyPort}`,
`--xpack.actions.enabledActionTypes=${JSON.stringify(['.bedrock', '.gen-ai', '.gemini'])}`,
],
},
testFiles: [require.resolve('..')],
junit: {
reportName: 'GenAI - Invoke AI Tests - ESS Env - Trial License',
},
};
}

19
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/trial_license_complete_tier/configs/serverless.config.ts
View file

@ -1,19 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { createTestConfig } from '../../../../../config/serverless/config.base';
export default createTestConfig({
kbnTestServerArgs: [
// used for connector simulators
`--xpack.actions.proxyUrl=http://localhost:6200`,
],
testFiles: [require.resolve('..')],
junit: {
reportName: 'GenAI - Invoke AI Tests - Serverless Env - Complete Tier',
},
});

15
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/trial_license_complete_tier/index.ts
View file

@ -1,15 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { FtrProviderContext } from '../../../../ftr_provider_context';
export default function ({ loadTestFile }: FtrProviderContext) {
// this is the test suite for the inaptly named post_actions_connector_execute route
describe('GenAI - Invoke AI', function () {
loadTestFile(require.resolve('./basic'));
});
}

73
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/utils/create_connector.ts
View file

@ -1,73 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import type SuperTest from 'supertest';
import {
ELASTIC_HTTP_VERSION_HEADER,
X_ELASTIC_INTERNAL_ORIGIN_REQUEST,
} from '@kbn/core-http-common';
import { getUrlPrefix } from './space_test_utils';
import { ObjectRemover } from './object_remover';
const connectorSetup = {
bedrock: {
connectorTypeId: '.bedrock',
name: 'A bedrock action',
secrets: {
accessKey: 'bedrockAccessKey',
secret: 'bedrockSecret',
},
config: {
defaultModel: 'anthropic.claude-v2',
},
},
openai: {
connectorTypeId: '.gen-ai',
name: 'An openai action',
secrets: {
apiKey: 'genAiApiKey',
},
config: {
apiProvider: 'OpenAI',
},
},
};
/**
* Creates a connector
* @param supertest The supertest agent.
* @param apiUrl The url of the api
* @param connectorType The type of connector to create
* @param spaceId The space id
*/
export const createConnector = async (
supertest: SuperTest.Agent,
objectRemover: ObjectRemover,
apiUrl: string,
connectorType: 'bedrock' | 'openai',
spaceId?: string
) => {
const { connectorTypeId, config, name, secrets } = connectorSetup[connectorType];
const result = await supertest
.post(`${getUrlPrefix(spaceId ?? 'default')}/api/actions/connector`)
.set('kbn-xsrf', 'foo')
.set(ELASTIC_HTTP_VERSION_HEADER, '1')
.set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana')
.send({
name,
connector_type_id: connectorTypeId,
config: { ...config, apiUrl },
secrets,
})
.expect(200);
const { body } = result;
objectRemover.add(spaceId ?? 'default', body.id, 'connector', 'actions');
return body.id;
};

55
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/utils/object_remover.ts
View file

@ -1,55 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import {
ELASTIC_HTTP_VERSION_HEADER,
X_ELASTIC_INTERNAL_ORIGIN_REQUEST,
} from '@kbn/core-http-common';
import { getUrlPrefix } from './space_test_utils';
interface ObjectToRemove {
spaceId: string;
id: string;
type: string;
plugin: string;
isInternal?: boolean;
}
export class ObjectRemover {
private readonly supertest: any;
private objectsToRemove: ObjectToRemove[] = [];
constructor(supertest: any) {
this.supertest = supertest;
}
add(
spaceId: ObjectToRemove['spaceId'],
id: ObjectToRemove['id'],
type: ObjectToRemove['type'],
plugin: ObjectToRemove['plugin'],
isInternal?: ObjectToRemove['isInternal']
) {
this.objectsToRemove.push({ spaceId, id, type, plugin, isInternal });
}
async removeAll() {
await Promise.all(
this.objectsToRemove.map(({ spaceId, id, type, plugin, isInternal }) => {
return this.supertest
.delete(
`${getUrlPrefix(spaceId)}/${isInternal ? 'internal' : 'api'}/${plugin}/${type}/${id}`
)
.set('kbn-xsrf', 'foo')
.set(ELASTIC_HTTP_VERSION_HEADER, '1')
.set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana')
.expect(plugin === 'saved_objects' ? 200 : 204);
})
);
this.objectsToRemove = [];
}
}

34
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/utils/post_actions_client_execute.ts
View file

@ -1,34 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import type SuperTest from 'supertest';
import {
ELASTIC_HTTP_VERSION_HEADER,
X_ELASTIC_INTERNAL_ORIGIN_REQUEST,
} from '@kbn/core-http-common';
import { Response } from 'superagent';
/**
* Executes an invoke AI action
* @param connectorId The connector id
* @param args The arguments to pass to the action
* @param supertest The supertest agent
*/
export const postActionsClientExecute = async (
connectorId: string,
args: any,
supertest: SuperTest.Agent
): Promise<Response> => {
const response = await supertest
.post(`/internal/elastic_assistant/actions/connector/${connectorId}/_execute`)
.set('kbn-xsrf', 'true')
.set(ELASTIC_HTTP_VERSION_HEADER, '1')
.set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana')
.send(args);
return response;
};

10
x-pack/test/security_solution_api_integration/test_suites/genai/invoke_ai/utils/space_test_utils.ts
View file

@ -1,10 +0,0 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
export function getUrlPrefix(spaceId: string) {
return spaceId && spaceId !== 'default' ? `/s/${spaceId}` : ``;
}