github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-23 09:19:04 -04:00
Code Issues Projects Releases Packages Wiki Activity

[Security solution] AI Assistant Telemetry for Knowledge Base (#173552)

Browse source
This commit is contained in:
Steph Milovic 2023-12-22 13:26:28 -06:00 • committed by GitHub
parent 4cdb3c9894
commit fa47b572f3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
55 changed files with 851 additions and 177 deletions
Download patch file Download diff file Expand all files Collapse all files

8
api_docs/elastic_assistant.devdocs.json
View file

@ -221,10 +221,10 @@
},
{
"parentPluginId": "elasticAssistant",
"id": "def-server.AssistantToolParams.assistantLangChain",
"id": "def-server.AssistantToolParams.isEnabledKnowledgeBase",
"type": "boolean",
"tags": [],
"label": "assistantLangChain",
"label": "isEnabledKnowledgeBase",
"description": [],
"path": "x-pack/plugins/elastic_assistant/server/types.ts",
"deprecated": false,
@ -1546,7 +1546,7 @@
"section": "def-common.KibanaRequest",
"text": "KibanaRequest"
},
"<unknown, unknown, { params: { subActionParams: { messages: { role: \"user\" | \"system\" | \"assistant\"; content: string; }[]; } & { model?: string | undefined; n?: number | undefined; stop?: string | string[] | null | undefined; temperature?: number | undefined; }; subAction: string; }; alertsIndexPattern: string | undefined; allow: string[] | undefined; allowReplacement: string[] | undefined; assistantLangChain: boolean; replacements: { [x: string]: string; } | undefined; size: number | undefined; }, any>"
"<unknown, unknown, { params: { subActionParams: { messages: { role: \"user\" | \"system\" | \"assistant\"; content: string; }[]; } & { model?: string | undefined; n?: number | undefined; stop?: string | string[] | null | undefined; temperature?: number | undefined; }; subAction: string; }; alertsIndexPattern: string | undefined; allow: string[] | undefined; allowReplacement: string[] | undefined; isEnabledKnowledgeBase: boolean; replacements: { [x: string]: string; } | undefined; size: number | undefined; }, any>"
],
"path": "x-pack/plugins/elastic_assistant/server/types.ts",
"deprecated": false,
@ -1876,4 +1876,4 @@
"misc": [],
"objects": []
}
}
}

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

@ -19,8 +19,8 @@ describe('AlertsSettings', () => {
it('updates the knowledgeBase settings when the switch is toggled', () => {
const knowledgeBase: KnowledgeBaseConfig = {
alerts: false,
assistantLangChain: false,
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};
const setUpdatedKnowledgeBaseSettings = jest.fn();
@ -36,8 +36,8 @@ describe('AlertsSettings', () => {
fireEvent.click(alertsSwitch);
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
alerts: true,
assistantLangChain: false,
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
});
});
@ -45,8 +45,8 @@ describe('AlertsSettings', () => {
it('updates the knowledgeBase settings when the alerts range slider is changed', () => {
const setUpdatedKnowledgeBaseSettings = jest.fn();
const knowledgeBase: KnowledgeBaseConfig = {
alerts: true,
assistantLangChain: false,
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};
@ -61,17 +61,17 @@ describe('AlertsSettings', () => {
fireEvent.change(rangeSlider, { target: { value: '10' } });
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
alerts: true,
assistantLangChain: false,
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: false,
latestAlerts: 10,
});
});
it('enables the alerts range slider when knowledgeBase.alerts is true', () => {
it('enables the alerts range slider when knowledgeBase.isEnabledRAGAlerts is true', () => {
const setUpdatedKnowledgeBaseSettings = jest.fn();
const knowledgeBase: KnowledgeBaseConfig = {
alerts: true, // <-- true
assistantLangChain: false,
isEnabledRAGAlerts: true, // <-- true
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};
@ -85,11 +85,11 @@ describe('AlertsSettings', () => {
expect(screen.getByTestId('alertsRange')).not.toBeDisabled();
});
it('disables the alerts range slider when knowledgeBase.alerts is false', () => {
it('disables the alerts range slider when knowledgeBase.isEnabledRAGAlerts is false', () => {
const setUpdatedKnowledgeBaseSettings = jest.fn();
const knowledgeBase: KnowledgeBaseConfig = {
alerts: false, // <-- false
assistantLangChain: false,
isEnabledRAGAlerts: false, // <-- false
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};

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

@ -40,7 +40,7 @@ const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSetting
(event: EuiSwitchEvent) => {
setUpdatedKnowledgeBaseSettings({
...knowledgeBase,
alerts: event.target.checked,
isEnabledRAGAlerts: event.target.checked,
});
},
[knowledgeBase, setUpdatedKnowledgeBaseSettings]
@ -58,7 +58,7 @@ const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSetting
`}
>
<EuiSwitch
checked={knowledgeBase.alerts}
checked={knowledgeBase.isEnabledRAGAlerts}
compressed
data-test-subj="alertsSwitch"
label={i18n.ALERTS_LABEL}
@ -81,7 +81,7 @@ const AlertsSettingsComponent = ({ knowledgeBase, setUpdatedKnowledgeBaseSetting
aria-label={i18n.ALERTS_RANGE}
compressed
data-test-subj="alertsRange"
disabled={!knowledgeBase.alerts}
disabled={!knowledgeBase.isEnabledRAGAlerts}
id={inputRangeSliderId}
max={MAX_LATEST_ALERTS}
min={MIN_LATEST_ALERTS}

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

@ -35,9 +35,9 @@ const messages: Message[] = [
{ content: 'This is a test', role: 'user', timestamp: new Date().toLocaleString() },
];
const fetchConnectorArgs: FetchConnectorExecuteAction = {
alerts: false,
isEnabledRAGAlerts: false,
apiConfig,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
assistantStreamingEnabled: true,
http: mockHttp,
messages,
@ -49,13 +49,13 @@ describe('API tests', () => {
});
describe('fetchConnectorExecuteAction', () => {
it('calls the internal assistant API when assistantLangChain is true', async () => {
it('calls the internal assistant API when isEnabledKnowledgeBase is true', async () => {
await fetchConnectorExecuteAction(fetchConnectorArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"assistantLangChain":true}',
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
headers: { 'Content-Type': 'application/json' },
method: 'POST',
signal: undefined,
@ -63,10 +63,10 @@ describe('API tests', () => {
);
});
it('calls the actions connector api with streaming when assistantStreamingEnabled is true when assistantLangChain is false', async () => {
it('calls the actions connector api with streaming when assistantStreamingEnabled is true when isEnabledKnowledgeBase is false', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
assistantLangChain: false,
isEnabledKnowledgeBase: false,
};
await fetchConnectorExecuteAction(testProps);
@ -74,7 +74,7 @@ describe('API tests', () => {
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeStream"},"assistantLangChain":false}',
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeStream"},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":false}',
method: 'POST',
asResponse: true,
rawResponse: true,
@ -86,7 +86,7 @@ describe('API tests', () => {
it('calls the actions connector with the expected optional request parameters', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
alerts: true,
isEnabledRAGAlerts: true,
alertsIndexPattern: '.alerts-security.alerts-default',
allow: ['a', 'b', 'c'],
allowReplacement: ['b', 'c'],
@ -99,7 +99,7 @@ describe('API tests', () => {
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"assistantLangChain":true,"alertsIndexPattern":".alerts-security.alerts-default","allow":["a","b","c"],"allowReplacement":["b","c"],"replacements":{"auuid":"real.hostname"},"size":30}',
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":true,"alertsIndexPattern":".alerts-security.alerts-default","allow":["a","b","c"],"allowReplacement":["b","c"],"replacements":{"auuid":"real.hostname"},"size":30}',
headers: {
'Content-Type': 'application/json',
},
@ -109,10 +109,10 @@ describe('API tests', () => {
);
});
it('calls the actions connector api with invoke when assistantStreamingEnabled is false when assistantLangChain is false', async () => {
it('calls the actions connector api with invoke when assistantStreamingEnabled is false when isEnabledKnowledgeBase is false', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
assistantLangChain: false,
isEnabledKnowledgeBase: false,
assistantStreamingEnabled: false,
};
@ -121,7 +121,7 @@ describe('API tests', () => {
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"assistantLangChain":false}',
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":false}',
method: 'POST',
headers: {
'Content-Type': 'application/json',
@ -131,11 +131,11 @@ describe('API tests', () => {
);
});
it('calls the actions connector api with invoke when assistantStreamingEnabled is true when assistantLangChain is false and alerts is true', async () => {
it('calls the actions connector api with invoke when assistantStreamingEnabled is true when isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
assistantLangChain: false,
alerts: true,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: true,
};
await fetchConnectorExecuteAction(testProps);
@ -143,7 +143,7 @@ describe('API tests', () => {
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/actions/connector/foo/_execute',
{
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"assistantLangChain":false}',
body: '{"params":{"subActionParams":{"model":"gpt-4","messages":[{"role":"user","content":"This is a test"}],"n":1,"stop":null,"temperature":0.2},"subAction":"invokeAI"},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
method: 'POST',
headers: {
'Content-Type': 'application/json',
@ -168,7 +168,7 @@ describe('API tests', () => {
});
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
assistantLangChain: false,
isEnabledKnowledgeBase: false,
assistantStreamingEnabled: false,
};
@ -186,7 +186,7 @@ describe('API tests', () => {
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
assistantLangChain: false,
isEnabledKnowledgeBase: false,
};
const result = await fetchConnectorExecuteAction(testProps);
@ -205,7 +205,7 @@ describe('API tests', () => {
});
const testProps: FetchConnectorExecuteAction = {
...fetchConnectorArgs,
assistantLangChain: false,
isEnabledKnowledgeBase: false,
};
const result = await fetchConnectorExecuteAction(testProps);
@ -225,7 +225,7 @@ describe('API tests', () => {
expect(result).toEqual({ response: API_ERROR, isStream: false, isError: true });
});
it('returns the value of the action_input property when assistantLangChain is true, and `content` has properly prefixed and suffixed JSON with the action_input property', async () => {
it('returns the value of the action_input property when isEnabledKnowledgeBase is true, and `content` has properly prefixed and suffixed JSON with the action_input property', async () => {
const response = '```json\n{"action_input": "value from action_input"}\n```';
(mockHttp.fetch as jest.Mock).mockResolvedValue({
@ -242,7 +242,7 @@ describe('API tests', () => {
});
});
it('returns the original content when assistantLangChain is true, and `content` has properly formatted JSON WITHOUT the action_input property', async () => {
it('returns the original content when isEnabledKnowledgeBase is true, and `content` has properly formatted JSON WITHOUT the action_input property', async () => {
const response = '```json\n{"some_key": "some value"}\n```';
(mockHttp.fetch as jest.Mock).mockResolvedValue({
@ -255,7 +255,7 @@ describe('API tests', () => {
expect(result).toEqual({ response, isStream: false, isError: false });
});
it('returns the original when assistantLangChain is true, and `content` is not JSON', async () => {
it('returns the original when isEnabledKnowledgeBase is true, and `content` is not JSON', async () => {
const response = 'plain text content';
(mockHttp.fetch as jest.Mock).mockResolvedValue({

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

@ -19,11 +19,11 @@ import {
import { PerformEvaluationParams } from './settings/evaluation_settings/use_perform_evaluation';
export interface FetchConnectorExecuteAction {
alerts: boolean;
isEnabledRAGAlerts: boolean;
alertsIndexPattern?: string;
allow?: string[];
allowReplacement?: string[];
assistantLangChain: boolean;
isEnabledKnowledgeBase: boolean;
assistantStreamingEnabled: boolean;
apiConfig: Conversation['apiConfig'];
http: HttpSetup;
@ -45,11 +45,11 @@ export interface FetchConnectorExecuteResponse {
}
export const fetchConnectorExecuteAction = async ({
alerts,
isEnabledRAGAlerts,
alertsIndexPattern,
allow,
allowReplacement,
assistantLangChain,
isEnabledKnowledgeBase,
assistantStreamingEnabled,
http,
messages,
@ -82,9 +82,9 @@ export const fetchConnectorExecuteAction = async ({
// tracked here: https://github.com/elastic/security-team/issues/7363
// In part 3 I will make enhancements to langchain to introduce streaming
// Once implemented, invokeAI can be removed
const isStream = assistantStreamingEnabled && !assistantLangChain && !alerts;
const isStream = assistantStreamingEnabled && !isEnabledKnowledgeBase && !isEnabledRAGAlerts;
const optionalRequestParams = getOptionalRequestParams({
alerts,
isEnabledRAGAlerts,
alertsIndexPattern,
allow,
allowReplacement,
@ -98,7 +98,8 @@ export const fetchConnectorExecuteAction = async ({
subActionParams: body,
subAction: 'invokeStream',
},
assistantLangChain,
isEnabledKnowledgeBase,
isEnabledRAGAlerts,
...optionalRequestParams,
}
: {
@ -106,7 +107,8 @@ export const fetchConnectorExecuteAction = async ({
subActionParams: body,
subAction: 'invokeAI',
},
assistantLangChain,
isEnabledKnowledgeBase,
isEnabledRAGAlerts,
...optionalRequestParams,
};
@ -187,8 +189,8 @@ export const fetchConnectorExecuteAction = async ({
return {
response: hasParsableResponse({
alerts,
assistantLangChain,
isEnabledRAGAlerts,
isEnabledKnowledgeBase,
})
? getFormattedMessageContent(response.data)
: response.data,

1
x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.test.tsx
View file

@ -15,6 +15,7 @@ const assistantTelemetry = {
reportAssistantInvoked,
reportAssistantMessageSent: () => {},
reportAssistantQuickPrompt: () => {},
reportAssistantSettingToggled: () => {},
};
describe('AssistantOverlay', () => {
beforeEach(() => {

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

@ -237,7 +237,7 @@ describe('getBlockBotConversation', () => {
describe('getOptionalRequestParams', () => {
it('should return an empty object when alerts is false', () => {
const params = {
alerts: false, // <-- false
isEnabledRAGAlerts: false, // <-- false
alertsIndexPattern: 'indexPattern',
allow: ['a', 'b', 'c'],
allowReplacement: ['b', 'c'],
@ -252,7 +252,7 @@ describe('getBlockBotConversation', () => {
it('should return the optional request params when alerts is true', () => {
const params = {
alerts: true,
isEnabledRAGAlerts: true,
alertsIndexPattern: 'indexPattern',
allow: ['a', 'b', 'c'],
allowReplacement: ['b', 'c'],
@ -273,7 +273,7 @@ describe('getBlockBotConversation', () => {
it('should return (only) the optional request params that are defined when some optional params are not provided', () => {
const params = {
alerts: true,
isEnabledRAGAlerts: true,
allow: ['a', 'b', 'c'], // all the others are undefined
};
@ -286,37 +286,37 @@ describe('getBlockBotConversation', () => {
});
describe('hasParsableResponse', () => {
it('returns true when just assistantLangChain is true', () => {
it('returns true when just isEnabledKnowledgeBase is true', () => {
const result = hasParsableResponse({
alerts: false,
assistantLangChain: true,
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: true,
});
expect(result).toBe(true);
});
it('returns true when just alerts is true', () => {
it('returns true when just isEnabledRAGAlerts is true', () => {
const result = hasParsableResponse({
alerts: true,
assistantLangChain: false,
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: false,
});
expect(result).toBe(true);
});
it('returns true when both assistantLangChain and alerts are true', () => {
it('returns true when both isEnabledKnowledgeBase and isEnabledRAGAlerts are true', () => {
const result = hasParsableResponse({
alerts: true,
assistantLangChain: true,
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: true,
});
expect(result).toBe(true);
});
it('returns false when both assistantLangChain and alerts are false', () => {
it('returns false when both isEnabledKnowledgeBase and isEnabledRAGAlerts are false', () => {
const result = hasParsableResponse({
alerts: false,
assistantLangChain: false,
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
});
expect(result).toBe(false);

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

@ -97,14 +97,14 @@ interface OptionalRequestParams {
}
export const getOptionalRequestParams = ({
alerts,
isEnabledRAGAlerts,
alertsIndexPattern,
allow,
allowReplacement,
replacements,
size,
}: {
alerts: boolean;
isEnabledRAGAlerts: boolean;
alertsIndexPattern?: string;
allow?: string[];
allowReplacement?: string[];
@ -118,7 +118,7 @@ export const getOptionalRequestParams = ({
const optionalSize = size ? { size } : undefined;
// the settings toggle must be enabled:
if (!alerts) {
if (!isEnabledRAGAlerts) {
return {}; // don't send any optional params
}
@ -132,9 +132,9 @@ export const getOptionalRequestParams = ({
};
export const hasParsableResponse = ({
alerts,
assistantLangChain,
isEnabledRAGAlerts,
isEnabledKnowledgeBase,
}: {
alerts: boolean;
assistantLangChain: boolean;
}): boolean => assistantLangChain || alerts;
isEnabledRAGAlerts: boolean;
isEnabledKnowledgeBase: boolean;
}): boolean => isEnabledKnowledgeBase || isEnabledRAGAlerts;

4
x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt_editor/system_prompt/select_system_prompt/index.test.tsx
View file

@ -48,6 +48,10 @@ const mockUseAssistantContext = {
},
],
setAllSystemPrompts: jest.fn(),
knowledgeBase: {
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
},
};
jest.mock('../../../../assistant_context', () => {
const original = jest.requireActual('../../../../assistant_context');

2
x-pack/packages/kbn-elastic-assistant/impl/assistant/quick_prompts/quick_prompts.test.tsx
View file

@ -26,7 +26,7 @@ const mockUseAssistantContext = {
promptContexts: {},
allQuickPrompts: MOCK_QUICK_PROMPTS,
knowledgeBase: {
assistantLangChain: true,
isEnabledKnowledgeBase: true,
},
};

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

@ -41,7 +41,7 @@ export const QuickPrompts: React.FC<QuickPromptsProps> = React.memo(
const contextFilteredQuickPrompts = useMemo(() => {
const registeredPromptContextTitles = Object.values(promptContexts).map((pc) => pc.category);
// If KB is enabled, include KNOWLEDGE_BASE_CATEGORY so KB dependent quick prompts are shown
if (knowledgeBase.assistantLangChain) {
if (knowledgeBase.isEnabledKnowledgeBase) {
registeredPromptContextTitles.push(KNOWLEDGE_BASE_CATEGORY);
}
return allQuickPrompts.filter((quickPrompt) => {
@ -54,7 +54,7 @@ export const QuickPrompts: React.FC<QuickPromptsProps> = React.memo(
});
}
});
}, [allQuickPrompts, knowledgeBase.assistantLangChain, promptContexts]);
}, [allQuickPrompts, knowledgeBase.isEnabledKnowledgeBase, promptContexts]);
// Overflow state
const [isOverflowPopoverOpen, setIsOverflowPopoverOpen] = useState(false);

56
x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/use_settings_updater/use_settings_updater.test.tsx
View file

@ -33,15 +33,18 @@ const setConversationsMock = jest.fn();
const setDefaultAllowMock = jest.fn();
const setDefaultAllowReplacementMock = jest.fn();
const setKnowledgeBaseMock = jest.fn();
const reportAssistantSettingToggled = jest.fn();
const mockValues = {
assistantTelemetry: { reportAssistantSettingToggled },
conversations: mockConversations,
allSystemPrompts: mockSystemPrompts,
allQuickPrompts: mockQuickPrompts,
defaultAllow: initialDefaultAllow,
defaultAllowReplacement: initialDefaultAllowReplacement,
knowledgeBase: {
assistantLangChain: true,
isEnabledRAGAlerts: true,
isEnabledKnowledgeBase: true,
latestAlerts: DEFAULT_LATEST_ALERTS,
},
setAllQuickPrompts: setAllQuickPromptsMock,
setConversations: setConversationsMock,
@ -58,8 +61,8 @@ const updatedValues = {
defaultAllow: ['allow2'],
defaultAllowReplacement: ['replacement2'],
knowledgeBase: {
alerts: false,
assistantLangChain: false,
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
},
};
@ -73,6 +76,9 @@ jest.mock('../../../assistant_context', () => {
});
describe('useSettingsUpdater', () => {
beforeEach(() => {
jest.clearAllMocks();
});
it('should set all state variables to their initial values when resetSettings is called', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useSettingsUpdater());
@ -144,4 +150,46 @@ describe('useSettingsUpdater', () => {
expect(setKnowledgeBaseMock).toHaveBeenCalledWith(updatedValues.knowledgeBase);
});
});
it('should track which toggles have been updated when saveSettings is called', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useSettingsUpdater());
await waitForNextUpdate();
const { setUpdatedKnowledgeBaseSettings } = result.current;
setUpdatedKnowledgeBaseSettings(updatedValues.knowledgeBase);
result.current.saveSettings();
expect(reportAssistantSettingToggled).toHaveBeenCalledWith({
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
});
});
});
it('should track only toggles that updated', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useSettingsUpdater());
await waitForNextUpdate();
const { setUpdatedKnowledgeBaseSettings } = result.current;
setUpdatedKnowledgeBaseSettings({
...updatedValues.knowledgeBase,
isEnabledKnowledgeBase: true,
});
result.current.saveSettings();
expect(reportAssistantSettingToggled).toHaveBeenCalledWith({
isEnabledRAGAlerts: false,
});
});
});
it('if no toggles update, do not track anything', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useSettingsUpdater());
await waitForNextUpdate();
const { setUpdatedKnowledgeBaseSettings } = result.current;
setUpdatedKnowledgeBaseSettings(mockValues.knowledgeBase);
result.current.saveSettings();
expect(reportAssistantSettingToggled).not.toHaveBeenCalledWith();
});
});
});

18
x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/use_settings_updater/use_settings_updater.tsx
View file

@ -34,6 +34,7 @@ export const useSettingsUpdater = (): UseSettingsUpdater => {
const {
allQuickPrompts,
allSystemPrompts,
assistantTelemetry,
conversations,
defaultAllow,
defaultAllowReplacement,
@ -92,10 +93,27 @@ export const useSettingsUpdater = (): UseSettingsUpdater => {
setAllQuickPrompts(updatedQuickPromptSettings);
setAllSystemPrompts(updatedSystemPromptSettings);
setConversations(updatedConversationSettings);
const didUpdateKnowledgeBase =
knowledgeBase.isEnabledKnowledgeBase !== updatedKnowledgeBaseSettings.isEnabledKnowledgeBase;
const didUpdateRAGAlerts =
knowledgeBase.isEnabledRAGAlerts !== updatedKnowledgeBaseSettings.isEnabledRAGAlerts;
if (didUpdateKnowledgeBase || didUpdateRAGAlerts) {
assistantTelemetry?.reportAssistantSettingToggled({
...(didUpdateKnowledgeBase
? { isEnabledKnowledgeBase: updatedKnowledgeBaseSettings.isEnabledKnowledgeBase }
: {}),
...(didUpdateRAGAlerts
? { isEnabledRAGAlerts: updatedKnowledgeBaseSettings.isEnabledRAGAlerts }
: {}),
});
}
setKnowledgeBase(updatedKnowledgeBaseSettings);
setDefaultAllow(updatedDefaultAllow);
setDefaultAllowReplacement(updatedDefaultAllowReplacement);
}, [
assistantTelemetry,
knowledgeBase.isEnabledRAGAlerts,
knowledgeBase.isEnabledKnowledgeBase,
setAllQuickPrompts,
setAllSystemPrompts,
setConversations,

4
x-pack/packages/kbn-elastic-assistant/impl/assistant/types.ts
View file

@ -17,7 +17,7 @@ export interface Prompt {
}
export interface KnowledgeBaseConfig {
alerts: boolean;
assistantLangChain: boolean;
isEnabledRAGAlerts: boolean;
isEnabledKnowledgeBase: boolean;
latestAlerts: number;
}

40
x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/index.test.tsx
View file

@ -36,6 +36,9 @@ const mockConvo = {
};
describe('useConversation', () => {
beforeEach(() => {
jest.clearAllMocks();
});
it('should append a message to an existing conversation when called with valid conversationId and message', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useConversation(), {
@ -63,6 +66,43 @@ describe('useConversation', () => {
});
});
it('should report telemetry when a message has been sent', async () => {
await act(async () => {
const reportAssistantMessageSent = jest.fn();
const { result, waitForNextUpdate } = renderHook(() => useConversation(), {
wrapper: ({ children }) => (
<TestProviders
providerContext={{
getInitialConversations: () => ({
[alertConvo.id]: alertConvo,
[welcomeConvo.id]: welcomeConvo,
}),
assistantTelemetry: {
reportAssistantInvoked: () => {},
reportAssistantQuickPrompt: () => {},
reportAssistantSettingToggled: () => {},
reportAssistantMessageSent,
},
}}
>
{children}
</TestProviders>
),
});
await waitForNextUpdate();
result.current.appendMessage({
conversationId: welcomeConvo.id,
message,
});
expect(reportAssistantMessageSent).toHaveBeenCalledWith({
conversationId: 'Welcome',
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
role: 'user',
});
});
});
it('should create a new conversation when called with valid conversationId and message', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useConversation(), {

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

@ -75,7 +75,12 @@ interface UseConversation {
}
export const useConversation = (): UseConversation => {
const { allSystemPrompts, assistantTelemetry, setConversations } = useAssistantContext();
const {
allSystemPrompts,
assistantTelemetry,
knowledgeBase: { isEnabledKnowledgeBase, isEnabledRAGAlerts },
setConversations,
} = useAssistantContext();
/**
* Removes the last message of conversation[] for a given conversationId
@ -140,7 +145,12 @@ export const useConversation = (): UseConversation => {
*/
const appendMessage = useCallback(
({ conversationId, message }: AppendMessageProps): Message[] => {
assistantTelemetry?.reportAssistantMessageSent({ conversationId, role: message.role });
assistantTelemetry?.reportAssistantMessageSent({
conversationId,
role: message.role,
isEnabledKnowledgeBase,
isEnabledRAGAlerts,
});
let messages: Message[] = [];
setConversations((prev: Record<string, Conversation>) => {
const prevConversation: Conversation | undefined = prev[conversationId];
@ -161,7 +171,7 @@ export const useConversation = (): UseConversation => {
});
return messages;
},
[assistantTelemetry, setConversations]
[isEnabledKnowledgeBase, isEnabledRAGAlerts, assistantTelemetry, setConversations]
);
const appendReplacements = useCallback(

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

@ -47,12 +47,12 @@ export const useSendMessages = (): UseSendMessages => {
try {
return await fetchConnectorExecuteAction({
alerts: knowledgeBase.alerts, // settings toggle
isEnabledRAGAlerts: knowledgeBase.isEnabledRAGAlerts, // settings toggle
alertsIndexPattern,
allow: defaultAllow,
allowReplacement: defaultAllowReplacement,
apiConfig,
assistantLangChain: knowledgeBase.assistantLangChain,
isEnabledKnowledgeBase: knowledgeBase.isEnabledKnowledgeBase,
assistantStreamingEnabled,
http,
replacements,
@ -69,8 +69,8 @@ export const useSendMessages = (): UseSendMessages => {
assistantStreamingEnabled,
defaultAllow,
defaultAllowReplacement,
knowledgeBase.alerts,
knowledgeBase.assistantLangChain,
knowledgeBase.isEnabledRAGAlerts,
knowledgeBase.isEnabledKnowledgeBase,
knowledgeBase.latestAlerts,
]
);

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

@ -17,7 +17,7 @@ export const KNOWLEDGE_BASE_LOCAL_STORAGE_KEY = 'knowledgeBase';
export const DEFAULT_LATEST_ALERTS = 20;
export const DEFAULT_KNOWLEDGE_BASE_SETTINGS: KnowledgeBaseConfig = {
alerts: false,
assistantLangChain: false,
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
};

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

@ -321,7 +321,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
docLinks,
getComments,
http,
knowledgeBase: localStorageKnowledgeBase ?? DEFAULT_KNOWLEDGE_BASE_SETTINGS,
knowledgeBase: { ...DEFAULT_KNOWLEDGE_BASE_SETTINGS, ...localStorageKnowledgeBase },
modelEvaluatorEnabled,
promptContexts,
nameSpace,

11
x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx
View file

@ -67,8 +67,17 @@ export interface Conversation {
export interface AssistantTelemetry {
reportAssistantInvoked: (params: { invokedBy: string; conversationId: string }) => void;
reportAssistantMessageSent: (params: { conversationId: string; role: string }) => void;
reportAssistantMessageSent: (params: {
conversationId: string;
role: string;
isEnabledKnowledgeBase: boolean;
isEnabledRAGAlerts: boolean;
}) => void;
reportAssistantQuickPrompt: (params: { conversationId: string; promptTitle: string }) => void;
reportAssistantSettingToggled: (params: {
isEnabledKnowledgeBase?: boolean;
isEnabledRAGAlerts?: boolean;
}) => void;
}
export interface AssistantAvailability {

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

@ -38,8 +38,8 @@ jest.mock('../assistant_context', () => {
const setUpdatedKnowledgeBaseSettings = jest.fn();
const defaultProps = {
knowledgeBase: {
assistantLangChain: true,
alerts: false,
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
},
setUpdatedKnowledgeBaseSettings,
@ -117,16 +117,16 @@ describe('Knowledge base settings', () => {
fireEvent.click(getByTestId('esqlEnableButton'));
expect(mockSetup).toHaveBeenCalledWith('esql');
});
it('On disable lang chain, set assistantLangChain to false', () => {
it('On disable lang chain, set isEnabledKnowledgeBase to false', () => {
const { getByTestId } = render(
<TestProviders>
<KnowledgeBaseSettings {...defaultProps} />
</TestProviders>
);
fireEvent.click(getByTestId('assistantLangChainSwitch'));
fireEvent.click(getByTestId('isEnabledKnowledgeBaseSwitch'));
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
alerts: false,
assistantLangChain: false,
isEnabledRAGAlerts: false,
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
});
@ -138,17 +138,17 @@ describe('Knowledge base settings', () => {
<KnowledgeBaseSettings
{...defaultProps}
knowledgeBase={{
assistantLangChain: false,
alerts: false,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
}}
/>
</TestProviders>
);
fireEvent.click(getByTestId('assistantLangChainSwitch'));
fireEvent.click(getByTestId('isEnabledKnowledgeBaseSwitch'));
expect(setUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({
assistantLangChain: true,
alerts: false,
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
});

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

@ -63,11 +63,11 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
// Resource availability state
const isLoadingKb = isLoading || isFetching || isSettingUpKB || isDeletingUpKB;
const isKnowledgeBaseAvailable = knowledgeBase.assistantLangChain && kbStatus?.elser_exists;
const isKnowledgeBaseAvailable = knowledgeBase.isEnabledKnowledgeBase && kbStatus?.elser_exists;
const isESQLAvailable =
knowledgeBase.assistantLangChain && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled;
knowledgeBase.isEnabledKnowledgeBase && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled;
// Prevent enabling if elser doesn't exist, but always allow to disable
const isSwitchDisabled = !kbStatus?.elser_exists && !knowledgeBase.assistantLangChain;
const isSwitchDisabled = !kbStatus?.elser_exists && !knowledgeBase.isEnabledKnowledgeBase;
// Calculated health state for EuiHealth component
const elserHealth = isElserEnabled ? 'success' : 'subdued';
@ -75,13 +75,13 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
const esqlHealth = isESQLEnabled ? 'success' : 'subdued';
//////////////////////////////////////////////////////////////////////////////////////////
// Main `Knowledge Base` switch, which toggles the `assistantLangChain` UI feature toggle
// Main `Knowledge Base` switch, which toggles the `isEnabledKnowledgeBase` UI feature toggle
// setting that is saved to localstorage
const onEnableAssistantLangChainChange = useCallback(
(event: EuiSwitchEvent) => {
setUpdatedKnowledgeBaseSettings({
...knowledgeBase,
assistantLangChain: event.target.checked,
isEnabledKnowledgeBase: event.target.checked,
});
// If enabling and ELSER exists, try to set up automatically
@ -92,16 +92,16 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
[kbStatus?.elser_exists, knowledgeBase, setUpdatedKnowledgeBaseSettings, setupKB]
);
const assistantLangChainSwitch = useMemo(() => {
const isEnabledKnowledgeBaseSwitch = useMemo(() => {
return isLoadingKb ? (
<EuiLoadingSpinner size="s" />
) : (
<EuiToolTip content={isSwitchDisabled && i18n.KNOWLEDGE_BASE_TOOLTIP} position={'right'}>
<EuiSwitch
showLabel={false}
data-test-subj="assistantLangChainSwitch"
data-test-subj="isEnabledKnowledgeBaseSwitch"
disabled={isSwitchDisabled}
checked={knowledgeBase.assistantLangChain}
checked={knowledgeBase.isEnabledKnowledgeBase}
onChange={onEnableAssistantLangChainChange}
label={i18n.KNOWLEDGE_BASE_LABEL}
compressed
@ -111,7 +111,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
}, [
isLoadingKb,
isSwitchDisabled,
knowledgeBase.assistantLangChain,
knowledgeBase.isEnabledKnowledgeBase,
onEnableAssistantLangChainChange,
]);
@ -221,7 +221,7 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
}
`}
>
{assistantLangChainSwitch}
{isEnabledKnowledgeBaseSwitch}
</EuiFormRow>
<EuiSpacer size="s" />

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

@ -26,6 +26,7 @@ export const createMockClients = () => {
actions: actionsClientMock.create(),
getRegisteredTools: jest.fn(),
logger: loggingSystemMock.createLogger(),
telemetry: coreMock.createSetup().analytics,
},
savedObjectsClient: core.savedObjects.client,
@ -75,6 +76,7 @@ const createElasticAssistantRequestContextMock = (
actions: clients.elasticAssistant.actions as unknown as ActionsPluginStart,
getRegisteredTools: jest.fn(),
logger: clients.elasticAssistant.logger,
telemetry: clients.elasticAssistant.telemetry,
};
};

34
x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.test.ts
View file

@ -20,6 +20,11 @@ import {
} from './elasticsearch_store';
import { mockMsearchResponse } from '../../../__mocks__/msearch_response';
import { mockQueryText } from '../../../__mocks__/query_text';
import { coreMock } from '@kbn/core/server/mocks';
import {
KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT,
KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT,
} from '../../telemetry/event_based_telemetry';
jest.mock('uuid', () => ({
v4: jest.fn(),
@ -36,10 +41,12 @@ jest.mock('@kbn/core/server', () => ({
const mockEsClient = elasticsearchServiceMock.createElasticsearchClient();
const mockLogger = loggingSystemMock.createLogger();
const reportEvent = jest.fn();
const mockTelemetry = { ...coreMock.createSetup().analytics, reportEvent };
const KB_INDEX = '.elastic-assistant-kb';
const getElasticsearchStore = () => {
return new ElasticsearchStore(mockEsClient, KB_INDEX, mockLogger);
return new ElasticsearchStore(mockEsClient, KB_INDEX, mockLogger, mockTelemetry);
};
describe('ElasticsearchStore', () => {
@ -415,5 +422,30 @@ describe('ElasticsearchStore', () => {
],
});
});
it('Reports successful telemetry event', async () => {
mockEsClient.msearch.mockResolvedValue(mockMsearchResponse);
await esStore.similaritySearch(mockQueryText);
expect(reportEvent).toHaveBeenCalledWith(KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT.eventType, {
model: '.elser_model_2',
resourceAccessed: 'esql',
responseTime: 142,
resultCount: 2,
});
});
it('Reports error telemetry event', async () => {
mockEsClient.msearch.mockRejectedValue(new Error('Oh no!'));
await esStore.similaritySearch(mockQueryText);
expect(reportEvent).toHaveBeenCalledWith(KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT.eventType, {
model: '.elser_model_2',
resourceAccessed: 'esql',
errorMessage: 'Oh no!',
});
});
});
});

23
x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts
View file

@ -5,7 +5,7 @@
* 2.0.
*/
import { ElasticsearchClient, Logger } from '@kbn/core/server';
import { type AnalyticsServiceSetup, ElasticsearchClient, Logger } from '@kbn/core/server';
import { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types';
import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
import { Callbacks } from 'langchain/callbacks';
@ -13,6 +13,7 @@ import { Document } from 'langchain/document';
import { VectorStore } from 'langchain/vectorstores/base';
import * as uuid from 'uuid';
import { transformError } from '@kbn/securitysolution-es-utils';
import { ElasticsearchEmbeddings } from '../embeddings/elasticsearch_embeddings';
import { FlattenedHit, getFlattenedHits } from './helpers/get_flattened_hits';
import { getMsearchQueryBody } from './helpers/get_msearch_query_body';
@ -25,6 +26,10 @@ import {
KNOWLEDGE_BASE_INGEST_PIPELINE,
} from '../../../routes/knowledge_base/constants';
import { getRequiredKbDocsTermsQueryDsl } from './helpers/get_required_kb_docs_terms_query_dsl';
import {
KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT,
KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT,
} from '../../telemetry/event_based_telemetry';
interface CreatePipelineParams {
id?: string;
@ -59,6 +64,7 @@ export class ElasticsearchStore extends VectorStore {
private readonly esClient: ElasticsearchClient;
private readonly index: string;
private readonly logger: Logger;
private readonly telemetry: AnalyticsServiceSetup;
private readonly model: string;
private readonly kbResource: string;
@ -70,6 +76,7 @@ export class ElasticsearchStore extends VectorStore {
esClient: ElasticsearchClient,
index: string,
logger: Logger,
telemetry: AnalyticsServiceSetup,
model?: string,
kbResource?: string | undefined
) {
@ -77,6 +84,7 @@ export class ElasticsearchStore extends VectorStore {
this.esClient = esClient;
this.index = index ?? KNOWLEDGE_BASE_INDEX_PATTERN;
this.logger = logger;
this.telemetry = telemetry;
this.model = model ?? '.elser_model_2';
this.kbResource = kbResource ?? ESQL_RESOURCE;
}
@ -222,6 +230,13 @@ export class ElasticsearchStore extends VectorStore {
return getFlattenedHits(maybeEsqlMsearchResponse);
});
this.telemetry.reportEvent(KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT.eventType, {
model: this.model,
resourceAccessed: this.kbResource,
resultCount: results.length,
responseTime: result.took ?? 0,
});
this.logger.debug(
`Similarity search metadata source:\n${JSON.stringify(
results.map((r) => r?.metadata?.source ?? '(missing metadata.source)'),
@ -232,6 +247,12 @@ export class ElasticsearchStore extends VectorStore {
return results;
} catch (e) {
const error = transformError(e);
this.telemetry.reportEvent(KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT.eventType, {
model: this.model,
resourceAccessed: this.kbResource,
errorMessage: error.message,
});
this.logger.error(e);
return [];
}

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

@ -7,6 +7,7 @@
import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server';
import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
import { coreMock } from '@kbn/core/server/mocks';
import { KibanaRequest } from '@kbn/core/server';
import { loggerMock } from '@kbn/logging-mocks';
@ -55,6 +56,7 @@ const mockRequest: KibanaRequest<unknown, unknown, any, any> = {} as KibanaReque
const mockActions: ActionsPluginStart = {} as ActionsPluginStart;
const mockLogger = loggerMock.create();
const mockTelemetry = coreMock.createSetup().analytics;
const esClientMock = elasticsearchServiceMock.createScopedClusterClient().asCurrentUser;
describe('callAgentExecutor', () => {
@ -69,7 +71,7 @@ describe('callAgentExecutor', () => {
it('creates an instance of ActionsClientLlm with the expected context from the request', async () => {
await callAgentExecutor({
actions: mockActions,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
connectorId: mockConnectorId,
esClient: esClientMock,
langChainMessages,
@ -77,6 +79,7 @@ describe('callAgentExecutor', () => {
onNewReplacements: jest.fn(),
request: mockRequest,
kbResource: ESQL_RESOURCE,
telemetry: mockTelemetry,
});
expect(ActionsClientLlm).toHaveBeenCalledWith({
@ -90,7 +93,7 @@ describe('callAgentExecutor', () => {
it('kicks off the chain with (only) the last message', async () => {
await callAgentExecutor({
actions: mockActions,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
connectorId: mockConnectorId,
esClient: esClientMock,
langChainMessages,
@ -98,6 +101,7 @@ describe('callAgentExecutor', () => {
onNewReplacements: jest.fn(),
request: mockRequest,
kbResource: ESQL_RESOURCE,
telemetry: mockTelemetry,
});
// We don't care about the `config` argument, so we use `expect.anything()`
@ -114,7 +118,7 @@ describe('callAgentExecutor', () => {
await callAgentExecutor({
actions: mockActions,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
connectorId: mockConnectorId,
esClient: esClientMock,
langChainMessages: onlyOneMessage,
@ -122,6 +126,7 @@ describe('callAgentExecutor', () => {
onNewReplacements: jest.fn(),
request: mockRequest,
kbResource: ESQL_RESOURCE,
telemetry: mockTelemetry,
});
// We don't care about the `config` argument, so we use `expect.anything()`
@ -136,7 +141,7 @@ describe('callAgentExecutor', () => {
it('returns the expected response body', async () => {
const result: ResponseBody = await callAgentExecutor({
actions: mockActions,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
connectorId: mockConnectorId,
esClient: esClientMock,
langChainMessages,
@ -144,6 +149,7 @@ describe('callAgentExecutor', () => {
onNewReplacements: jest.fn(),
request: mockRequest,
kbResource: ESQL_RESOURCE,
telemetry: mockTelemetry,
});
expect(result).toEqual({

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

@ -30,7 +30,7 @@ export const callAgentExecutor = async ({
alertsIndexPattern,
allow,
allowReplacement,
assistantLangChain,
isEnabledKnowledgeBase,
assistantTools = [],
connectorId,
elserId,
@ -43,6 +43,7 @@ export const callAgentExecutor = async ({
replacements,
request,
size,
telemetry,
traceOptions,
}: AgentExecutorParams): AgentExecutorResponse => {
const llm = new ActionsClientLlm({ actions, connectorId, request, llmType, logger });
@ -63,6 +64,7 @@ export const callAgentExecutor = async ({
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
logger,
telemetry,
elserId,
kbResource
);
@ -77,7 +79,7 @@ export const callAgentExecutor = async ({
allow,
allowReplacement,
alertsIndexPattern,
assistantLangChain,
isEnabledKnowledgeBase,
chain,
esClient,
modelExists,

2
x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts
View file

@ -36,6 +36,7 @@ export const callOpenAIFunctionsExecutor = async ({
request,
elserId,
kbResource,
telemetry,
traceOptions,
}: AgentExecutorParams): AgentExecutorResponse => {
const llm = new ActionsClientLlm({ actions, connectorId, request, llmType, logger });
@ -56,6 +57,7 @@ export const callOpenAIFunctionsExecutor = async ({
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
logger,
telemetry,
elserId,
kbResource
);

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

@ -11,6 +11,7 @@ import { BaseMessage } from 'langchain/schema';
import { Logger } from '@kbn/logging';
import { KibanaRequest } from '@kbn/core-http-server';
import type { LangChainTracer } from 'langchain/callbacks';
import type { AnalyticsServiceSetup } from '@kbn/core-analytics-server';
import { RequestBody, ResponseBody } from '../types';
import type { AssistantTool } from '../../../types';
@ -19,7 +20,7 @@ export interface AgentExecutorParams {
actions: ActionsPluginStart;
allow?: string[];
allowReplacement?: string[];
assistantLangChain: boolean;
isEnabledKnowledgeBase: boolean;
assistantTools?: AssistantTool[];
connectorId: string;
esClient: ElasticsearchClient;
@ -33,6 +34,7 @@ export interface AgentExecutorParams {
size?: number;
elserId?: string;
traceOptions?: TraceOptions;
telemetry: AnalyticsServiceSetup;
}
export type AgentExecutorResponse = Promise<ResponseBody>;

2
x-pack/plugins/elastic_assistant/server/lib/langchain/llm/actions_client_llm.test.ts
View file

@ -51,7 +51,7 @@ const mockRequest: KibanaRequest<unknown, unknown, RequestBody> = {
},
subAction: 'invokeAI',
},
assistantLangChain: true,
isEnabledKnowledgeBase: true,
},
} as KibanaRequest<unknown, unknown, RequestBody>;

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

@ -0,0 +1,127 @@
/*
* 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 { EventTypeOpts } from '@kbn/analytics-client';
export const KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT: EventTypeOpts<{
model: string;
resourceAccessed: string;
resultCount: number;
responseTime: number;
}> = {
eventType: 'knowledge_base_execution_success',
schema: {
model: {
type: 'keyword',
_meta: {
description: 'ELSER model used to execute the knowledge base query',
},
},
resourceAccessed: {
type: 'keyword',
_meta: {
description: 'Which knowledge base resource was accessed',
},
},
resultCount: {
type: 'long',
_meta: {
description: 'Number of documents returned from Elasticsearch',
},
},
responseTime: {
type: 'long',
_meta: {
description: `How long it took for Elasticsearch to respond to the knowledge base query`,
},
},
},
};
export const KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT: EventTypeOpts<{
model: string;
resourceAccessed: string;
errorMessage: string;
}> = {
eventType: 'knowledge_base_execution_error',
schema: {
model: {
type: 'keyword',
_meta: {
description: 'ELSER model used to execute the knowledge base query',
},
},
resourceAccessed: {
type: 'keyword',
_meta: {
description: 'Which knowledge base resource was accessed',
},
},
errorMessage: {
type: 'keyword',
_meta: {
description: 'Error message from Elasticsearch',
},
},
},
};
export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{
isEnabledKnowledgeBase: boolean;
isEnabledRAGAlerts: boolean;
}> = {
eventType: 'invoke_assistant_success',
schema: {
isEnabledKnowledgeBase: {
type: 'boolean',
_meta: {
description: 'Is Knowledge Base enabled',
},
},
isEnabledRAGAlerts: {
type: 'boolean',
_meta: {
description: 'Is RAG Alerts enabled',
},
},
},
};
export const INVOKE_ASSISTANT_ERROR_EVENT: EventTypeOpts<{
errorMessage: string;
isEnabledKnowledgeBase: boolean;
isEnabledRAGAlerts: boolean;
}> = {
eventType: 'invoke_assistant_error',
schema: {
errorMessage: {
type: 'keyword',
_meta: {
description: 'Error message from Elasticsearch',
},
},
isEnabledKnowledgeBase: {
type: 'boolean',
_meta: {
description: 'Is Knowledge Base enabled',
},
},
isEnabledRAGAlerts: {
type: 'boolean',
_meta: {
description: 'Is RAG Alerts enabled',
},
},
},
};
export const events: Array<EventTypeOpts<{ [key: string]: unknown }>> = [
KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT,
KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT,
INVOKE_ASSISTANT_SUCCESS_EVENT,
INVOKE_ASSISTANT_ERROR_EVENT,
];

7
x-pack/plugins/elastic_assistant/server/plugin.ts
View file

@ -14,9 +14,11 @@ import {
IContextProvider,
KibanaRequest,
SavedObjectsClientContract,
type AnalyticsServiceSetup,
} from '@kbn/core/server';
import { once } from 'lodash';
import { events } from './lib/telemetry/event_based_telemetry';
import {
AssistantTool,
ElasticAssistantPluginSetup,
@ -40,6 +42,7 @@ interface CreateRouteHandlerContextParams {
core: CoreSetup<ElasticAssistantPluginStart, unknown>;
logger: Logger;
getRegisteredTools: GetRegisteredTools;
telemetry: AnalyticsServiceSetup;
}
export class ElasticAssistantPlugin
@ -61,6 +64,7 @@ export class ElasticAssistantPlugin
core,
logger,
getRegisteredTools,
telemetry,
}: CreateRouteHandlerContextParams): IContextProvider<
ElasticAssistantRequestHandlerContext,
typeof PLUGIN_ID
@ -72,6 +76,7 @@ export class ElasticAssistantPlugin
actions: pluginsStart.actions,
getRegisteredTools,
logger,
telemetry,
};
};
};
@ -87,8 +92,10 @@ export class ElasticAssistantPlugin
getRegisteredTools: (pluginName: string) => {
return appContextService.getRegisteredTools(pluginName);
},
telemetry: core.analytics,
})
);
events.forEach((eventConfig) => core.analytics.registerEventType(eventConfig));
const getElserId: GetElser = once(
async (request: KibanaRequest, savedObjectsClient: SavedObjectsClientContract) => {

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

@ -5,7 +5,7 @@
* 2.0.
*/
import { IRouter, KibanaRequest, Logger } from '@kbn/core/server';
import { IRouter, KibanaRequest } from '@kbn/core/server';
import { transformError } from '@kbn/securitysolution-es-utils';
import { v4 as uuidv4 } from 'uuid';
@ -53,9 +53,11 @@ export const postEvaluateRoute = (
query: buildRouteValidation(PostEvaluatePathQuery),
},
},
// TODO: Limit route based on experimental feature
async (context, request, response) => {
// TODO: Limit route based on experimental feature
const logger: Logger = (await context.elasticAssistant).logger;
const assistantContext = await context.elasticAssistant;
const logger = assistantContext.logger;
const telemetry = assistantContext.telemetry;
try {
const evaluationId = uuidv4();
const {
@ -112,7 +114,8 @@ export const postEvaluateRoute = (
// Default ELSER model
const elserId = await getElser(request, (await context.core).savedObjects.getClient());
// Skeleton request to satisfy `subActionParams` spread in `ActionsClientLlm`
// Skeleton request from route to pass to the agents
// params will be passed to the actions executor
const skeletonRequest: KibanaRequest<unknown, unknown, RequestBody> = {
...request,
body: {
@ -127,7 +130,8 @@ export const postEvaluateRoute = (
},
replacements: {},
size: DEFAULT_SIZE,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
},
};
@ -146,7 +150,7 @@ export const postEvaluateRoute = (
agentEvaluator: (langChainMessages, exampleId) =>
AGENT_EXECUTOR_MAP[agentName]({
actions,
assistantLangChain: true,
isEnabledKnowledgeBase: true,
assistantTools,
connectorId,
esClient,
@ -156,6 +160,7 @@ export const postEvaluateRoute = (
logger,
request: skeletonRequest,
kbResource: ESQL_RESOURCE,
telemetry,
traceOptions: {
exampleId,
projectName,

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

@ -38,7 +38,9 @@ export const deleteKnowledgeBaseRoute = (
},
async (context, request, response) => {
const resp = buildResponse(response);
const logger = (await context.elasticAssistant).logger;
const assistantContext = await context.elasticAssistant;
const logger = assistantContext.logger;
const telemetry = assistantContext.telemetry;
try {
const kbResource =
@ -46,7 +48,12 @@ export const deleteKnowledgeBaseRoute = (
// Get a scoped esClient for deleting the Knowledge Base index, pipeline, and documents
const esClient = (await context.core).elasticsearch.client.asCurrentUser;
const esStore = new ElasticsearchStore(esClient, KNOWLEDGE_BASE_INDEX_PATTERN, logger);
const esStore = new ElasticsearchStore(
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
logger,
telemetry
);
if (kbResource === ESQL_RESOURCE) {
// For now, tearing down the Knowledge Base is fine, but will want to support removing specific assets based

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

@ -41,7 +41,9 @@ export const getKnowledgeBaseStatusRoute = (
},
async (context, request, response) => {
const resp = buildResponse(response);
const logger = (await context.elasticAssistant).logger;
const assistantContext = await context.elasticAssistant;
const logger = assistantContext.logger;
const telemetry = assistantContext.telemetry;
try {
// Get a scoped esClient for finding the status of the Knowledge Base index, pipeline, and documents
@ -52,6 +54,7 @@ export const getKnowledgeBaseStatusRoute = (
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
logger,
telemetry,
elserId,
kbResource
);

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

@ -40,17 +40,21 @@ export const postKnowledgeBaseRoute = (
},
async (context, request, response) => {
const resp = buildResponse(response);
const logger = (await context.elasticAssistant).logger;
const assistantContext = await context.elasticAssistant;
const logger = assistantContext.logger;
const telemetry = assistantContext.telemetry;
try {
const core = await context.core;
// Get a scoped esClient for creating the Knowledge Base index, pipeline, and documents
const esClient = (await context.core).elasticsearch.client.asCurrentUser;
const elserId = await getElser(request, (await context.core).savedObjects.getClient());
const esClient = core.elasticsearch.client.asCurrentUser;
const elserId = await getElser(request, core.savedObjects.getClient());
const kbResource = getKbResource(request);
const esStore = new ElasticsearchStore(
esClient,
KNOWLEDGE_BASE_INDEX_PATTERN,
logger,
telemetry,
elserId,
kbResource
);

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

@ -15,16 +15,26 @@ 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';
jest.mock('../lib/build_response', () => ({
buildResponse: jest.fn().mockImplementation((x) => x),
}));
jest.mock('../lib/executor', () => ({
executeAction: jest.fn().mockImplementation((x) => ({
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
})),
executeAction: jest.fn().mockImplementation(async ({ connectorId }) => {
if (connectorId === 'mock-connector-id') {
return {
connector_id: 'mock-connector-id',
data: mockActionResponse,
status: 'ok',
};
} else {
throw new Error('simulated error');
}
}),
}));
jest.mock('../lib/langchain/execute_custom_llm_chain', () => ({
@ -53,11 +63,13 @@ jest.mock('../lib/langchain/execute_custom_llm_chain', () => ({
),
}));
const reportEvent = jest.fn();
const mockContext = {
elasticAssistant: {
actions: jest.fn(),
getRegisteredTools: jest.fn(() => []),
logger: loggingSystemMock.createLogger(),
telemetry: { ...coreMock.createSetup().analytics, reportEvent },
},
core: {
elasticsearch: {
@ -90,7 +102,8 @@ const mockRequest = {
},
subAction: 'invokeAI',
},
assistantLangChain: true,
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: false,
},
};
@ -106,7 +119,7 @@ describe('postActionsConnectorExecuteRoute', () => {
jest.clearAllMocks();
});
it('returns the expected response when assistantLangChain=false', async () => {
it('returns the expected response when isEnabledKnowledgeBase=false', async () => {
const mockRouter = {
post: jest.fn().mockImplementation(async (_, handler) => {
const result = await handler(
@ -115,7 +128,7 @@ describe('postActionsConnectorExecuteRoute', () => {
...mockRequest,
body: {
...mockRequest.body,
assistantLangChain: false,
isEnabledKnowledgeBase: false,
},
},
mockResponse
@ -137,7 +150,7 @@ describe('postActionsConnectorExecuteRoute', () => {
);
});
it('returns the expected response when assistantLangChain=true', async () => {
it('returns the expected response when isEnabledKnowledgeBase=true', async () => {
const mockRouter = {
post: jest.fn().mockImplementation(async (_, handler) => {
const result = await handler(mockContext, mockRequest, mockResponse);
@ -181,4 +194,219 @@ describe('postActionsConnectorExecuteRoute', () => {
mockGetElser
);
});
it('reports success events to telemetry - kb on, RAG alerts off', async () => {
const mockRouter = {
post: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, mockRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: 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,
allow: ['@timestamp'],
allowReplacement: ['host.name'],
replacements: {},
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
post: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, ragRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
});
}),
};
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,
allow: ['@timestamp'],
allowReplacement: ['host.name'],
replacements: {},
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
post: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, req, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
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 = {
post: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, req, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
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 = {
post: 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,
});
}),
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
it('reports error events to telemetry - kb on, RAG alerts on', async () => {
const badRequest = {
...mockRequest,
params: { connectorId: 'bad-connector-id' },
body: {
...mockRequest.body,
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
post: jest.fn().mockImplementation(async (_, handler) => {
await handler(mockContext, badRequest, mockResponse);
expect(reportEvent).toHaveBeenCalledWith(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
errorMessage: 'simulated error',
isEnabledKnowledgeBase: true,
isEnabledRAGAlerts: true,
});
}),
};
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,
allow: ['@timestamp'],
allowReplacement: ['host.name'],
replacements: {},
isEnabledRAGAlerts: true,
},
};
const mockRouter = {
post: 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,
});
}),
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
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 = {
post: 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,
});
}),
};
await postActionsConnectorExecuteRoute(
mockRouter as unknown as IRouter<ElasticAssistantRequestHandlerContext>,
mockGetElser
);
});
});

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

@ -8,6 +8,10 @@
import { IRouter, Logger } from '@kbn/core/server';
import { transformError } from '@kbn/securitysolution-es-utils';
import {
INVOKE_ASSISTANT_ERROR_EVENT,
INVOKE_ASSISTANT_SUCCESS_EVENT,
} from '../lib/telemetry/event_based_telemetry';
import { executeAction } from '../lib/executor';
import { POST_ACTIONS_CONNECTOR_EXECUTE } from '../../common/constants';
import {
@ -39,7 +43,9 @@ export const postActionsConnectorExecuteRoute = (
},
async (context, request, response) => {
const resp = buildResponse(response);
const logger: Logger = (await context.elasticAssistant).logger;
const assistantContext = await context.elasticAssistant;
const logger: Logger = assistantContext.logger;
const telemetry = assistantContext.telemetry;
try {
const connectorId = decodeURIComponent(request.params.connectorId);
@ -48,16 +54,25 @@ export const postActionsConnectorExecuteRoute = (
const actions = (await context.elasticAssistant).actions;
// if not langchain, call execute action directly and return the response:
if (!request.body.assistantLangChain && !requestHasRequiredAnonymizationParams(request)) {
if (
!request.body.isEnabledKnowledgeBase &&
!requestHasRequiredAnonymizationParams(request)
) {
logger.debug('Executing via actions framework directly');
const result = await executeAction({ actions, request, connectorId });
telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase,
isEnabledRAGAlerts: request.body.isEnabledRAGAlerts,
});
return response.ok({
body: result,
});
}
// TODO: Add `traceId` to actions request when calling via langchain
logger.debug('Executing via langchain, assistantLangChain: true');
logger.debug(
`Executing via langchain, isEnabledKnowledgeBase: ${request.body.isEnabledKnowledgeBase}, isEnabledRAGAlerts: ${request.body.isEnabledRAGAlerts}`
);
// Fetch any tools registered by the request's originating plugin
const pluginName = getPluginNameFromRequest({
@ -87,7 +102,7 @@ export const postActionsConnectorExecuteRoute = (
allow: request.body.allow,
allowReplacement: request.body.allowReplacement,
actions,
assistantLangChain: request.body.assistantLangChain,
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase,
assistantTools,
connectorId,
elserId,
@ -99,8 +114,13 @@ export const postActionsConnectorExecuteRoute = (
request,
replacements: request.body.replacements,
size: request.body.size,
telemetry,
});
telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase,
isEnabledRAGAlerts: request.body.isEnabledRAGAlerts,
});
return response.ok({
body: {
...langChainResponseBody,
@ -110,6 +130,11 @@ export const postActionsConnectorExecuteRoute = (
} catch (err) {
logger.error(err);
const error = transformError(err);
telemetry.reportEvent(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase,
isEnabledRAGAlerts: request.body.isEnabledRAGAlerts,
errorMessage: error.message,
});
return resp.error({
body: error.message,

3
x-pack/plugins/elastic_assistant/server/schemas/post_actions_connector_execute.ts
View file

@ -37,7 +37,8 @@ export const PostActionsConnectorExecuteBody = t.type({
alertsIndexPattern: t.union([t.string, t.undefined]),
allow: t.union([t.array(t.string), t.undefined]),
allowReplacement: t.union([t.array(t.string), t.undefined]),
assistantLangChain: t.boolean,
isEnabledKnowledgeBase: t.boolean,
isEnabledRAGAlerts: t.boolean,
replacements: t.union([t.record(t.string, t.string), t.undefined]),
size: t.union([t.number, t.undefined]),
});

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

@ -10,6 +10,7 @@ import type {
PluginStartContract as ActionsPluginStart,
} from '@kbn/actions-plugin/server';
import type {
AnalyticsServiceSetup,
CustomRequestHandlerContext,
KibanaRequest,
Logger,
@ -57,6 +58,7 @@ export interface ElasticAssistantApiRequestHandlerContext {
actions: ActionsPluginStart;
getRegisteredTools: GetRegisteredTools;
logger: Logger;
telemetry: AnalyticsServiceSetup;
}
/**
@ -88,7 +90,7 @@ export interface AssistantToolParams {
alertsIndexPattern?: string;
allow?: string[];
allowReplacement?: string[];
assistantLangChain: boolean;
isEnabledKnowledgeBase: boolean;
chain: RetrievalQAChain;
esClient: ElasticsearchClient;
modelExists: boolean;

2
x-pack/plugins/elastic_assistant/tsconfig.json
View file

@ -13,6 +13,7 @@
"../../../typings/**/*"
],
"kbn_references": [
"@kbn/analytics-client",
"@kbn/core",
"@kbn/core-http-server",
"@kbn/licensing-plugin",
@ -32,6 +33,7 @@
"@kbn/stack-connectors-plugin",
"@kbn/ml-plugin",
"@kbn/apm-utils",
"@kbn/core-analytics-server",
],
"exclude": [
"target/**/*",

1
x-pack/plugins/security_solution/public/assistant/use_assistant_telemetry/index.test.tsx
View file

@ -27,6 +27,7 @@ const mockedTelemetry = {
reportAssistantInvoked,
reportAssistantMessageSent,
reportAssistantQuickPrompt,
reportAssistantSettingToggled: () => {},
};
jest.mock('../use_conversation_store', () => {

1
x-pack/plugins/security_solution/public/assistant/use_assistant_telemetry/index.tsx
View file

@ -43,5 +43,6 @@ export const useAssistantTelemetry = (): AssistantTelemetry => {
reportTelemetry({ fn: telemetry.reportAssistantMessageSent, params }),
reportAssistantQuickPrompt: (params) =>
reportTelemetry({ fn: telemetry.reportAssistantQuickPrompt, params }),
reportAssistantSettingToggled: (params) => telemetry.reportAssistantSettingToggled(params),
};
};

1
x-pack/plugins/security_solution/public/common/lib/telemetry/constants.ts
View file

@ -44,6 +44,7 @@ export enum TelemetryEventTypes {
AssistantInvoked = 'Assistant Invoked',
AssistantMessageSent = 'Assistant Message Sent',
AssistantQuickPrompt = 'Assistant Quick Prompt',
AssistantSettingToggled = 'Assistant Setting Toggled',
EntityDetailsClicked = 'Entity Details Clicked',
EntityAlertsClicked = 'Entity Alerts Clicked',
EntityRiskFiltered = 'Entity Risk Filtered',

34
x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/index.ts
View file

@ -45,6 +45,20 @@ export const assistantMessageSentEvent: TelemetryEvent = {
optional: false,
},
},
isEnabledKnowledgeBase: {
type: 'boolean',
_meta: {
description: 'Is knowledge base enabled',
optional: false,
},
},
isEnabledRAGAlerts: {
type: 'boolean',
_meta: {
description: 'Is RAG on Alerts enabled',
optional: false,
},
},
},
};
@ -67,3 +81,23 @@ export const assistantQuickPrompt: TelemetryEvent = {
},
},
};
export const assistantSettingToggledEvent: TelemetryEvent = {
eventType: TelemetryEventTypes.AssistantSettingToggled,
schema: {
isEnabledKnowledgeBase: {
type: 'boolean',
_meta: {
description: 'Is knowledge base enabled',
optional: true,
},
},
isEnabledRAGAlerts: {
type: 'boolean',
_meta: {
description: 'Is RAG on Alerts enabled',
optional: true,
},
},
},
};

12
x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/types.ts
View file

@ -16,6 +16,8 @@ export interface ReportAssistantInvokedParams {
export interface ReportAssistantMessageSentParams {
conversationId: string;
role: string;
isEnabledKnowledgeBase: boolean;
isEnabledRAGAlerts: boolean;
}
export interface ReportAssistantQuickPromptParams {
@ -23,9 +25,15 @@ export interface ReportAssistantQuickPromptParams {
promptTitle: string;
}
export interface ReportAssistantSettingToggledParams {
isEnabledKnowledgeBase?: boolean;
isEnabledRAGAlerts?: boolean;
}
export type ReportAssistantTelemetryEventParams =
| ReportAssistantInvokedParams
| ReportAssistantMessageSentParams
| ReportAssistantSettingToggledParams
| ReportAssistantQuickPromptParams;
export type AssistantTelemetryEvent =
@ -33,6 +41,10 @@ export type AssistantTelemetryEvent =
eventType: TelemetryEventTypes.AssistantInvoked;
schema: RootSchema<ReportAssistantInvokedParams>;
}
| {
eventType: TelemetryEventTypes.AssistantSettingToggled;
schema: RootSchema<ReportAssistantSettingToggledParams>;
}
| {
eventType: TelemetryEventTypes.AssistantMessageSent;
schema: RootSchema<ReportAssistantMessageSentParams>;

2
x-pack/plugins/security_solution/public/common/lib/telemetry/events/telemetry_events.ts
View file

@ -18,6 +18,7 @@ import {
} from './entity_analytics';
import {
assistantInvokedEvent,
assistantSettingToggledEvent,
assistantMessageSentEvent,
assistantQuickPrompt,
} from './ai_assistant';
@ -138,6 +139,7 @@ export const telemetryEvents = [
assistantInvokedEvent,
assistantMessageSentEvent,
assistantQuickPrompt,
assistantSettingToggledEvent,
entityClickedEvent,
entityAlertsClickedEvent,
entityRiskFilteredEvent,

1
x-pack/plugins/security_solution/public/common/lib/telemetry/telemetry_client.mock.ts
View file

@ -14,6 +14,7 @@ export const createTelemetryClientMock = (): jest.Mocked<TelemetryClientStart> =
reportAssistantInvoked: jest.fn(),
reportAssistantMessageSent: jest.fn(),
reportAssistantQuickPrompt: jest.fn(),
reportAssistantSettingToggled: jest.fn(),
reportEntityDetailsClicked: jest.fn(),
reportEntityAlertsClicked: jest.fn(),
reportEntityRiskFiltered: jest.fn(),

9
x-pack/plugins/security_solution/public/common/lib/telemetry/telemetry_client.ts
View file

@ -23,6 +23,7 @@ import type {
ReportAssistantInvokedParams,
ReportAssistantMessageSentParams,
ReportAssistantQuickPromptParams,
ReportAssistantSettingToggledParams,
} from './types';
import { TelemetryEventTypes } from './constants';
@ -54,10 +55,14 @@ export class TelemetryClient implements TelemetryClientStart {
public reportAssistantMessageSent = ({
conversationId,
isEnabledKnowledgeBase,
isEnabledRAGAlerts,
role,
}: ReportAssistantMessageSentParams) => {
this.analytics.reportEvent(TelemetryEventTypes.AssistantMessageSent, {
conversationId,
isEnabledKnowledgeBase,
isEnabledRAGAlerts,
role,
});
};
@ -72,6 +77,10 @@ export class TelemetryClient implements TelemetryClientStart {
});
};
public reportAssistantSettingToggled = (params: ReportAssistantSettingToggledParams) => {
this.analytics.reportEvent(TelemetryEventTypes.AssistantSettingToggled, params);
};
public reportEntityDetailsClicked = ({ entity }: ReportEntityDetailsClickedParams) => {
this.analytics.reportEvent(TelemetryEventTypes.EntityDetailsClicked, {
entity,

2
x-pack/plugins/security_solution/public/common/lib/telemetry/types.ts
View file

@ -34,6 +34,7 @@ import type {
ReportAssistantInvokedParams,
ReportAssistantQuickPromptParams,
ReportAssistantMessageSentParams,
ReportAssistantSettingToggledParams,
} from './events/ai_assistant/types';
export * from './events/ai_assistant/types';
@ -92,6 +93,7 @@ export interface TelemetryClientStart {
reportAssistantInvoked(params: ReportAssistantInvokedParams): void;
reportAssistantMessageSent(params: ReportAssistantMessageSentParams): void;
reportAssistantQuickPrompt(params: ReportAssistantQuickPromptParams): void;
reportAssistantSettingToggled(params: ReportAssistantSettingToggledParams): void;
reportEntityDetailsClicked(params: ReportEntityDetailsClickedParams): void;
reportEntityAlertsClicked(params: ReportEntityAlertsClickedParams): void;

8
x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts
View file

@ -22,7 +22,7 @@ describe('AlertCountsTool', () => {
const replacements = { key: 'value' };
const request = {
body: {
assistantLangChain: false,
isEnabledKnowledgeBase: false,
alertsIndexPattern: '.alerts-security.alerts-default',
allow: ['@timestamp', 'cloud.availability_zone', 'user.name'],
allowReplacement: ['user.name'],
@ -30,11 +30,11 @@ describe('AlertCountsTool', () => {
size: 20,
},
} as unknown as KibanaRequest<unknown, unknown, RequestBody>;
const assistantLangChain = true;
const isEnabledKnowledgeBase = true;
const chain = {} as unknown as RetrievalQAChain;
const modelExists = true;
const rest = {
assistantLangChain,
isEnabledKnowledgeBase,
chain,
modelExists,
};
@ -57,7 +57,7 @@ describe('AlertCountsTool', () => {
it('returns false when the request is missing required anonymization parameters', () => {
const requestMissingAnonymizationParams = {
body: {
assistantLangChain: false,
isEnabledKnowledgeBase: false,
alertsIndexPattern: '.alerts-security.alerts-default',
size: 20,
},

24
x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts
View file

@ -19,7 +19,7 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
} as unknown as ElasticsearchClient;
const request = {
body: {
assistantLangChain: false,
isEnabledKnowledgeBase: false,
alertsIndexPattern: '.alerts-security.alerts-default',
allow: ['@timestamp', 'cloud.availability_zone', 'user.name'],
allowReplacement: ['user.name'],
@ -34,9 +34,9 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
};
describe('isSupported', () => {
it('returns false if assistantLangChain is false', () => {
it('returns false if isEnabledKnowledgeBase is false', () => {
const params = {
assistantLangChain: false,
isEnabledKnowledgeBase: false,
modelExists: true,
...rest,
};
@ -46,7 +46,7 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
it('returns false if modelExists is false (the ELSER model is not installed)', () => {
const params = {
assistantLangChain: true,
isEnabledKnowledgeBase: true,
modelExists: false, // <-- ELSER model is not installed
...rest,
};
@ -54,9 +54,9 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
expect(ESQL_KNOWLEDGE_BASE_TOOL.isSupported(params)).toBe(false);
});
it('returns true if assistantLangChain and modelExists are true', () => {
it('returns true if isEnabledKnowledgeBase and modelExists are true', () => {
const params = {
assistantLangChain: true,
isEnabledKnowledgeBase: true,
modelExists: true,
...rest,
};
@ -66,9 +66,9 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
});
describe('getTool', () => {
it('returns null if assistantLangChain is false', () => {
it('returns null if isEnabledKnowledgeBase is false', () => {
const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({
assistantLangChain: false,
isEnabledKnowledgeBase: false,
modelExists: true,
...rest,
});
@ -78,7 +78,7 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
it('returns null if modelExists is false (the ELSER model is not installed)', () => {
const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({
assistantLangChain: true,
isEnabledKnowledgeBase: true,
modelExists: false, // <-- ELSER model is not installed
...rest,
});
@ -86,9 +86,9 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
expect(tool).toBeNull();
});
it('should return a Tool instance if assistantLangChain and modelExists are true', () => {
it('should return a Tool instance if isEnabledKnowledgeBase and modelExists are true', () => {
const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({
assistantLangChain: true,
isEnabledKnowledgeBase: true,
modelExists: true,
...rest,
});
@ -98,7 +98,7 @@ describe('EsqlLanguageKnowledgeBaseTool', () => {
it('should return a tool with the expected tags', () => {
const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({
assistantLangChain: true,
isEnabledKnowledgeBase: true,
modelExists: true,
...rest,
}) as DynamicTool;

4
x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts
View file

@ -18,8 +18,8 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = {
'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language.',
sourceRegister: APP_UI_ID,
isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => {
const { assistantLangChain, modelExists } = params;
return assistantLangChain && modelExists;
const { isEnabledKnowledgeBase, modelExists } = params;
return isEnabledKnowledgeBase && modelExists;
},
getTool(params: AssistantToolParams) {
if (!this.isSupported(params)) return null;

8
x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts
View file

@ -24,7 +24,7 @@ describe('OpenAndAcknowledgedAlertsTool', () => {
const replacements = { key: 'value' };
const request = {
body: {
assistantLangChain: false,
isEnabledKnowledgeBase: false,
alertsIndexPattern: '.alerts-security.alerts-default',
allow: ['@timestamp', 'cloud.availability_zone', 'user.name'],
allowReplacement: ['user.name'],
@ -32,11 +32,11 @@ describe('OpenAndAcknowledgedAlertsTool', () => {
size: 20,
},
} as unknown as KibanaRequest<unknown, unknown, RequestBody>;
const assistantLangChain = true;
const isEnabledKnowledgeBase = true;
const chain = {} as unknown as RetrievalQAChain;
const modelExists = true;
const rest = {
assistantLangChain,
isEnabledKnowledgeBase,
esClient,
chain,
modelExists,
@ -60,7 +60,7 @@ describe('OpenAndAcknowledgedAlertsTool', () => {
it('returns false when the request is missing required anonymization parameters', () => {
const requestMissingAnonymizationParams = {
body: {
assistantLangChain: false,
isEnabledKnowledgeBase: false,
alertsIndexPattern: '.alerts-security.alerts-default',
size: 20,
},

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

@ -436,7 +436,8 @@ export default function bedrockTest({ getService }: FtrProviderContext) {
],
},
},
assistantLangChain: false,
isEnabledKnowledgeBase: false,
isEnabledRAGAlerts: false,
})
.pipe(passThrough);
const responseBuffer: Uint8Array[] = [];