From 1a42e3b33e1e2ec439e0b7fb02a4d1b48e38844f Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 16 Oct 2024 17:32:47 +0200 Subject: [PATCH 01/41] [Security Assistant] Fix KB output fields --- .../index_entry_editor.tsx | 109 ++++++++++++++---- 1 file changed, 84 insertions(+), 25 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx index 7475ea55ca5fc..550861bcbffd9 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx @@ -17,7 +17,7 @@ import { EuiSuperSelect, } from '@elastic/eui'; import useAsync from 'react-use/lib/useAsync'; -import React, { useCallback } from 'react'; +import React, { useCallback, useMemo } from 'react'; import { IndexEntry } from '@kbn/elastic-assistant-common'; import { DataViewsContract } from '@kbn/data-views-plugin/public'; import * as i18n from './translations'; @@ -96,29 +96,37 @@ export const IndexEntryEditor: React.FC = React.memo( })); }, [dataViews]); - const fieldOptions = useAsync(async () => { - const fields = await dataViews.getFieldsForWildcard({ - pattern: entry?.index ?? '', - fieldTypes: ['semantic_text'], - }); + const indexFields = useAsync( + async () => + dataViews.getFieldsForWildcard({ + pattern: entry?.index ?? '', + }), + [] + ); - return fields - .filter((field) => field.esTypes?.includes('semantic_text')) - .map((field) => ({ + const fieldOptions = useMemo( + () => + indexFields?.value + ?.filter((field) => field.esTypes?.includes('semantic_text')) + .map((field) => ({ + 'data-test-subj': field.name, + label: field.name, + value: field.name, + })) ?? [], + [indexFields?.value] + ); + + const outputFieldOptions = useMemo( + () => + indexFields?.value?.map((field) => ({ 'data-test-subj': field.name, label: field.name, value: field.name, - })); - }, [entry]); - - const setIndex = useCallback( - async (e: Array>) => { - setEntry((prevEntry) => ({ ...prevEntry, index: e[0]?.value })); - }, - [setEntry] + })) ?? [], + [indexFields?.value] ); - const onCreateOption = (searchValue: string) => { + const onCreateIndexOption = (searchValue: string) => { const normalizedSearchValue = searchValue.trim().toLowerCase(); if (!normalizedSearchValue) { @@ -131,7 +139,6 @@ export const IndexEntryEditor: React.FC = React.memo( }; setIndex([newOption]); - setField([{ label: '', value: '' }]); }; const onCreateFieldOption = (searchValue: string) => { @@ -170,6 +177,52 @@ export const IndexEntryEditor: React.FC = React.memo( [setEntry] ); + // Field + const setOutputFields = useCallback( + async (e: Array>) => { + setEntry((prevEntry) => ({ + ...prevEntry, + outputFields: e + ?.filter((option) => !!option.value) + .map((option) => option.value as string), + })); + }, + [setEntry] + ); + + const setIndex = useCallback( + async (e: Array>) => { + setEntry((prevEntry) => ({ ...prevEntry, index: e[0]?.value })); + setField([]); + setOutputFields([]); + }, + [setEntry, setField, setOutputFields] + ); + + const onCreateOutputFieldsOption = useCallback( + (searchValue: string) => { + const normalizedSearchValue = searchValue.trim().toLowerCase(); + + if (!normalizedSearchValue) { + return; + } + + const newOption: EuiComboBoxOptionOption = { + label: searchValue, + value: searchValue, + }; + + setOutputFields([ + ...(entry?.outputFields?.map((field) => ({ + label: field, + value: field, + })) ?? []), + newOption, + ]); + }, + [entry?.outputFields, setOutputFields] + ); + return ( = React.memo( aria-label={i18n.ENTRY_INDEX_NAME_INPUT_LABEL} isClearable={true} singleSelection={{ asPlainText: true }} - onCreateOption={onCreateOption} + onCreateOption={onCreateIndexOption} fullWidth options={indexOptions.value ?? []} selectedOptions={ @@ -228,7 +281,7 @@ export const IndexEntryEditor: React.FC = React.memo( singleSelection={{ asPlainText: true }} onCreateOption={onCreateFieldOption} fullWidth - options={fieldOptions.value ?? []} + options={fieldOptions} selectedOptions={ entry?.field ? [ @@ -281,11 +334,17 @@ export const IndexEntryEditor: React.FC = React.memo( ({ + label: field, + value: field, + })) ?? [] + } + onChange={setOutputFields} /> From 6bbd0376c2fb626aa5ae39c9acf4960e0adbbda2 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 16 Oct 2024 22:56:10 +0200 Subject: [PATCH 02/41] update tests --- .../index_entry_editor.test.tsx | 52 +++++++++---------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.test.tsx index d4634cdf4c563..e4656b10d1d31 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.test.tsx @@ -42,10 +42,10 @@ describe('IndexEntryEditor', () => { jest.clearAllMocks(); }); - it('renders the form fields with initial values', () => { + it('renders the form fields with initial values', async () => { const { getByDisplayValue } = render(); - waitFor(() => { + await waitFor(() => { expect(getByDisplayValue('Test Entry')).toBeInTheDocument(); expect(getByDisplayValue('Test Description')).toBeInTheDocument(); expect(getByDisplayValue('Test Query Description')).toBeInTheDocument(); @@ -54,35 +54,37 @@ describe('IndexEntryEditor', () => { }); }); - it('updates the name field on change', () => { + it('updates the name field on change', async () => { const { getByTestId } = render(); - waitFor(() => { + await waitFor(() => { const nameInput = getByTestId('entry-name'); fireEvent.change(nameInput, { target: { value: 'New Entry Name' } }); + expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); - - expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); - it('updates the description field on change', () => { + it('updates the description field on change', async () => { const { getByTestId } = render(); - waitFor(() => { + + await waitFor(() => { const descriptionInput = getByTestId('entry-description'); fireEvent.change(descriptionInput, { target: { value: 'New Description' } }); }); - expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); + await waitFor(() => { + expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); + }); }); - it('updates the query description field on change', () => { + it('updates the query description field on change', async () => { const { getByTestId } = render(); - waitFor(() => { + + await waitFor(() => { const queryDescriptionInput = getByTestId('query-description'); fireEvent.change(queryDescriptionInput, { target: { value: 'New Query Description' } }); + expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); - - expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); it('displays sharing options and updates on selection', async () => { @@ -91,8 +93,6 @@ describe('IndexEntryEditor', () => { await waitFor(() => { fireEvent.click(getByTestId('sharing-select')); fireEvent.click(getByTestId('sharing-private-option')); - }); - await waitFor(() => { expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); }); @@ -100,28 +100,25 @@ describe('IndexEntryEditor', () => { it('fetches index options and updates on selection', async () => { const { getAllByTestId, getByTestId } = render(); - await waitFor(() => expect(mockDataViews.getIndices).toHaveBeenCalled()); - await waitFor(() => { + expect(mockDataViews.getIndices).toHaveBeenCalled(); fireEvent.click(getByTestId('index-combobox')); fireEvent.click(getAllByTestId('comboBoxToggleListButton')[0]); + fireEvent.click(getByTestId('index-2')); + expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); - fireEvent.click(getByTestId('index-2')); - - expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); }); it('fetches field options based on selected index and updates on selection', async () => { const { getByTestId, getAllByTestId } = render(); - await waitFor(() => + await waitFor(() => { expect(mockDataViews.getFieldsForWildcard).toHaveBeenCalledWith({ pattern: 'index-1', - fieldTypes: ['semantic_text'], - }) - ); + }); + }); - await waitFor(() => { + await waitFor(async () => { fireEvent.click(getByTestId('index-combobox')); fireEvent.click(getAllByTestId('comboBoxToggleListButton')[0]); }); @@ -135,7 +132,10 @@ describe('IndexEntryEditor', () => { within(getByTestId('entry-combobox')).getByTestId('comboBoxSearchInput'), 'field-3' ); - expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); + + await waitFor(() => { + expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function)); + }); }); it('disables the field combo box if no index is selected', () => { From bbad040ec66501b401f696743df54de1b5484830 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 18 Oct 2024 14:24:34 +0200 Subject: [PATCH 03/41] revert cog wheel --- .../impl/assistant/index.test.tsx | 107 +++++++++++++++ .../settings/assistant_settings.test.tsx | 13 +- .../assistant/settings/assistant_settings.tsx | 124 +++++++++++++++++- .../assistant_settings_button.test.tsx | 8 ++ .../settings/assistant_settings_button.tsx | 51 +++++-- 5 files changed, 289 insertions(+), 14 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx index 08bac25c0a522..1ef2db7b26c03 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx @@ -24,6 +24,7 @@ import { Conversation } from '../assistant_context/types'; import * as all from './chat_send/use_chat_send'; import { useConversation } from './use_conversation'; import { AIConnector } from '../connectorland/connector_selector'; +import { omit } from 'lodash'; jest.mock('../connectorland/use_load_connectors'); jest.mock('../connectorland/connector_setup'); @@ -140,6 +141,112 @@ describe('Assistant', () => { >); }); + describe('persistent storage', () => { + it('should refetchCurrentUserConversations after settings save button click', async () => { + const chatSendSpy = jest.spyOn(all, 'useChatSend'); + await renderAssistant(); + + fireEvent.click(screen.getByTestId('settings')); + + jest.mocked(useFetchCurrentUserConversations).mockReturnValue({ + data: { + ...mockData, + welcome_id: { + ...mockData.welcome_id, + apiConfig: { newProp: true }, + }, + }, + isLoading: false, + refetch: jest.fn().mockResolvedValue({ + isLoading: false, + data: { + ...mockData, + welcome_id: { + ...mockData.welcome_id, + apiConfig: { newProp: true }, + }, + }, + }), + isFetched: true, + } as unknown as DefinedUseQueryResult, unknown>); + + await act(async () => { + fireEvent.click(screen.getByTestId('save-button')); + }); + + expect(chatSendSpy).toHaveBeenLastCalledWith( + expect.objectContaining({ + currentConversation: { + apiConfig: { newProp: true }, + category: 'assistant', + id: mockData.welcome_id.id, + messages: [], + title: 'Welcome', + replacements: {}, + }, + }) + ); + }); + + it('should refetchCurrentUserConversations after settings save button click, but do not update convos when refetch returns bad results', async () => { + jest.mocked(useFetchCurrentUserConversations).mockReturnValue({ + data: mockData, + isLoading: false, + refetch: jest.fn().mockResolvedValue({ + isLoading: false, + data: omit(mockData, 'welcome_id'), + }), + isFetched: true, + } as unknown as DefinedUseQueryResult, unknown>); + const chatSendSpy = jest.spyOn(all, 'useChatSend'); + await renderAssistant(); + + fireEvent.click(screen.getByTestId('settings')); + await act(async () => { + fireEvent.click(screen.getByTestId('save-button')); + }); + + expect(chatSendSpy).toHaveBeenLastCalledWith( + expect.objectContaining({ + currentConversation: { + apiConfig: { connectorId: '123' }, + replacements: {}, + category: 'assistant', + id: mockData.welcome_id.id, + messages: [], + title: 'Welcome', + }, + }) + ); + }); + + it('should delete conversation when delete button is clicked', async () => { + await renderAssistant(); + const deleteButton = screen.getAllByTestId('delete-option')[0]; + await act(async () => { + fireEvent.click(deleteButton); + }); + + await act(async () => { + fireEvent.click(screen.getByTestId('confirmModalConfirmButton')); + }); + + await waitFor(() => { + expect(mockDeleteConvo).toHaveBeenCalledWith(mockData.electric_sheep_id.id); + }); + }); + it('should refetchCurrentUserConversations after clear chat history button click', async () => { + await renderAssistant(); + fireEvent.click(screen.getByTestId('chat-context-menu')); + fireEvent.click(screen.getByTestId('clear-chat')); + fireEvent.click(screen.getByTestId('confirmModalConfirmButton')); + await waitFor(() => { + expect(clearConversation).toHaveBeenCalled(); + expect(refetchResults).toHaveBeenCalled(); + }); + }); + }); + describe('when selected conversation changes and some connectors are loaded', () => { it('should persist the conversation id to local storage', async () => { const getConversation = jest.fn().mockResolvedValue(mockData.electric_sheep_id); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx index 14bfcb4cdbbec..c9f4f07d83b11 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx @@ -38,7 +38,7 @@ const mockContext = { basePromptContexts: MOCK_QUICK_PROMPTS, setSelectedSettingsTab, http: {}, - assistantFeatures: { assistantModelEvaluation: true }, + assistantFeatures: { assistantModelEvaluation: true, assistantKnowledgeBaseByDefault: false }, selectedSettingsTab: 'CONVERSATIONS_TAB', assistantAvailability: { isAssistantEnabled: true, @@ -136,6 +136,17 @@ describe('AssistantSettings', () => { QUICK_PROMPTS_TAB, SYSTEM_PROMPTS_TAB, ])('%s', (tab) => { + it('Opens the tab on button click', () => { + (useAssistantContext as jest.Mock).mockImplementation(() => ({ + ...mockContext, + selectedSettingsTab: tab === CONVERSATIONS_TAB ? ANONYMIZATION_TAB : CONVERSATIONS_TAB, + })); + const { getByTestId } = render(, { + wrapper, + }); + fireEvent.click(getByTestId(`${tab}-button`)); + expect(setSelectedSettingsTab).toHaveBeenCalledWith(tab); + }); it('renders with the correct tab open', () => { (useAssistantContext as jest.Mock).mockImplementation(() => ({ ...mockContext, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx index f325e411bae2b..350780ea5b168 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx @@ -9,10 +9,14 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react'; import { EuiButton, EuiButtonEmpty, + EuiIcon, EuiModal, EuiModalFooter, + EuiKeyPadMenu, + EuiKeyPadMenuItem, EuiPage, EuiPageBody, + EuiPageSidebar, EuiSplitPanel, } from '@elastic/eui'; @@ -76,7 +80,16 @@ export const AssistantSettings: React.FC = React.memo( conversations, conversationsLoaded, }) => { - const { http, toasts, selectedSettingsTab, setSelectedSettingsTab } = useAssistantContext(); + const { + assistantFeatures: { + assistantModelEvaluation: modelEvaluatorEnabled, + assistantKnowledgeBaseByDefault, + }, + http, + toasts, + selectedSettingsTab, + setSelectedSettingsTab, + } = useAssistantContext(); useEffect(() => { if (selectedSettingsTab == null) { @@ -201,6 +214,115 @@ export const AssistantSettings: React.FC = React.memo( return ( + {!assistantKnowledgeBaseByDefault && ( + + + setSelectedSettingsTab(CONVERSATIONS_TAB)} + data-test-subj={`${CONVERSATIONS_TAB}-button`} + > + <> + + + + + setSelectedSettingsTab(QUICK_PROMPTS_TAB)} + data-test-subj={`${QUICK_PROMPTS_TAB}-button`} + > + <> + + + + + setSelectedSettingsTab(SYSTEM_PROMPTS_TAB)} + data-test-subj={`${SYSTEM_PROMPTS_TAB}-button`} + > + + + + setSelectedSettingsTab(ANONYMIZATION_TAB)} + data-test-subj={`${ANONYMIZATION_TAB}-button`} + > + + + setSelectedSettingsTab(KNOWLEDGE_BASE_TAB)} + data-test-subj={`${KNOWLEDGE_BASE_TAB}-button`} + > + + + {modelEvaluatorEnabled && ( + setSelectedSettingsTab(EVALUATION_TAB)} + data-test-subj={`${EVALUATION_TAB}-button`} + > + + + )} + + + )} + { jest.clearAllMocks(); }); + it('Clicking the settings gear opens the conversations tab', () => { + const { getByTestId } = render(); + fireEvent.click(getByTestId('settings')); + expect(setSelectedSettingsTab).toHaveBeenCalledWith(CONVERSATIONS_TAB); + expect(setIsSettingsModalVisible).toHaveBeenCalledWith(true); + }); + it('Settings modal is visible and calls correct actions per click', () => { const { getByTestId } = render( diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.tsx index 40bf1e740ab60..3d6544643ba3e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.tsx @@ -6,6 +6,7 @@ */ import React, { useCallback } from 'react'; +import { EuiButtonIcon, EuiToolTip } from '@elastic/eui'; import { QueryObserverResult, RefetchOptions, RefetchQueryFilters } from '@tanstack/react-query'; import { DataStreamApis } from '../use_data_stream_apis'; import { AIConnector } from '../../connectorland/connector_selector'; @@ -13,6 +14,7 @@ import { Conversation } from '../../..'; import { AssistantSettings } from './assistant_settings'; import * as i18n from './translations'; import { useAssistantContext } from '../../assistant_context'; +import { CONVERSATIONS_TAB } from './const'; interface Props { defaultConnector?: AIConnector; @@ -45,7 +47,11 @@ export const AssistantSettingsButton: React.FC = React.memo( refetchCurrentUserConversations, refetchPrompts, }) => { - const { toasts } = useAssistantContext(); + const { + assistantFeatures: { assistantKnowledgeBaseByDefault }, + toasts, + setSelectedSettingsTab, + } = useAssistantContext(); // Modal control functions const cleanupAndCloseModal = useCallback(() => { @@ -73,18 +79,39 @@ export const AssistantSettingsButton: React.FC = React.memo( [cleanupAndCloseModal, refetchCurrentUserConversations, refetchPrompts, toasts] ); + const handleShowConversationSettings = useCallback(() => { + setSelectedSettingsTab(CONVERSATIONS_TAB); + setIsSettingsModalVisible(true); + }, [setIsSettingsModalVisible, setSelectedSettingsTab]); + return ( - isSettingsModalVisible && ( - - ) + <> + {!assistantKnowledgeBaseByDefault && ( + + + + )} + + {isSettingsModalVisible && ( + + )} + ); } ); From bfe06936a6555082f73d4ebc22c329e5d3b5fdc2 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 18 Oct 2024 16:41:43 +0200 Subject: [PATCH 04/41] fix nav --- .../features/src/assistant/kibana_features.ts | 6 ++++++ .../packages/security-solution/navigation/src/constants.ts | 1 + 2 files changed, 7 insertions(+) diff --git a/x-pack/packages/security-solution/features/src/assistant/kibana_features.ts b/x-pack/packages/security-solution/features/src/assistant/kibana_features.ts index e7bafd5595316..81cf7d18af129 100644 --- a/x-pack/packages/security-solution/features/src/assistant/kibana_features.ts +++ b/x-pack/packages/security-solution/features/src/assistant/kibana_features.ts @@ -26,6 +26,9 @@ export const getAssistantBaseKibanaFeature = (): BaseKibanaFeatureConfig => ({ app: [ASSISTANT_FEATURE_ID, 'kibana'], catalogue: [APP_ID], minimumLicense: 'enterprise', + management: { + kibana: ['securityAiAssistantManagement'], + }, privileges: { all: { api: ['elasticAssistant'], @@ -36,6 +39,9 @@ export const getAssistantBaseKibanaFeature = (): BaseKibanaFeatureConfig => ({ read: [], }, ui: [], + management: { + kibana: ['securityAiAssistantManagement'], + }, }, read: { // No read-only mode currently supported diff --git a/x-pack/packages/security-solution/navigation/src/constants.ts b/x-pack/packages/security-solution/navigation/src/constants.ts index 249b834f929bd..4d045f8e2fcf6 100644 --- a/x-pack/packages/security-solution/navigation/src/constants.ts +++ b/x-pack/packages/security-solution/navigation/src/constants.ts @@ -94,6 +94,7 @@ export enum ExternalPageName { managementFiles = 'management:filesManagement', managementSpaces = 'management:spaces', managementSettings = 'management:settings', + managementSecurityAiAssistantManagement = 'management:securityAiAssistantManagement', // Cloud UI // These are links to Cloud UI outside Kibana // Special Format: : From 0098227aa294492ddb28f302bd54c004e51f9d18 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 18 Oct 2024 16:57:49 +0200 Subject: [PATCH 05/41] fix tests --- .../impl/assistant/settings/assistant_settings_button.test.tsx | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.test.tsx index 027eb1c20b80f..6a0909d344968 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_button.test.tsx @@ -32,6 +32,7 @@ const testProps = { const setSelectedSettingsTab = jest.fn(); const mockUseAssistantContext = { setSelectedSettingsTab, + assistantFeatures: {}, }; jest.mock('../../assistant_context', () => { const original = jest.requireActual('../../assistant_context'); From 7ff4a2e02773a746d9f632f1ef27e89c1dbacb55 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 18 Oct 2024 17:03:01 +0200 Subject: [PATCH 06/41] cleanup --- x-pack/packages/security-solution/navigation/src/constants.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/packages/security-solution/navigation/src/constants.ts b/x-pack/packages/security-solution/navigation/src/constants.ts index 4d045f8e2fcf6..249b834f929bd 100644 --- a/x-pack/packages/security-solution/navigation/src/constants.ts +++ b/x-pack/packages/security-solution/navigation/src/constants.ts @@ -94,7 +94,6 @@ export enum ExternalPageName { managementFiles = 'management:filesManagement', managementSpaces = 'management:spaces', managementSettings = 'management:settings', - managementSecurityAiAssistantManagement = 'management:securityAiAssistantManagement', // Cloud UI // These are links to Cloud UI outside Kibana // Special Format: : From feee0ce9763c6a25291569a0ae9dafc1836e04cb Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 21 Oct 2024 10:53:30 +0200 Subject: [PATCH 07/41] Migrate to use semantic_text --- .../src/field_maps/types.ts | 1 + .../impl/capabilities/index.ts | 2 +- .../server/__mocks__/msearch_query.ts | 10 ++-- .../server/__mocks__/vector_search_query.ts | 10 ++-- .../anonymization_fields/helpers.ts | 14 +++-- .../conversations/helpers.ts | 14 +++-- .../conversations/update_conversation.ts | 2 +- .../create_knowledge_base_entry.ts | 55 +++---------------- .../field_maps_configuration.ts | 6 ++ .../knowledge_base/helpers.ts | 8 +-- .../knowledge_base/types.ts | 3 + .../prompts/helpers.ts | 14 +++-- .../server/ai_assistant_service/index.ts | 6 +- .../lib/data_stream/documents_data_writer.ts | 19 ++++--- .../content_loaders/security_labs_loader.ts | 17 ++++-- .../helpers/get_vector_search_query.test.ts | 30 ++++------ .../helpers/get_vector_search_query.ts | 8 +-- .../knowledge_base/post_knowledge_base.ts | 2 +- .../common/experimental_features.ts | 2 +- .../scripts/run_cypress/parallel.ts | 5 +- .../test/security_solution_cypress/config.ts | 2 +- 21 files changed, 100 insertions(+), 130 deletions(-) diff --git a/packages/kbn-data-stream-adapter/src/field_maps/types.ts b/packages/kbn-data-stream-adapter/src/field_maps/types.ts index 4f42a6c6b686d..921df9cac238b 100644 --- a/packages/kbn-data-stream-adapter/src/field_maps/types.ts +++ b/packages/kbn-data-stream-adapter/src/field_maps/types.ts @@ -53,5 +53,6 @@ export interface FieldMap { scaling_factor?: number; dynamic?: boolean | 'strict'; properties?: Record; + inference_id?: string; }; } diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index c1c101fd74cd8..54c24f6ce7b8f 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/msearch_query.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/msearch_query.ts index e411dfaa2f1ef..ae5adcfab61aa 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/msearch_query.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/msearch_query.ts @@ -34,12 +34,10 @@ export const mSearchQueryBody: MsearchQueryBody = { ], must: [ { - text_expansion: { - 'vector.tokens': { - model_id: '.elser_model_2', - model_text: - 'Generate an ESQL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.', - }, + semantic: { + field: 'semantic_text', + query: + 'Generate an ESQL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.', }, }, ], diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts index 30fbd0ad2c58f..f8966604e637d 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts @@ -26,12 +26,10 @@ export const mockVectorSearchQuery: QueryDslQueryContainer = { ], must: [ { - text_expansion: { - 'vector.tokens': { - model_id: '.elser_model_2', - model_text: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.', - }, + semantic: { + field: 'inference_field', + query: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.', }, }, ], diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts index 9a4a3b6e1c0ce..0f577df4e56e1 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts @@ -99,7 +99,8 @@ export const getUpdateScript = ({ isPatch?: boolean; }) => { return { - source: ` + script: { + source: ` if (params.assignEmpty == true || params.containsKey('allowed')) { ctx._source.allowed = params.allowed; } @@ -108,11 +109,12 @@ export const getUpdateScript = ({ } ctx._source.updated_at = params.updated_at; `, - lang: 'painless', - params: { - ...anonymizationField, // when assigning undefined in painless, it will remove property and wil set it to null - // for patch we don't want to remove unspecified value in payload - assignEmpty: !(isPatch ?? true), + lang: 'painless', + params: { + ...anonymizationField, // when assigning undefined in painless, it will remove property and wil set it to null + // for patch we don't want to remove unspecified value in payload + assignEmpty: !(isPatch ?? true), + }, }, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts index 9e52b4a7414a6..bdd1107942cc1 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts @@ -15,7 +15,8 @@ export const getUpdateScript = ({ isPatch?: boolean; }) => { return { - source: ` + script: { + source: ` if (params.assignEmpty == true || params.containsKey('api_config')) { if (ctx._source.api_config != null) { if (params.assignEmpty == true || params.api_config.containsKey('connector_id')) { @@ -70,11 +71,12 @@ export const getUpdateScript = ({ } ctx._source.updated_at = params.updated_at; `, - lang: 'painless', - params: { - ...conversation, // when assigning undefined in painless, it will remove property and wil set it to null - // for patch we don't want to remove unspecified value in payload - assignEmpty: !(isPatch ?? true), + lang: 'painless', + params: { + ...conversation, // when assigning undefined in painless, it will remove property and wil set it to null + // for patch we don't want to remove unspecified value in payload + assignEmpty: !(isPatch ?? true), + }, }, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts index 807fea2decd99..7e9ee336f6fe1 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts @@ -76,7 +76,7 @@ export const updateConversation = async ({ }, }, refresh: true, - script: getUpdateScript({ conversation: params, isPatch }), + script: getUpdateScript({ conversation: params, isPatch }).script, }); if (response.failures && response.failures.length > 0) { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts index 23f73501b1056..838227b7c8e15 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts @@ -139,55 +139,11 @@ export const getUpdateScript = ({ entry: UpdateKnowledgeBaseEntrySchema; isPatch?: boolean; }) => { + // Cannot use script for updating documents with semantic_text fields return { - source: ` - if (params.assignEmpty == true || params.containsKey('name')) { - ctx._source.name = params.name; - } - if (params.assignEmpty == true || params.containsKey('type')) { - ctx._source.type = params.type; - } - if (params.assignEmpty == true || params.containsKey('users')) { - ctx._source.users = params.users; - } - if (params.assignEmpty == true || params.containsKey('query_description')) { - ctx._source.query_description = params.query_description; - } - if (params.assignEmpty == true || params.containsKey('input_schema')) { - ctx._source.input_schema = params.input_schema; - } - if (params.assignEmpty == true || params.containsKey('output_fields')) { - ctx._source.output_fields = params.output_fields; - } - if (params.assignEmpty == true || params.containsKey('kb_resource')) { - ctx._source.kb_resource = params.kb_resource; - } - if (params.assignEmpty == true || params.containsKey('required')) { - ctx._source.required = params.required; - } - if (params.assignEmpty == true || params.containsKey('source')) { - ctx._source.source = params.source; - } - if (params.assignEmpty == true || params.containsKey('text')) { - ctx._source.text = params.text; - } - if (params.assignEmpty == true || params.containsKey('description')) { - ctx._source.description = params.description; - } - if (params.assignEmpty == true || params.containsKey('field')) { - ctx._source.field = params.field; - } - if (params.assignEmpty == true || params.containsKey('index')) { - ctx._source.index = params.index; - } - ctx._source.updated_at = params.updated_at; - ctx._source.updated_by = params.updated_by; - `, - lang: 'painless', - params: { - ...entry, // when assigning undefined in painless, it will remove property and wil set it to null - // for patch we don't want to remove unspecified value in payload - assignEmpty: !(isPatch ?? true), + doc: { + ...entry, + semantic_text: entry.text, }, }; }; @@ -247,6 +203,8 @@ export const transformToCreateSchema = ({ required: entry.required ?? false, source: entry.source, text: entry.text, + semantic_text: entry.text, + // deprecated vector: undefined, }; }; @@ -289,6 +247,7 @@ export const transformToLegacyCreateSchema = ({ }, ], ...entry, + // deprecated vector: undefined, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index 0712664bbfeed..75206a61c674e 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -169,6 +169,12 @@ export const knowledgeBaseFieldMapV2: FieldMap = { required: false, }, // Embeddings field + semantic_text: { + type: 'semantic_text', + array: false, + required: false, + inference_id: '.elser-2', + }, vector: { type: 'object', array: false, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index de76a38135f0b..a7c406011670f 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -118,11 +118,9 @@ export const getKBVectorSearchQuery = ({ bool: { must: [ { - text_expansion: { - 'vector.tokens': { - model_id: modelId, - model_text: query, - }, + semantic: { + field: 'semantic_text', + query, }, }, ...requiredFilter, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts index 3de1a15d79b2a..443d03941ccdd 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts @@ -27,6 +27,7 @@ export interface EsDocumentEntry { required: boolean; source: string; text: string; + semantic_text?: string; vector?: { tokens: Record; model_id: string; @@ -99,6 +100,7 @@ export interface UpdateKnowledgeBaseEntrySchema { required?: boolean; source?: string; text?: string; + semantic_text?: string; vector?: { tokens: Record; model_id: string; @@ -135,6 +137,7 @@ export interface CreateKnowledgeBaseEntrySchema { required?: boolean; source?: string; text?: string; + semantic_text?: string; vector?: { tokens: Record; model_id: string; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/prompts/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/prompts/helpers.ts index a4534972c8478..eb71270127b2a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/prompts/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/prompts/helpers.ts @@ -143,7 +143,8 @@ export const getUpdateScript = ({ isPatch?: boolean; }) => { return { - source: ` + script: { + source: ` if (params.assignEmpty == true || params.containsKey('content')) { ctx._source.content = params.content; } @@ -158,11 +159,12 @@ export const getUpdateScript = ({ } ctx._source.updated_at = params.updated_at; `, - lang: 'painless', - params: { - ...prompt, // when assigning undefined in painless, it will remove property and wil set it to null - // for patch we don't want to remove unspecified value in payload - assignEmpty: !(isPatch ?? true), + lang: 'painless', + params: { + ...prompt, // when assigning undefined in painless, it will remove property and wil set it to null + // for patch we don't want to remove unspecified value in payload + assignEmpty: !(isPatch ?? true), + }, }, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index bfdf8b96f44b0..e6bf145931c02 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -97,7 +97,7 @@ export class AIAssistantService { this.knowledgeBaseDataStream = this.createDataStream({ resource: 'knowledgeBase', kibanaVersion: options.kibanaVersion, - fieldMap: knowledgeBaseFieldMap, // TODO: use V2 if FF is enabled + fieldMap: knowledgeBaseFieldMapV2, // TODO: use V2 if FF is enabled }); this.promptsDataStream = this.createDataStream({ resource: 'prompts', @@ -151,7 +151,7 @@ export class AIAssistantService { name: this.resourceNames.indexTemplate[resource], componentTemplateRefs: [this.resourceNames.componentTemplate[resource]], // Apply `default_pipeline` if pipeline exists for resource - ...(resource in this.resourceNames.pipelines + ...(resource in this.resourceNames.pipelines && !this.hasInitializedV2KnowledgeBase ? { template: { settings: { @@ -202,7 +202,7 @@ export class AIAssistantService { id: this.resourceNames.pipelines.knowledgeBase, }); // TODO: When FF is removed, ensure pipeline is re-created for those upgrading - if (!pipelineCreated || this.v2KnowledgeBaseEnabled) { + if (!pipelineCreated && !this.v2KnowledgeBaseEnabled) { this.options.logger.debug( `Installing ingest pipeline - ${this.resourceNames.pipelines.knowledgeBase}` ); diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 87ec80568dbdd..47c37162fc688 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -34,7 +34,10 @@ interface BulkParams { documentsToCreate?: TCreateParams[]; documentsToUpdate?: TUpdateParams[]; documentsToDelete?: string[]; - getUpdateScript?: (document: TUpdateParams, updatedAt: string) => Script; + getUpdateScript?: ( + document: TUpdateParams, + updatedAt: string + ) => { script?: Script; doc?: TUpdateParams }; authenticatedUser?: AuthenticatedUser; } @@ -73,8 +76,8 @@ export class DocumentsDataWriter implements DocumentsDataWriter { body: await this.buildBulkOperations(params), }, { - // Increasing timout to 2min as KB docs were failing to load after 30s - requestTimeout: 120000, + // Increasing timout to 10min as KB docs were failing to load after 30s + requestTimeout: 600000, } ); @@ -146,7 +149,10 @@ export class DocumentsDataWriter implements DocumentsDataWriter { private getUpdateDocumentsQuery = async ( documentsToUpdate: TUpdateParams[], - getUpdateScript: (document: TUpdateParams, updatedAt: string) => Script, + getUpdateScript: ( + document: TUpdateParams, + updatedAt: string + ) => { script?: Script; doc?: TUpdateParams }, authenticatedUser?: AuthenticatedUser ) => { const updatedAt = new Date().toISOString(); @@ -191,10 +197,7 @@ export class DocumentsDataWriter implements DocumentsDataWriter { _source: true, }, }, - { - script: getUpdateScript(document, updatedAt), - upsert: { counter: 1 }, - }, + getUpdateScript(document, updatedAt), ]); }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts index 10566b3e5a1d5..facd89b3ada55 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts @@ -11,7 +11,9 @@ import { TextLoader } from 'langchain/document_loaders/fs/text'; import { resolve } from 'path'; import { Document } from 'langchain/document'; import { Metadata } from '@kbn/elastic-assistant-common'; +import pMap from 'p-map'; +import { chunk } from 'lodash'; import { addRequiredKbResourceMetadata } from './add_required_kb_resource_metadata'; import { SECURITY_LABS_RESOURCE } from '../../../routes/knowledge_base/constants'; import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; @@ -42,10 +44,17 @@ export const loadSecurityLabs = async ( logger.info(`Loading ${docs.length} Security Labs docs into the Knowledge Base`); - const response = await kbDataClient.addKnowledgeBaseDocuments({ - documents: docs, - global: true, - }); + const response = ( + await pMap( + chunk(docs, 30), + (partialDocs) => + kbDataClient.addKnowledgeBaseDocuments({ + documents: partialDocs, + global: true, + }), + { concurrency: 1 } + ) + ).flat(); logger.info(`Loaded ${response?.length ?? 0} Security Labs docs into the Knowledge Base`); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts index da6a7227953f2..148ed7ffab362 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts @@ -25,12 +25,10 @@ describe('getVectorSearchQuery', () => { filter: undefined, must: [ { - text_expansion: { - 'vector.tokens': { - model_id: '.elser_model_2', - model_text: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', - }, + semantic: { + field: 'semantic_text', + query: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', }, }, ], @@ -52,12 +50,10 @@ describe('getVectorSearchQuery', () => { filter: undefined, must: [ { - text_expansion: { - 'vector.tokens': { - model_id: '.elser_model_2', - model_text: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', - }, + semantic: { + field: 'semantic_text', + query: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', }, }, ], @@ -102,12 +98,10 @@ describe('getVectorSearchQuery', () => { filter, must: [ { - text_expansion: { - 'vector.tokens': { - model_id: '.elser_model_2', - model_text: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', - }, + semantic: { + field: 'semantic_text', + query: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', }, }, ], diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts index 613ee5c501560..d249119537dbc 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts @@ -37,11 +37,9 @@ export const getVectorSearchQuery = ({ must_not: [...mustNotTerms], must: [ { - text_expansion: { - 'vector.tokens': { - model_id: modelId, - model_text: query, - }, + semantic: { + field: 'semantic_text', + query, }, }, ], diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 96317da303ac1..fae436b5bf3f1 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -20,7 +20,7 @@ import { isV2KnowledgeBaseEnabled } 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 -const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes +const ROUTE_HANDLER_TIMEOUT = 20 * 60 * 1000; // 10 * 60 seconds = 20 minutes /** * Load Knowledge Base index, pipeline, and resources (collection of documents) diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 5e438669916c6..4572eb0255796 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -111,7 +111,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables new Knowledge Base Entries features, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, /** * Enables the Managed User section inside the new user details flyout. diff --git a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts index 0f93e4fceb10c..6174fd125ae7b 100644 --- a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts +++ b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts @@ -348,10 +348,7 @@ ${JSON.stringify( procs, config, installDir: options?.installDir, - extraKbnOpts: - options?.installDir || options?.ci || !isOpen - ? [] - : ['--dev', '--no-dev-config', '--no-dev-credentials'], + extraKbnOpts: options?.installDir || options?.ci || !isOpen ? [] : ['--dev'], onEarlyExit, inspect: argv.inspect, }); diff --git a/x-pack/test/security_solution_cypress/config.ts b/x-pack/test/security_solution_cypress/config.ts index f02968945087d..2fd39abac8b7c 100644 --- a/x-pack/test/security_solution_cypress/config.ts +++ b/x-pack/test/security_solution_cypress/config.ts @@ -45,7 +45,7 @@ export default async function ({ readConfigFile }: FtrConfigProviderContext) { '--xpack.alerting.rules.minimumScheduleInterval.value=1s', '--xpack.ruleRegistry.unsafe.legacyMultiTenancy.enabled=true', // mock cloud to enable the guided onboarding tour in e2e tests - '--xpack.cloud.id=test', + // '--xpack.cloud.id=test', `--home.disableWelcomeScreen=true`, // Specify which version of the detection-rules package to install // `--xpack.securitySolution.prebuiltRulesPackageVersion=8.3.1`, From e3d801ba24c15a2e3677e6fc46203b1ae0bd2738 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 21 Oct 2024 10:56:03 +0200 Subject: [PATCH 08/41] fix --- .../server/__mocks__/vector_search_query.ts | 2 +- .../security_solution/scripts/run_cypress/parallel.ts | 5 ++++- x-pack/test/security_solution_cypress/config.ts | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts index f8966604e637d..04263c5d242bb 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts @@ -27,7 +27,7 @@ export const mockVectorSearchQuery: QueryDslQueryContainer = { must: [ { semantic: { - field: 'inference_field', + field: 'semantic_text', query: 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.', }, diff --git a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts index 6174fd125ae7b..0f93e4fceb10c 100644 --- a/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts +++ b/x-pack/plugins/security_solution/scripts/run_cypress/parallel.ts @@ -348,7 +348,10 @@ ${JSON.stringify( procs, config, installDir: options?.installDir, - extraKbnOpts: options?.installDir || options?.ci || !isOpen ? [] : ['--dev'], + extraKbnOpts: + options?.installDir || options?.ci || !isOpen + ? [] + : ['--dev', '--no-dev-config', '--no-dev-credentials'], onEarlyExit, inspect: argv.inspect, }); diff --git a/x-pack/test/security_solution_cypress/config.ts b/x-pack/test/security_solution_cypress/config.ts index 2fd39abac8b7c..f02968945087d 100644 --- a/x-pack/test/security_solution_cypress/config.ts +++ b/x-pack/test/security_solution_cypress/config.ts @@ -45,7 +45,7 @@ export default async function ({ readConfigFile }: FtrConfigProviderContext) { '--xpack.alerting.rules.minimumScheduleInterval.value=1s', '--xpack.ruleRegistry.unsafe.legacyMultiTenancy.enabled=true', // mock cloud to enable the guided onboarding tour in e2e tests - // '--xpack.cloud.id=test', + '--xpack.cloud.id=test', `--home.disableWelcomeScreen=true`, // Specify which version of the detection-rules package to install // `--xpack.securitySolution.prebuiltRulesPackageVersion=8.3.1`, From 55e8c5568bfc4041d5c78b0d7c37e37cbb8cd85e Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Mon, 21 Oct 2024 09:51:25 +0000 Subject: [PATCH 09/41] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- .../elastic_assistant/server/ai_assistant_service/index.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index e6bf145931c02..79dc5ccc54e77 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -26,10 +26,7 @@ import { conversationsFieldMap } from '../ai_assistant_data_clients/conversation import { assistantPromptsFieldMap } from '../ai_assistant_data_clients/prompts/field_maps_configuration'; import { assistantAnonymizationFieldsFieldMap } from '../ai_assistant_data_clients/anonymization_fields/field_maps_configuration'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; -import { - knowledgeBaseFieldMap, - knowledgeBaseFieldMapV2, -} from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; +import { knowledgeBaseFieldMapV2 } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; import { AIAssistantKnowledgeBaseDataClient, GetAIAssistantKnowledgeBaseDataClientParams, From 8cfbf61f2ce582893a26069186cec3992713d967 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 22 Oct 2024 09:59:52 +0200 Subject: [PATCH 10/41] WIP --- packages/kbn-data-stream-adapter/src/field_maps/types.ts | 1 + .../knowledge_base/field_maps_configuration.ts | 1 + .../elastic_assistant/server/ai_assistant_service/index.ts | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/kbn-data-stream-adapter/src/field_maps/types.ts b/packages/kbn-data-stream-adapter/src/field_maps/types.ts index 921df9cac238b..7c8cbbcce3854 100644 --- a/packages/kbn-data-stream-adapter/src/field_maps/types.ts +++ b/packages/kbn-data-stream-adapter/src/field_maps/types.ts @@ -54,5 +54,6 @@ export interface FieldMap { dynamic?: boolean | 'strict'; properties?: Record; inference_id?: string; + copy_to?: string; }; } diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index 75206a61c674e..1d31a924eaa09 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -140,6 +140,7 @@ export const knowledgeBaseFieldMapV2: FieldMap = { type: 'keyword', array: false, required: false, + copy_to: 'semantic_text', }, // Discriminator: 'document' | 'index' type: { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index e6bf145931c02..eb8d7dca1aa83 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -151,7 +151,7 @@ export class AIAssistantService { name: this.resourceNames.indexTemplate[resource], componentTemplateRefs: [this.resourceNames.componentTemplate[resource]], // Apply `default_pipeline` if pipeline exists for resource - ...(resource in this.resourceNames.pipelines && !this.hasInitializedV2KnowledgeBase + ...(resource in this.resourceNames.pipelines ? { template: { settings: { From cffe767dae4d8656ca83a7209ff43984622ff833 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 22 Oct 2024 14:58:24 +0200 Subject: [PATCH 11/41] fix --- .../assistant/chat_send/use_chat_send.tsx | 3 +- .../index_entry_editor.tsx | 17 ++++++++- .../field_maps_configuration.ts | 4 +- .../knowledge_base/index.ts | 31 ++++++++++++--- .../knowledge_base/ingest_pipeline.ts | 38 ++++++++++++------- .../server/ai_assistant_service/helpers.ts | 3 ++ .../server/ai_assistant_service/index.ts | 21 ++++++++-- 7 files changed, 90 insertions(+), 27 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 4ea376518b5a7..e0a6717a72ed3 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -62,7 +62,8 @@ export const useChatSend = ({ kbStatus?.elser_exists && kbStatus?.index_exists && kbStatus?.pipeline_exists && - kbStatus?.security_labs_exists; + kbStatus?.security_labs_exists && + !kbStatus?.is_setup_in_progress; // Handles sending latest user prompt to API const handleSendMessage = useCallback( diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx index 550861bcbffd9..4109e1f608c98 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx @@ -96,12 +96,20 @@ export const IndexEntryEditor: React.FC = React.memo( })); }, [dataViews]); + const indexExists = useAsync(async () => { + if (!entry?.index?.length) return true; + + return !!(await dataViews.getExistingIndices([entry.index])).length; + }, [entry?.index]); + + console.error('indexExists', indexExists); + const indexFields = useAsync( async () => dataViews.getFieldsForWildcard({ pattern: entry?.index ?? '', }), - [] + [entry?.index] ); const fieldOptions = useMemo( @@ -251,7 +259,12 @@ export const IndexEntryEditor: React.FC = React.memo( fullWidth /> - + {"Index doesn't exist"}} + > ({ +export const knowledgeBaseIngestPipeline = ({ + id, + modelId, + v2KnowledgeBaseEnabled, +}: { + id: string; + modelId: string; + v2KnowledgeBaseEnabled: boolean; +}) => ({ id, description: 'Embedding pipeline for Elastic AI Assistant ELSER Knowledge Base', - processors: [ - { - inference: { - if: 'ctx?.text != null', - model_id: modelId, - input_output: [ - { - input_field: 'text', - output_field: 'vector.tokens', + processors: !v2KnowledgeBaseEnabled + ? [ + { + inference: { + if: 'ctx?.text != null', + model_id: modelId, + input_output: [ + { + input_field: 'text', + output_field: 'vector.tokens', + }, + ], }, - ], - }, - }, - ], + }, + ] + : [], }); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts index 07da930320712..93338174364fc 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts @@ -54,6 +54,7 @@ interface CreatePipelineParams { esClient: ElasticsearchClient; id: string; modelId: string; + v2KnowledgeBaseEnabled: boolean; } /** @@ -70,12 +71,14 @@ export const createPipeline = async ({ esClient, id, modelId, + v2KnowledgeBaseEnabled, }: CreatePipelineParams): Promise => { try { const response = await esClient.ingest.putPipeline( knowledgeBaseIngestPipeline({ id, modelId, + v2KnowledgeBaseEnabled, }) ); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index ee6f4b4dacf96..e0ee261576a4f 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -26,7 +26,10 @@ import { conversationsFieldMap } from '../ai_assistant_data_clients/conversation import { assistantPromptsFieldMap } from '../ai_assistant_data_clients/prompts/field_maps_configuration'; import { assistantAnonymizationFieldsFieldMap } from '../ai_assistant_data_clients/anonymization_fields/field_maps_configuration'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; -import { knowledgeBaseFieldMapV2 } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; +import { + knowledgeBaseFieldMap, + knowledgeBaseFieldMapV2, +} from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; import { AIAssistantKnowledgeBaseDataClient, GetAIAssistantKnowledgeBaseDataClientParams, @@ -94,7 +97,7 @@ export class AIAssistantService { this.knowledgeBaseDataStream = this.createDataStream({ resource: 'knowledgeBase', kibanaVersion: options.kibanaVersion, - fieldMap: knowledgeBaseFieldMapV2, // TODO: use V2 if FF is enabled + fieldMap: knowledgeBaseFieldMap, }); this.promptsDataStream = this.createDataStream({ resource: 'prompts', @@ -144,11 +147,15 @@ export class AIAssistantService { fieldMap, }); + console.error('CREATE DATA STREAM', resource, this.v2KnowledgeBaseEnabled); + newDataStream.setIndexTemplate({ name: this.resourceNames.indexTemplate[resource], componentTemplateRefs: [this.resourceNames.componentTemplate[resource]], // Apply `default_pipeline` if pipeline exists for resource - ...(resource in this.resourceNames.pipelines + ...(resource in this.resourceNames.pipelines && + // Remove this param and initialization when the `assistantKnowledgeBaseByDefault` feature flag is removed + !(resource === 'knowledgeBase' && this.v2KnowledgeBaseEnabled) ? { template: { settings: { @@ -199,7 +206,12 @@ export class AIAssistantService { id: this.resourceNames.pipelines.knowledgeBase, }); // TODO: When FF is removed, ensure pipeline is re-created for those upgrading - if (!pipelineCreated && !this.v2KnowledgeBaseEnabled) { + if ( + // Install for v1 + (!this.v2KnowledgeBaseEnabled && !pipelineCreated) || + // Upgrade from v1 to v2 + (pipelineCreated && this.v2KnowledgeBaseEnabled) + ) { this.options.logger.debug( `Installing ingest pipeline - ${this.resourceNames.pipelines.knowledgeBase}` ); @@ -207,6 +219,7 @@ export class AIAssistantService { esClient, id: this.resourceNames.pipelines.knowledgeBase, modelId: await this.getElserId(), + v2KnowledgeBaseEnabled: this.v2KnowledgeBaseEnabled, }); this.options.logger.debug(`Installed ingest pipeline: ${response}`); From e251ca2a7c19a9cb08a983303fc9a47c1705f9c0 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 22 Oct 2024 15:03:23 +0200 Subject: [PATCH 12/41] fix --- .../kbn-elastic-assistant-common/impl/capabilities/index.ts | 2 +- .../plugins/security_solution/common/experimental_features.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 54c24f6ce7b8f..c1c101fd74cd8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 4572eb0255796..5e438669916c6 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -111,7 +111,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables new Knowledge Base Entries features, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, /** * Enables the Managed User section inside the new user details flyout. From 033cc5717ed7c689ebea46fd825778381288b94f Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 22 Oct 2024 15:11:49 +0200 Subject: [PATCH 13/41] fix --- .../knowledge_base/index.ts | 2 +- .../server/ai_assistant_service/index.ts | 2 -- .../helpers/get_vector_search_query.test.ts | 20 +++++++++++-------- .../helpers/get_vector_search_query.ts | 8 +++++--- .../knowledge_base/post_knowledge_base.ts | 2 +- 5 files changed, 19 insertions(+), 15 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index b9d91c48df5b8..df1b526a20da5 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -175,7 +175,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }); } } catch (error) { - this.options.logger.error(`Error creating ELSER inference endpoint '${elserId}':\n${error}`); + this.options.logger.error(`Error deploying ELSER model '${elserId}':\n${error}`); throw new Error(`Error deploying ELSER model '${elserId}':\n${error}`); } }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index e0ee261576a4f..a7b54dd5ca4be 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -147,8 +147,6 @@ export class AIAssistantService { fieldMap, }); - console.error('CREATE DATA STREAM', resource, this.v2KnowledgeBaseEnabled); - newDataStream.setIndexTemplate({ name: this.resourceNames.indexTemplate[resource], componentTemplateRefs: [this.resourceNames.componentTemplate[resource]], diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts index 148ed7ffab362..4c92bb91fc0da 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts @@ -25,10 +25,12 @@ describe('getVectorSearchQuery', () => { filter: undefined, must: [ { - semantic: { - field: 'semantic_text', - query: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', + text_expansion: { + 'vector.tokens': { + model_id: '.elser_model_2', + model_text: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', + }, }, }, ], @@ -50,10 +52,12 @@ describe('getVectorSearchQuery', () => { filter: undefined, must: [ { - semantic: { - field: 'semantic_text', - query: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', + text_expansion: { + 'vector.tokens': { + model_id: '.elser_model_2', + model_text: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', + }, }, }, ], diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts index d249119537dbc..613ee5c501560 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.ts @@ -37,9 +37,11 @@ export const getVectorSearchQuery = ({ must_not: [...mustNotTerms], must: [ { - semantic: { - field: 'semantic_text', - query, + text_expansion: { + 'vector.tokens': { + model_id: modelId, + model_text: query, + }, }, }, ], diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index fae436b5bf3f1..cbb3ce1c10a6c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -20,7 +20,7 @@ import { isV2KnowledgeBaseEnabled } 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 -const ROUTE_HANDLER_TIMEOUT = 20 * 60 * 1000; // 10 * 60 seconds = 20 minutes +const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 20 minutes /** * Load Knowledge Base index, pipeline, and resources (collection of documents) From 82e4ca352273f3ed17fb83a0a193328953b71532 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 22 Oct 2024 18:17:55 +0200 Subject: [PATCH 14/41] filter out tools with missing indices --- .../knowledge_base/index.ts | 109 ++++++++++++++---- 1 file changed, 89 insertions(+), 20 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index df1b526a20da5..d5e76ffdccdb6 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -25,6 +25,8 @@ import pRetry from 'p-retry'; import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey'; import { StructuredTool } from '@langchain/core/tools'; import { ElasticsearchClient } from '@kbn/core/server'; +import { IndexPatternsFetcher } from '@kbn/data-views-plugin/server'; +import { map } from 'lodash'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { AssistantToolParams, GetElser } from '../../types'; import { @@ -214,6 +216,51 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } }; + public isInferenceEndpointExists = async () => { + try { + const esClient = await this.options.elasticsearchClientPromise; + + return await esClient.inference.get({ + inference_id: ASSISTANT_ELSER_INFERENCE_ID, + task_type: 'sparse_embedding', + }); + } catch (error) { + return false; + } + }; + + public createInferenceEndpoint = async () => { + const elserId = await this.options.getElserId(); + this.options.logger.debug(`Deploying ELSER model '${elserId}'...`); + try { + const esClient = await this.options.elasticsearchClientPromise; + if (this.isV2KnowledgeBaseEnabled) { + await esClient.inference.put({ + task_type: 'sparse_embedding', + inference_id: ASSISTANT_ELSER_INFERENCE_ID, + inference_config: { + service: 'elasticsearch', + service_settings: { + adaptive_allocations: { + enabled: true, + min_number_of_allocations: 0, + max_number_of_allocations: 8, + }, + num_threads: 1, + model_id: elserId, + }, + task_settings: {}, + }, + }); + } + } catch (error) { + this.options.logger.error( + `Error creating inference endpoint for ELSER model '${elserId}':\n${error}` + ); + throw new Error(`Error creating inference endpoint for ELSER model '${elserId}':\n${error}`); + } + }; + /** * Downloads and deploys recommended ELSER (if not already), then loads ES|QL docs * @@ -280,19 +327,34 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { this.options.logger.debug(`ELSER model '${elserId}' is already installed`); } - const isDeployed = await this.isModelDeployed(); - if (!isDeployed) { - await this.deployModel(); - await pRetry( - async () => - (await this.isModelDeployed()) - ? Promise.resolve() - : Promise.reject(new Error('Model not deployed')), - { minTimeout: 2000, retries: 10 } - ); - this.options.logger.debug(`ELSER model '${elserId}' successfully deployed!`); + if (!this.isV2KnowledgeBaseEnabled) { + const isDeployed = await this.isModelDeployed(); + if (!isDeployed) { + await this.deployModel(); + await pRetry( + async () => + (await this.isModelDeployed()) + ? Promise.resolve() + : Promise.reject(new Error('Model not deployed')), + { minTimeout: 2000, retries: 10 } + ); + this.options.logger.debug(`ELSER model '${elserId}' successfully deployed!`); + } else { + this.options.logger.debug(`ELSER model '${elserId}' is already deployed`); + } } else { - this.options.logger.debug(`ELSER model '${elserId}' is already deployed`); + const inferenceExists = await this.isInferenceEndpointExists(); + if (!inferenceExists) { + await this.createInferenceEndpoint(); + + this.options.logger.debug( + `Inference endpoint for ELSER model '${elserId}' successfully deployed!` + ); + } else { + this.options.logger.debug( + `Inference endpoint for ELSER model '${elserId}' is already deployed` + ); + } } this.options.logger.debug(`Checking if Knowledge Base docs have been loaded...`); @@ -616,14 +678,21 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { if (results) { const entries = transformESSearchToKnowledgeBaseEntry(results.data) as IndexEntry[]; - return entries.map((indexEntry) => { - return getStructuredToolForIndexEntry({ - indexEntry, - esClient, - logger: this.options.logger, - elserId, - }); - }); + const indexPatternFetcher = new IndexPatternsFetcher(esClient); + const existingIndices = await indexPatternFetcher.getExistingIndices(map(entries, 'index')); + return ( + entries + // Filter out any IndexEntries that don't have an existing index + .filter((entry) => existingIndices.includes(entry.index)) + .map((indexEntry) => { + return getStructuredToolForIndexEntry({ + indexEntry, + esClient, + logger: this.options.logger, + elserId, + }); + }) + ); } } catch (e) { this.options.logger.error(`kbDataClient.getAssistantTools() - Failed to fetch IndexEntries`); From 70ef30813fc665b3f792e94d40ca187ea8b38d62 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 22 Oct 2024 18:26:30 +0200 Subject: [PATCH 15/41] fix --- .../index_entry_editor.tsx | 2 -- .../helpers/get_vector_search_query.test.ts | 10 ++++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx index 4109e1f608c98..1d5611d3b0373 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx @@ -102,8 +102,6 @@ export const IndexEntryEditor: React.FC = React.memo( return !!(await dataViews.getExistingIndices([entry.index])).length; }, [entry?.index]); - console.error('indexExists', indexExists); - const indexFields = useAsync( async () => dataViews.getFieldsForWildcard({ diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts index 4c92bb91fc0da..da6a7227953f2 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/helpers/get_vector_search_query.test.ts @@ -102,10 +102,12 @@ describe('getVectorSearchQuery', () => { filter, must: [ { - semantic: { - field: 'semantic_text', - query: - 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', + text_expansion: { + 'vector.tokens': { + model_id: '.elser_model_2', + model_text: + 'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called follow_up that contains a value of true, otherwise, it should contain false. The user names should also be enriched with their respective group names.', + }, }, }, ], From 68a1f66715e038bf99218e7a2ca28ec6864e4984 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 22 Oct 2024 17:04:32 +0000 Subject: [PATCH 16/41] [CI] Auto-commit changed files from 'node scripts/notice' --- x-pack/plugins/elastic_assistant/tsconfig.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 747a58ed930d3..d3436f28a1d3e 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -48,7 +48,8 @@ "@kbn/apm-utils", "@kbn/std", "@kbn/zod", - "@kbn/inference-plugin" + "@kbn/inference-plugin", + "@kbn/data-views-plugin" ], "exclude": [ "target/**/*", From 185f44d4b7cfb939c10c994695697e3ada3d2888 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 23 Oct 2024 09:51:34 +0200 Subject: [PATCH 17/41] test --- .../server/ai_assistant_data_clients/knowledge_base/index.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index d5e76ffdccdb6..4c1361dc4b8b3 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -162,7 +162,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { adaptive_allocations: { enabled: true, min_number_of_allocations: 0, - max_number_of_allocations: 8, + max_number_of_allocations: 1, }, num_threads: 1, model_id: elserId, From 0e85e297819c6831af7d01ea713852df24dc765b Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 23 Oct 2024 11:24:49 +0200 Subject: [PATCH 18/41] test --- .../server/ai_assistant_data_clients/knowledge_base/index.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 4c1361dc4b8b3..8267b0780942b 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -244,7 +244,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { adaptive_allocations: { enabled: true, min_number_of_allocations: 0, - max_number_of_allocations: 8, + max_number_of_allocations: 1, }, num_threads: 1, model_id: elserId, From fc84417af67bb71ce8ab289acac8acf2fb3cbfcd Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 23 Oct 2024 15:35:38 +0200 Subject: [PATCH 19/41] cleanup --- .../knowledge_base/index.ts | 28 +++---------------- 1 file changed, 4 insertions(+), 24 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 8267b0780942b..1ae392bd531f7 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -152,30 +152,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { this.options.logger.debug(`Deploying ELSER model '${elserId}'...`); try { const esClient = await this.options.elasticsearchClientPromise; - if (this.isV2KnowledgeBaseEnabled) { - await esClient.inference.put({ - task_type: 'sparse_embedding', - inference_id: ASSISTANT_ELSER_INFERENCE_ID, - inference_config: { - service: 'elasticsearch', - service_settings: { - adaptive_allocations: { - enabled: true, - min_number_of_allocations: 0, - max_number_of_allocations: 1, - }, - num_threads: 1, - model_id: elserId, - }, - task_settings: {}, - }, - }); - } else { - await esClient.ml.startTrainedModelDeployment({ - model_id: elserId, - wait_for: 'fully_allocated', - }); - } + await esClient.ml.startTrainedModelDeployment({ + model_id: elserId, + wait_for: 'fully_allocated', + }); } catch (error) { this.options.logger.error(`Error deploying ELSER model '${elserId}':\n${error}`); throw new Error(`Error deploying ELSER model '${elserId}':\n${error}`); From d1ac1b6b26862706dc3bec64bc943e8b5a19fc25 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 23 Oct 2024 17:24:59 +0200 Subject: [PATCH 20/41] test --- .../ai_assistant_data_clients/knowledge_base/index.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 1ae392bd531f7..4769321c88ea7 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -221,11 +221,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { inference_config: { service: 'elasticsearch', service_settings: { - adaptive_allocations: { - enabled: true, - min_number_of_allocations: 0, - max_number_of_allocations: 1, - }, + num_allocations: 1, num_threads: 1, model_id: elserId, }, From 3786057d6b6cb2d01f82bb27972d42fcdee8fa9a Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 23 Oct 2024 19:33:42 +0200 Subject: [PATCH 21/41] test --- .../kbn-test/src/functional_test_runner/lib/config/schema.ts | 2 +- .../kbn-elastic-assistant-common/impl/capabilities/index.ts | 2 +- .../plugins/security_solution/common/experimental_features.ts | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts b/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts index bb240a0416479..dba2174003c7f 100644 --- a/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts +++ b/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts @@ -145,7 +145,7 @@ export const schema = Joi.object() grep: Joi.string(), invert: Joi.boolean().default(false), slow: Joi.number().default(30000), - timeout: Joi.number().default(INSPECTING ? 360000 * 100 : 360000), + timeout: Joi.number().default(INSPECTING ? 360000 * 100 : 360000 * 2), ui: Joi.string().default('bdd'), }) .default(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index c1c101fd74cd8..54c24f6ce7b8f 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 5e438669916c6..4572eb0255796 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -111,7 +111,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables new Knowledge Base Entries features, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, /** * Enables the Managed User section inside the new user details flyout. From dba8f0b21cb613b9a885f82bf32d391566d5eeaa Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 23 Oct 2024 21:42:47 +0200 Subject: [PATCH 22/41] cleanup --- .../functional_test_runner/lib/config/schema.ts | 2 +- .../impl/capabilities/index.ts | 2 +- .../use_knowledge_base_status.test.tsx | 1 + .../knowledge_base/use_knowledge_base_status.tsx | 3 ++- .../index_entry_editor.tsx | 13 +------------ .../knowledge_base/get_knowledge_base_status.ts | 14 ++++---------- .../common/experimental_features.ts | 2 +- .../configs/ess.config.ts | 4 ++++ .../entries/trial_license_complete_tier/entries.ts | 2 +- 9 files changed, 16 insertions(+), 27 deletions(-) diff --git a/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts b/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts index dba2174003c7f..bb240a0416479 100644 --- a/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts +++ b/packages/kbn-test/src/functional_test_runner/lib/config/schema.ts @@ -145,7 +145,7 @@ export const schema = Joi.object() grep: Joi.string(), invert: Joi.boolean().default(false), slow: Joi.number().default(30000), - timeout: Joi.number().default(INSPECTING ? 360000 * 100 : 360000 * 2), + timeout: Joi.number().default(INSPECTING ? 360000 * 100 : 360000), ui: Joi.string().default('bdd'), }) .default(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 54c24f6ce7b8f..c1c101fd74cd8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, assistantModelEvaluation: false, }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx index 80ce3d27d8dcb..83073b5770ba0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx @@ -34,6 +34,7 @@ const statusResponse = { elser_exists: true, index_exists: true, pipeline_exists: true, + security_labs_exists: true, }; const http = { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index 75e78f2a06948..60d714e3272d5 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx @@ -91,7 +91,8 @@ export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undef (kbStatus?.elser_exists && kbStatus?.security_labs_exists && kbStatus?.index_exists && - kbStatus?.pipeline_exists) ?? + kbStatus?.pipeline_exists && + !kbStatus.is_setup_in_progress) ?? false ); }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx index 1d5611d3b0373..6a8a266e5e026 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx @@ -96,12 +96,6 @@ export const IndexEntryEditor: React.FC = React.memo( })); }, [dataViews]); - const indexExists = useAsync(async () => { - if (!entry?.index?.length) return true; - - return !!(await dataViews.getExistingIndices([entry.index])).length; - }, [entry?.index]); - const indexFields = useAsync( async () => dataViews.getFieldsForWildcard({ @@ -257,12 +251,7 @@ export const IndexEntryEditor: React.FC = React.memo( fullWidth /> - {"Index doesn't exist"}} - > + { const es = getService('es'); const ml = getService('ml') as ReturnType; - describe('@ess Basic Security AI Assistant Knowledge Base Entries', () => { + describe.skip('@ess Basic Security AI Assistant Knowledge Base Entries', () => { before(async () => { await installTinyElser(ml); await setupKnowledgeBase(supertest, log); From 060f2873a955b0b67ee48be16520bcb0558462ef Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 24 Oct 2024 00:55:16 +0200 Subject: [PATCH 23/41] fix --- .../api/knowledge_base/use_knowledge_base_status.tsx | 9 +-------- .../impl/assistant/chat_send/use_chat_send.tsx | 6 +----- .../ai_assistant_data_clients/knowledge_base/index.ts | 6 +++++- .../server/lib/data_stream/documents_data_writer.ts | 4 ++-- .../routes/knowledge_base/get_knowledge_base_status.ts | 6 ++---- .../server/routes/knowledge_base/post_knowledge_base.ts | 2 +- 6 files changed, 12 insertions(+), 21 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index 60d714e3272d5..dc5d9c6116aea 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx @@ -87,12 +87,5 @@ export const useInvalidateKnowledgeBaseStatus = () => { * @param kbStatus ReadKnowledgeBaseResponse */ export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undefined): boolean => { - return ( - (kbStatus?.elser_exists && - kbStatus?.security_labs_exists && - kbStatus?.index_exists && - kbStatus?.pipeline_exists && - !kbStatus.is_setup_in_progress) ?? - false - ); + return (kbStatus?.elser_exists && kbStatus?.index_exists && kbStatus?.pipeline_exists) ?? false; }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index e0a6717a72ed3..7d335b76bd297 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -59,11 +59,7 @@ export const useChatSend = ({ const { clearConversation, removeLastMessage } = useConversation(); const { data: kbStatus } = useKnowledgeBaseStatus({ http }); const isSetupComplete = - kbStatus?.elser_exists && - kbStatus?.index_exists && - kbStatus?.pipeline_exists && - kbStatus?.security_labs_exists && - !kbStatus?.is_setup_in_progress; + kbStatus?.elser_exists && kbStatus?.index_exists && kbStatus?.pipeline_exists; // Handles sending latest user prompt to API const handleSendMessage = useCallback( diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 4769321c88ea7..5535cfe964fb9 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -221,7 +221,11 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { inference_config: { service: 'elasticsearch', service_settings: { - num_allocations: 1, + adaptive_allocations: { + enabled: true, + min_number_of_allocations: 0, + max_number_of_allocations: 4, + }, num_threads: 1, model_id: elserId, }, diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 47c37162fc688..b4097d45da5ab 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -76,8 +76,8 @@ export class DocumentsDataWriter implements DocumentsDataWriter { body: await this.buildBulkOperations(params), }, { - // Increasing timout to 10min as KB docs were failing to load after 30s - requestTimeout: 600000, + // Increasing timeout to 2min as KB docs were failing to load after 30s + requestTimeout: 120000, } ); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 89ce20ee4ef63..a2db80356b75e 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -67,14 +67,12 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); const body: ReadKnowledgeBaseResponse = { - elser_exists: modelExists, + elser_exists: modelExists && isModelDeployed, index_exists: indexExists, is_setup_in_progress: kbDataClient.isSetupInProgress, is_setup_available: setupAvailable, pipeline_exists: pipelineExists, - security_labs_exists: v2KnowledgeBaseEnabled - ? isModelDeployed && securityLabsExists - : true, + security_labs_exists: v2KnowledgeBaseEnabled ? securityLabsExists : true, }; return response.ok({ body }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index cbb3ce1c10a6c..96317da303ac1 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -20,7 +20,7 @@ import { isV2KnowledgeBaseEnabled } 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 -const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 20 minutes +const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes /** * Load Knowledge Base index, pipeline, and resources (collection of documents) From d329ca736a7f3626cfe7ab511d0343a3dd03a41d Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 24 Oct 2024 01:12:40 +0200 Subject: [PATCH 24/41] fix --- .../api/knowledge_base/use_knowledge_base_status.tsx | 8 +++++++- .../impl/assistant/chat_send/use_chat_send.tsx | 5 ++++- .../routes/knowledge_base/get_knowledge_base_status.ts | 5 +++-- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index dc5d9c6116aea..75e78f2a06948 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx @@ -87,5 +87,11 @@ export const useInvalidateKnowledgeBaseStatus = () => { * @param kbStatus ReadKnowledgeBaseResponse */ export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undefined): boolean => { - return (kbStatus?.elser_exists && kbStatus?.index_exists && kbStatus?.pipeline_exists) ?? false; + return ( + (kbStatus?.elser_exists && + kbStatus?.security_labs_exists && + kbStatus?.index_exists && + kbStatus?.pipeline_exists) ?? + false + ); }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 7d335b76bd297..4ea376518b5a7 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -59,7 +59,10 @@ export const useChatSend = ({ const { clearConversation, removeLastMessage } = useConversation(); const { data: kbStatus } = useKnowledgeBaseStatus({ http }); const isSetupComplete = - kbStatus?.elser_exists && kbStatus?.index_exists && kbStatus?.pipeline_exists; + kbStatus?.elser_exists && + kbStatus?.index_exists && + kbStatus?.pipeline_exists && + kbStatus?.security_labs_exists; // Handles sending latest user prompt to API const handleSendMessage = useCallback( diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index a2db80356b75e..2cc7373bf5b21 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -64,10 +64,11 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter const modelExists = await kbDataClient.isModelInstalled(); const setupAvailable = await kbDataClient.isSetupAvailable(); const isModelDeployed = await kbDataClient.isModelDeployed(); - const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); + const securityLabsExists = + isModelDeployed && (await kbDataClient.isSecurityLabsDocsLoaded()); const body: ReadKnowledgeBaseResponse = { - elser_exists: modelExists && isModelDeployed, + elser_exists: modelExists, index_exists: indexExists, is_setup_in_progress: kbDataClient.isSetupInProgress, is_setup_available: setupAvailable, From 6f2bd7300b8d37af794a56053375c5ac6e623ab7 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 24 Oct 2024 01:15:15 +0200 Subject: [PATCH 25/41] flags --- .../kbn-elastic-assistant-common/impl/capabilities/index.ts | 2 +- .../plugins/security_solution/common/experimental_features.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index c1c101fd74cd8..54c24f6ce7b8f 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, assistantModelEvaluation: false, }); diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 5e438669916c6..4572eb0255796 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -111,7 +111,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables new Knowledge Base Entries features, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: false, + assistantKnowledgeBaseByDefault: true, /** * Enables the Managed User section inside the new user details flyout. From 4e28c67510fa2d92031f8301c1cc166cd59d5eed Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 24 Oct 2024 09:03:54 +0200 Subject: [PATCH 26/41] fix --- .../knowledge_base/helpers.ts | 2 -- .../knowledge_base/index.ts | 23 +++++++++++-------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index a7c406011670f..3f3fc23f55170 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -37,7 +37,6 @@ export const isModelAlreadyExistsError = (error: Error) => { export const getKBVectorSearchQuery = ({ filter, kbResource, - modelId, query, required, user, @@ -45,7 +44,6 @@ export const getKBVectorSearchQuery = ({ }: { filter?: QueryDslQueryContainer | undefined; kbResource?: string | undefined; - modelId: string; query: string; required?: boolean | undefined; user: AuthenticatedUser; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 5535cfe964fb9..875ffd36ebcac 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -224,7 +224,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { adaptive_allocations: { enabled: true, min_number_of_allocations: 0, - max_number_of_allocations: 4, + max_number_of_allocations: 8, }, num_threads: 1, model_id: elserId, @@ -392,7 +392,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { user: authenticatedUser, entry: { type: DocumentEntryType.value, - name: 'unknown', + // TODO: parse title from document + name: '', text: doc.pageContent, kbResource: doc.metadata.kbResource ?? 'unknown', required: doc.metadata.required ?? false, @@ -451,12 +452,16 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { * Returns if Security Labs KB docs have been loaded */ public isSecurityLabsDocsLoaded = async (): Promise => { - const securityLabsDocs = await this.getKnowledgeBaseDocumentEntries({ - query: '', - kbResource: SECURITY_LABS_RESOURCE, - required: false, - }); - return securityLabsDocs.length > 0; + try { + const securityLabsDocs = await this.getKnowledgeBaseDocumentEntries({ + query: '', + kbResource: SECURITY_LABS_RESOURCE, + required: false, + }); + return !!securityLabsDocs.length; + } catch (e) { + return false; + } }; /** @@ -481,12 +486,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } const esClient = await this.options.elasticsearchClientPromise; - const modelId = await this.options.getElserId(); const vectorSearchQuery = getKBVectorSearchQuery({ filter, kbResource, - modelId, query, required, user, From 461eb7979325409fef4e5385b49bd078978befe6 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 24 Oct 2024 14:21:26 +0200 Subject: [PATCH 27/41] test --- .../server/ai_assistant_data_clients/knowledge_base/index.ts | 3 ++- .../entries/trial_license_complete_tier/entries.ts | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 875ffd36ebcac..c0e06836dd589 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -393,7 +393,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { entry: { type: DocumentEntryType.value, // TODO: parse title from document - name: '', + // Disabled so we don't copy that value with copy_to + // name: 'unknown', text: doc.pageContent, kbResource: doc.metadata.kbResource ?? 'unknown', required: doc.metadata.required ?? false, diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts index 79cd8232d4eb9..7cd44a21ce236 100644 --- a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts @@ -27,7 +27,7 @@ export default ({ getService }: FtrProviderContext) => { const es = getService('es'); const ml = getService('ml') as ReturnType; - describe.skip('@ess Basic Security AI Assistant Knowledge Base Entries', () => { + describe('@ess Basic Security AI Assistant Knowledge Base Entries', () => { before(async () => { await installTinyElser(ml); await setupKnowledgeBase(supertest, log); From 59226a578bbf96605e78a8c567c2a52d77444395 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Thu, 24 Oct 2024 15:34:57 +0200 Subject: [PATCH 28/41] fix --- .../knowledge_base/field_maps_configuration.ts | 1 - .../server/ai_assistant_data_clients/knowledge_base/index.ts | 4 +--- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index 51c816eaa8a02..6bb0a700c1015 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -142,7 +142,6 @@ export const knowledgeBaseFieldMapV2: FieldMap = { type: 'keyword', array: false, required: false, - copy_to: 'semantic_text', }, // Discriminator: 'document' | 'index' type: { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index c0e06836dd589..baf7de33b4614 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -392,9 +392,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { user: authenticatedUser, entry: { type: DocumentEntryType.value, - // TODO: parse title from document - // Disabled so we don't copy that value with copy_to - // name: 'unknown', + name: 'unknown', text: doc.pageContent, kbResource: doc.metadata.kbResource ?? 'unknown', required: doc.metadata.required ?? false, From 701a54a7fd5b24c0cc73f45efd075df1e0dc40b3 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 25 Oct 2024 13:59:22 +0200 Subject: [PATCH 29/41] fix --- .../knowledge_base/index.ts | 9 ++++++--- .../content_loaders/security_labs_loader.ts | 18 ++++-------------- .../get_knowledge_base_status.ts | 4 ++++ .../knowledge_base/post_knowledge_base.ts | 2 +- .../server/routes/request_context_factory.ts | 6 ++++++ .../plugins/elastic_assistant/server/types.ts | 1 + 6 files changed, 22 insertions(+), 18 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 4769321c88ea7..c0c3620d0b442 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -221,7 +221,11 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { inference_config: { service: 'elasticsearch', service_settings: { - num_allocations: 1, + adaptive_allocations: { + enabled: true, + min_number_of_allocations: 0, + max_number_of_allocations: 8, + }, num_threads: 1, model_id: elserId, }, @@ -339,7 +343,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { const labsDocsLoaded = await this.isSecurityLabsDocsLoaded(); if (!labsDocsLoaded) { this.options.logger.debug(`Loading Security Labs KB docs...`); - await loadSecurityLabs(this, this.options.logger); + loadSecurityLabs(this, this.options.logger); } else { this.options.logger.debug(`Security Labs Knowledge Base docs already loaded!`); } @@ -349,7 +353,6 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); throw new Error(`Error setting up Knowledge Base: ${e.message}`); } - this.options.setIsKBSetupInProgress(false); }; /** diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts index facd89b3ada55..0a5be262bc7ba 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts @@ -11,9 +11,6 @@ import { TextLoader } from 'langchain/document_loaders/fs/text'; import { resolve } from 'path'; import { Document } from 'langchain/document'; import { Metadata } from '@kbn/elastic-assistant-common'; -import pMap from 'p-map'; - -import { chunk } from 'lodash'; import { addRequiredKbResourceMetadata } from './add_required_kb_resource_metadata'; import { SECURITY_LABS_RESOURCE } from '../../../routes/knowledge_base/constants'; import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; @@ -44,17 +41,10 @@ export const loadSecurityLabs = async ( logger.info(`Loading ${docs.length} Security Labs docs into the Knowledge Base`); - const response = ( - await pMap( - chunk(docs, 30), - (partialDocs) => - kbDataClient.addKnowledgeBaseDocuments({ - documents: partialDocs, - global: true, - }), - { concurrency: 1 } - ) - ).flat(); + const response = await kbDataClient.addKnowledgeBaseDocuments({ + documents: docs, + global: true, + }); logger.info(`Loaded ${response?.length ?? 0} Security Labs docs into the Knowledge Base`); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 89ce20ee4ef63..911612def964b 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -66,6 +66,10 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter const isModelDeployed = await kbDataClient.isModelDeployed(); const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); + if (securityLabsExists && kbDataClient.isSetupInProgress) { + assistantContext.setIsKBSetupInProgress(false); + } + const body: ReadKnowledgeBaseResponse = { elser_exists: modelExists, index_exists: indexExists, diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index cbb3ce1c10a6c..96317da303ac1 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -20,7 +20,7 @@ import { isV2KnowledgeBaseEnabled } 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 -const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 20 minutes +const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes /** * Load Knowledge Base index, pipeline, and resources (collection of documents) diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index eeb1a5564d1cf..29e0ffd16d4e0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -145,6 +145,12 @@ export class RequestContextFactory implements IRequestContextFactory { currentUser, }); }), + + setIsKBSetupInProgress: memoize(() => { + if (this.assistantService.getIsKBSetupInProgress()) { + return this.assistantService.setIsKBSetupInProgress(false); + } + }), }; } } diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index e84b97ab43d7a..a864adfbfa870 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -131,6 +131,7 @@ export interface ElasticAssistantApiRequestHandlerContext { getAttackDiscoveryDataClient: () => Promise; getAIAssistantPromptsDataClient: () => Promise; getAIAssistantAnonymizationFieldsDataClient: () => Promise; + setIsKBSetupInProgress: (isInProgress: boolean) => void; inference: InferenceServerStart; telemetry: AnalyticsServiceSetup; } From 574f8342e0c1aeb33f89a88b13993c37bc5a11dc Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Fri, 25 Oct 2024 16:21:48 +0200 Subject: [PATCH 30/41] fix --- .../entries/trial_license_complete_tier/configs/ess.config.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts index b851e59ee0a9d..7954db769a6d5 100644 --- a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts @@ -48,6 +48,7 @@ export default async function ({ readConfigFile }: FtrConfigProviderContext) { esTestCluster: { ...functionalConfig.get('esTestCluster'), ssl: false, + esJavaOpts: '-Xms4g -Xmx4g', }, mochaOpts: { ...functionalConfig.get('mochaOpts'), From d2aa921e5b68113c933042abb637843a3f4204c3 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 28 Oct 2024 07:33:28 +0100 Subject: [PATCH 31/41] fix --- .../knowledge_base/crud_kb_route.gen.ts | 1 + .../use_knowledge_base_status.tsx | 18 +-- .../knowledge_base_settings.tsx | 4 +- .../index.tsx | 9 +- .../use_knowledge_base_table.tsx | 33 +++++- .../create_knowledge_base_entry.ts | 3 - .../knowledge_base/helpers.ts | 2 - .../knowledge_base/index.ts | 103 ++++++++++++++++-- .../knowledge_base/ingest_pipeline.ts | 1 - .../content_loaders/security_labs_loader.ts | 31 +++++- .../server/routes/knowledge_base/constants.ts | 1 + .../get_knowledge_base_status.ts | 7 +- .../plugins/elastic_assistant/server/types.ts | 1 - 13 files changed, 168 insertions(+), 46 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index fd599f5798cdc..aad215021da81 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -81,4 +81,5 @@ export const ReadKnowledgeBaseResponse = z.object({ is_setup_in_progress: z.boolean().optional(), pipeline_exists: z.boolean().optional(), security_labs_exists: z.boolean().optional(), + user_data_exists: z.boolean().optional(), }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index 75e78f2a06948..863595fc35cce 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx @@ -86,12 +86,12 @@ export const useInvalidateKnowledgeBaseStatus = () => { * * @param kbStatus ReadKnowledgeBaseResponse */ -export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undefined): boolean => { - return ( - (kbStatus?.elser_exists && - kbStatus?.security_labs_exists && - kbStatus?.index_exists && - kbStatus?.pipeline_exists) ?? - false - ); -}; +export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undefined): boolean => + (kbStatus?.elser_exists && + kbStatus?.index_exists && + kbStatus?.pipeline_exists && + // Allows to use UI while importing Security Labs docs + (kbStatus?.security_labs_exists || + kbStatus?.is_setup_in_progress || + kbStatus?.user_data_exists)) ?? + false; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx index a46ba652574f6..7041bf909601f 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx @@ -53,9 +53,9 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const isSecurityLabsEnabled = kbStatus?.security_labs_exists ?? false; const isKnowledgeBaseSetup = (isElserEnabled && - isSecurityLabsEnabled && kbStatus?.index_exists && - kbStatus?.pipeline_exists) ?? + kbStatus?.pipeline_exists && + (isSecurityLabsEnabled || kbStatus?.user_data_exists)) ?? false; const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false; const isSetupAvailable = kbStatus?.is_setup_available ?? false; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index e3a86c62d1222..33617314e0110 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -177,6 +177,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d const columns = useMemo( () => getColumns({ + isKbSetupInProgress: kbStatus?.is_setup_in_progress ?? false, isDeleteEnabled: (entry: KnowledgeBaseEntryResponse) => { return ( !isSystemEntry(entry) && (isGlobalEntry(entry) ? hasManageGlobalKnowledgeBase : true) @@ -197,7 +198,13 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d openFlyout(); }, }), - [entries.data, getColumns, hasManageGlobalKnowledgeBase, openFlyout] + [ + entries.data, + getColumns, + hasManageGlobalKnowledgeBase, + kbStatus?.is_setup_in_progress, + openFlyout, + ] ); // Refresh button diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx index d27853f6e8625..227802604f93c 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx @@ -5,7 +5,14 @@ * 2.0. */ -import { EuiAvatar, EuiBadge, EuiBasicTableColumn, EuiIcon, EuiText } from '@elastic/eui'; +import { + EuiAvatar, + EuiBadge, + EuiBasicTableColumn, + EuiIcon, + EuiText, + EuiLoadingSpinner, +} from '@elastic/eui'; import { css } from '@emotion/react'; import React, { useCallback, useMemo } from 'react'; import { FormattedDate } from '@kbn/i18n-react'; @@ -101,11 +108,13 @@ export const useKnowledgeBaseTable = () => { isEditEnabled, onDeleteActionClicked, onEditActionClicked, + isKbSetupInProgress, }: { isDeleteEnabled: (entry: KnowledgeBaseEntryResponse) => boolean; isEditEnabled: (entry: KnowledgeBaseEntryResponse) => boolean; onDeleteActionClicked: (entry: KnowledgeBaseEntryResponse) => void; onEditActionClicked: (entry: KnowledgeBaseEntryResponse) => void; + isKbSetupInProgress: boolean; }): Array> => { return [ { @@ -136,11 +145,23 @@ export const useKnowledgeBaseTable = () => { { name: i18n.COLUMN_ENTRIES, render: (entry: KnowledgeBaseEntryResponse) => { - return isSystemEntry(entry) - ? entry.text - : entry.type === DocumentEntryType.value - ? '1' - : '-'; + return isSystemEntry(entry) ? ( + <> + {`${entry.text}`} + {isKbSetupInProgress && ( + + )} + + ) : entry.type === DocumentEntryType.value ? ( + '1' + ) : ( + '-' + ); }, }, { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts index 838227b7c8e15..09bb5b291ef9a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts @@ -204,8 +204,6 @@ export const transformToCreateSchema = ({ source: entry.source, text: entry.text, semantic_text: entry.text, - // deprecated - vector: undefined, }; }; @@ -247,7 +245,6 @@ export const transformToLegacyCreateSchema = ({ }, ], ...entry, - // deprecated vector: undefined, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index a7c406011670f..3f3fc23f55170 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -37,7 +37,6 @@ export const isModelAlreadyExistsError = (error: Error) => { export const getKBVectorSearchQuery = ({ filter, kbResource, - modelId, query, required, user, @@ -45,7 +44,6 @@ export const getKBVectorSearchQuery = ({ }: { filter?: QueryDslQueryContainer | undefined; kbResource?: string | undefined; - modelId: string; query: string; required?: boolean | undefined; user: AuthenticatedUser; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index c0c3620d0b442..8b610e7cb5e5a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -40,6 +40,7 @@ import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY, SECURITY_LABS_RESOURCE, + USER_RESOURCE, } from '../../routes/knowledge_base/constants'; import { getKBVectorSearchQuery, @@ -47,7 +48,10 @@ import { isModelAlreadyExistsError, } from './helpers'; import { getKBUserFilter } from '../../routes/knowledge_base/entries/utils'; -import { loadSecurityLabs } from '../../lib/langchain/content_loaders/security_labs_loader'; +import { + loadSecurityLabs, + getSecurityLabsDocsCount, +} from '../../lib/langchain/content_loaders/security_labs_loader'; import { ASSISTANT_ELSER_INFERENCE_ID } from './field_maps_configuration'; /** @@ -286,8 +290,22 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { `Removed ${legacyESQL?.total} ESQL knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.` ); } + // Delete any existing Security Labs content + const securityLabsDocs = await esClient.deleteByQuery({ + index: this.indexTemplateAndPattern.alias, + query: { + bool: { + must: [{ terms: { kb_resource: [SECURITY_LABS_RESOURCE] } }], + }, + }, + }); + if (securityLabsDocs?.total) { + this.options.logger.info( + `Removed ${securityLabsDocs?.total} Security Labs knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.` + ); + } } catch (e) { - this.options.logger.info('No legacy ESQL knowledge base docs to delete'); + this.options.logger.info('No legacy ESQL or Security Labs knowledge base docs to delete'); } } @@ -343,7 +361,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { const labsDocsLoaded = await this.isSecurityLabsDocsLoaded(); if (!labsDocsLoaded) { this.options.logger.debug(`Loading Security Labs KB docs...`); - loadSecurityLabs(this, this.options.logger); + await loadSecurityLabs(this, this.options.logger); } else { this.options.logger.debug(`Security Labs Knowledge Base docs already loaded!`); } @@ -352,6 +370,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { this.options.setIsKBSetupInProgress(false); this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); throw new Error(`Error setting up Knowledge Base: ${e.message}`); + } finally { + this.options.setIsKBSetupInProgress(false); } }; @@ -447,15 +467,76 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }; /** - * Returns if Security Labs KB docs have been loaded + * Returns if user's KB docs exists + */ + + public isUserDataExists = async (): Promise => { + const user = this.options.currentUser; + if (user == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } + + const esClient = await this.options.elasticsearchClientPromise; + + try { + const vectorSearchQuery = getKBVectorSearchQuery({ + kbResource: USER_RESOURCE, + query: '', + required: false, + user, + v2KnowledgeBaseEnabled: this.options.v2KnowledgeBaseEnabled, + }); + + const result = await esClient.search({ + index: this.indexTemplateAndPattern.alias, + size: 0, + query: vectorSearchQuery, + track_total_hits: true, + }); + + return !!result.hits?.total; + } catch (e) { + return false; + } + }; + + /** + * Returns if allSecurity Labs KB docs have been loaded */ public isSecurityLabsDocsLoaded = async (): Promise => { - const securityLabsDocs = await this.getKnowledgeBaseDocumentEntries({ - query: '', - kbResource: SECURITY_LABS_RESOURCE, - required: false, - }); - return securityLabsDocs.length > 0; + const user = this.options.currentUser; + if (user == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } + + const expectedDocsCount = await getSecurityLabsDocsCount(); + + const esClient = await this.options.elasticsearchClientPromise; + + try { + const vectorSearchQuery = getKBVectorSearchQuery({ + kbResource: SECURITY_LABS_RESOURCE, + query: '', + required: false, + user, + v2KnowledgeBaseEnabled: this.options.v2KnowledgeBaseEnabled, + }); + + const result = await esClient.search({ + index: this.indexTemplateAndPattern.alias, + size: 0, + query: vectorSearchQuery, + track_total_hits: true, + }); + + return result.hits?.total === expectedDocsCount; + } catch (e) { + return false; + } }; /** @@ -480,12 +561,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } const esClient = await this.options.elasticsearchClientPromise; - const modelId = await this.options.getElserId(); const vectorSearchQuery = getKBVectorSearchQuery({ filter, kbResource, - modelId, query, required, user, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts index ec86885f26164..8f459848af420 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts @@ -5,7 +5,6 @@ * 2.0. */ -// TODO: Ensure old pipeline is updated/replaced export const knowledgeBaseIngestPipeline = ({ id, modelId, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts index 0a5be262bc7ba..bbb44c12019ea 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts @@ -5,12 +5,14 @@ * 2.0. */ +import globby from 'globby'; import { Logger } from '@kbn/core/server'; import { DirectoryLoader } from 'langchain/document_loaders/fs/directory'; import { TextLoader } from 'langchain/document_loaders/fs/text'; import { resolve } from 'path'; import { Document } from 'langchain/document'; import { Metadata } from '@kbn/elastic-assistant-common'; +import pMap from 'p-map'; import { addRequiredKbResourceMetadata } from './add_required_kb_resource_metadata'; import { SECURITY_LABS_RESOURCE } from '../../../routes/knowledge_base/constants'; import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; @@ -41,10 +43,22 @@ export const loadSecurityLabs = async ( logger.info(`Loading ${docs.length} Security Labs docs into the Knowledge Base`); - const response = await kbDataClient.addKnowledgeBaseDocuments({ - documents: docs, - global: true, - }); + /** + * Ingest Security Labs docs into the Knowledge Base one by one to avoid blocking + * Inference Endpoint for too long + */ + + const response = ( + await pMap( + docs, + (singleDoc) => + kbDataClient.addKnowledgeBaseDocuments({ + documents: [singleDoc], + global: true, + }), + { concurrency: 1 } + ) + ).flat(); logger.info(`Loaded ${response?.length ?? 0} Security Labs docs into the Knowledge Base`); @@ -54,3 +68,12 @@ export const loadSecurityLabs = async ( return false; } }; + +export const getSecurityLabsDocsCount = async (): Promise => { + try { + return (await globby(`${resolve(__dirname, '../../../knowledge_base/security_labs')}/**/*.md`)) + .length; + } catch (e) { + return 0; + } +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/constants.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/constants.ts index 89970611df0e9..8bf17027e751e 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/constants.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/constants.ts @@ -12,3 +12,4 @@ export const KNOWLEDGE_BASE_INGEST_PIPELINE = '.kibana-elastic-ai-assistant-kb-i export const ESQL_DOCS_LOADED_QUERY = 'You can chain processing commands, separated by a pipe character: `|`.'; export const SECURITY_LABS_RESOURCE = 'security_labs'; +export const USER_RESOURCE = 'user'; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 911612def964b..b757fab6e1956 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -64,11 +64,7 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter const modelExists = await kbDataClient.isModelInstalled(); const setupAvailable = await kbDataClient.isSetupAvailable(); const isModelDeployed = await kbDataClient.isModelDeployed(); - const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); - - if (securityLabsExists && kbDataClient.isSetupInProgress) { - assistantContext.setIsKBSetupInProgress(false); - } + const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); // && !kbDataClient.isSetupInProgress; const body: ReadKnowledgeBaseResponse = { elser_exists: modelExists, @@ -79,6 +75,7 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter security_labs_exists: v2KnowledgeBaseEnabled ? isModelDeployed && securityLabsExists : true, + user_data_exists: await kbDataClient.isUserDataExists(), }; return response.ok({ body }); diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index a864adfbfa870..e84b97ab43d7a 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -131,7 +131,6 @@ export interface ElasticAssistantApiRequestHandlerContext { getAttackDiscoveryDataClient: () => Promise; getAIAssistantPromptsDataClient: () => Promise; getAIAssistantAnonymizationFieldsDataClient: () => Promise; - setIsKBSetupInProgress: (isInProgress: boolean) => void; inference: InferenceServerStart; telemetry: AnalyticsServiceSetup; } From 211586a9e81d9b606627a760a5d960d47f1a122e Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 28 Oct 2024 14:51:12 +0100 Subject: [PATCH 32/41] fix --- .../entries/use_knowledge_base_entries.ts | 3 ++ .../use_knowledge_base_status.tsx | 2 + .../index.tsx | 1 + .../use_knowledge_base_table.tsx | 7 +++- .../setup_knowledge_base_button.tsx | 22 +++++++++- .../knowledge_base/index.ts | 41 +++++++++++-------- .../get_knowledge_base_status.ts | 2 +- 7 files changed, 55 insertions(+), 23 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts index b41119779b21d..6f67e5e71010f 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts @@ -24,6 +24,7 @@ export interface UseKnowledgeBaseEntriesParams { signal?: AbortSignal | undefined; toasts?: IToasts; enabled?: boolean; // For disabling if FF is off + isRefetching?: boolean; // For enabling polling } const defaultQuery: FindKnowledgeBaseEntriesRequestQuery = { @@ -56,6 +57,7 @@ export const useKnowledgeBaseEntries = ({ signal, toasts, enabled = false, + isRefetching = false, }: UseKnowledgeBaseEntriesParams) => useQuery( KNOWLEDGE_BASE_ENTRY_QUERY_KEY, @@ -73,6 +75,7 @@ export const useKnowledgeBaseEntries = ({ enabled, keepPreviousData: true, initialData: { page: 1, perPage: 100, total: 0, data: [] }, + refetchInterval: isRefetching ? 5000 : false, onError: (error: IHttpFetchError) => { if (error.name !== 'AbortError') { toasts?.addError(error, { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index 863595fc35cce..2b04b45c2359f 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx @@ -45,6 +45,8 @@ export const useKnowledgeBaseStatus = ({ { retry: false, keepPreviousData: true, + // Polling interval for Knowledge Base setup in progress + refetchInterval: (data) => (data?.is_setup_in_progress ? 5000 : false), // Deprecated, hoist to `queryCache` w/in `QueryClient. See: https://stackoverflow.com/a/76961109 onError: (error: IHttpFetchError) => { if (error.name !== 'AbortError') { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index 33617314e0110..3a3e822a5a3c0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -172,6 +172,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d http, toasts, enabled: enableKnowledgeBaseByDefault, + isRefetching: kbStatus?.is_setup_in_progress, }); const { getColumns } = useKnowledgeBaseTable(); const columns = useMemo( diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx index 227802604f93c..e079ef4aeae98 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx @@ -29,6 +29,7 @@ import * as i18n from './translations'; import { BadgesColumn } from '../../assistant/common/components/assistant_settings_management/badges'; import { useInlineActions } from '../../assistant/common/components/assistant_settings_management/inline_actions'; import { isSystemEntry } from './helpers'; +import { SetupKnowledgeBaseButton } from '../setup_knowledge_base_button'; const AuthorColumn = ({ entry }: { entry: KnowledgeBaseEntryResponse }) => { const { userProfileService } = useAssistantContext(); @@ -148,13 +149,15 @@ export const useKnowledgeBaseTable = () => { return isSystemEntry(entry) ? ( <> {`${entry.text}`} - {isKbSetupInProgress && ( + {isKbSetupInProgress ? ( + ) : ( + )} ) : entry.type === DocumentEntryType.value ? ( diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx index d697fc7120d01..948e45232028c 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx @@ -6,15 +6,16 @@ */ import React, { useCallback } from 'react'; -import { EuiButton, EuiButtonEmpty, EuiToolTip } from '@elastic/eui'; +import { EuiButton, EuiButtonIcon, EuiButtonEmpty, EuiToolTip } from '@elastic/eui'; import { i18n } from '@kbn/i18n'; +import { css } from '@emotion/react'; import { useAssistantContext } from '../..'; import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base'; import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status'; interface Props { - display?: 'mini'; + display?: 'mini' | 'refresh'; } /** @@ -48,6 +49,23 @@ export const SetupKnowledgeBaseButton: React.FC = React.memo(({ display } }) : undefined; + if (display === 'refresh') { + return ( + + ); + } + return ( {display === 'mini' ? ( diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 8b610e7cb5e5a..f26496f4a1d7c 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -8,6 +8,7 @@ import { MlTrainedModelDeploymentNodesStats, MlTrainedModelStats, + SearchTotalHits, } from '@elastic/elasticsearch/lib/api/types'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import type { KibanaRequest } from '@kbn/core-http-server'; @@ -176,35 +177,39 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { this.options.logger.debug(`Checking if ELSER model '${elserId}' is deployed...`); try { - const esClient = await this.options.elasticsearchClientPromise; - const getResponse = await esClient.ml.getTrainedModelsStats({ - model_id: elserId, - }); + if (this.isV2KnowledgeBaseEnabled) { + return await this.isInferenceEndpointExists(); + } else { + const esClient = await this.options.elasticsearchClientPromise; + const getResponse = await esClient.ml.getTrainedModelsStats({ + model_id: elserId, + }); - // For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986 - const isReadyESS = (stats: MlTrainedModelStats) => - stats.deployment_stats?.state === 'started' && - stats.deployment_stats?.allocation_status.state === 'fully_allocated'; + // For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986 + const isReadyESS = (stats: MlTrainedModelStats) => + stats.deployment_stats?.state === 'started' && + stats.deployment_stats?.allocation_status.state === 'fully_allocated'; - const isReadyServerless = (stats: MlTrainedModelStats) => - (stats.deployment_stats?.nodes as unknown as MlTrainedModelDeploymentNodesStats[]).some( - (node) => node.routing_state.routing_state === 'started' - ); + const isReadyServerless = (stats: MlTrainedModelStats) => + (stats.deployment_stats?.nodes as unknown as MlTrainedModelDeploymentNodesStats[]).some( + (node) => node.routing_state.routing_state === 'started' + ); - return getResponse.trained_model_stats.some( - (stats) => isReadyESS(stats) || isReadyServerless(stats) - ); + return getResponse.trained_model_stats.some( + (stats) => isReadyESS(stats) || isReadyServerless(stats) + ); + } } catch (e) { // Returns 404 if it doesn't exist return false; } }; - public isInferenceEndpointExists = async () => { + public isInferenceEndpointExists = async (): Promise => { try { const esClient = await this.options.elasticsearchClientPromise; - return await esClient.inference.get({ + return await !!esClient.inference.get({ inference_id: ASSISTANT_ELSER_INFERENCE_ID, task_type: 'sparse_embedding', }); @@ -533,7 +538,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { track_total_hits: true, }); - return result.hits?.total === expectedDocsCount; + return (result.hits?.total as SearchTotalHits).value === expectedDocsCount; } catch (e) { return false; } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index b757fab6e1956..4617f2f33d733 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -64,7 +64,7 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter const modelExists = await kbDataClient.isModelInstalled(); const setupAvailable = await kbDataClient.isSetupAvailable(); const isModelDeployed = await kbDataClient.isModelDeployed(); - const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); // && !kbDataClient.isSetupInProgress; + const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); const body: ReadKnowledgeBaseResponse = { elser_exists: modelExists, From f86beae308920672a912c336663917a4bea43add Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Mon, 28 Oct 2024 22:53:13 +0100 Subject: [PATCH 33/41] PR comments --- .../knowledge_base/index.ts | 8 ++++++-- .../knowledge_base/get_knowledge_base_status.ts | 16 +++++++++++----- .../server/routes/request_context_factory.ts | 6 ------ 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index f26496f4a1d7c..39daed82a11b7 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -200,6 +200,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { ); } } catch (e) { + this.options.logger.error(`Error checking if ELSER model '${elserId}' is deployed: ${e}`); // Returns 404 if it doesn't exist return false; } @@ -209,11 +210,14 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { try { const esClient = await this.options.elasticsearchClientPromise; - return await !!esClient.inference.get({ + return !!(await esClient.inference.get({ inference_id: ASSISTANT_ELSER_INFERENCE_ID, task_type: 'sparse_embedding', - }); + })); } catch (error) { + this.options.logger.error( + `Error checking if Inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} exists: ${error}` + ); return false; } }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 4617f2f33d733..a67d9ed5bb71f 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -64,7 +64,6 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter const modelExists = await kbDataClient.isModelInstalled(); const setupAvailable = await kbDataClient.isSetupAvailable(); const isModelDeployed = await kbDataClient.isModelDeployed(); - const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); const body: ReadKnowledgeBaseResponse = { elser_exists: modelExists, @@ -72,12 +71,19 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter is_setup_in_progress: kbDataClient.isSetupInProgress, is_setup_available: setupAvailable, pipeline_exists: pipelineExists, - security_labs_exists: v2KnowledgeBaseEnabled - ? isModelDeployed && securityLabsExists - : true, - user_data_exists: await kbDataClient.isUserDataExists(), }; + if (v2KnowledgeBaseEnabled && indexExists && isModelDeployed) { + const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); + return response.ok({ + body: { + ...body, + security_labs_exists: v2KnowledgeBaseEnabled ? securityLabsExists : true, + user_data_exists: await kbDataClient.isUserDataExists(), + }, + }); + } + return response.ok({ body }); } catch (err) { logger.error(err); diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index 29e0ffd16d4e0..eeb1a5564d1cf 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -145,12 +145,6 @@ export class RequestContextFactory implements IRequestContextFactory { currentUser, }); }), - - setIsKBSetupInProgress: memoize(() => { - if (this.assistantService.getIsKBSetupInProgress()) { - return this.assistantService.setIsKBSetupInProgress(false); - } - }), }; } } From 90cf431abe0f3a3f7d4271e45291d9a5239bb07e Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Mon, 28 Oct 2024 23:18:57 +0000 Subject: [PATCH 34/41] [CI] Auto-commit changed files from 'yarn openapi:generate' --- .../impl/schemas/knowledge_base/crud_kb_route.gen.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index aad215021da81..fd599f5798cdc 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -81,5 +81,4 @@ export const ReadKnowledgeBaseResponse = z.object({ is_setup_in_progress: z.boolean().optional(), pipeline_exists: z.boolean().optional(), security_labs_exists: z.boolean().optional(), - user_data_exists: z.boolean().optional(), }); From ae220af1fbc91725065c21ccd833abdedea90c63 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 29 Oct 2024 02:03:21 +0100 Subject: [PATCH 35/41] update test --- .../impl/schemas/knowledge_base/crud_kb_route.gen.ts | 1 + .../kbn-elastic-assistant/impl/assistant/index.test.tsx | 6 ++++-- .../ai_assistant_data_clients/knowledge_base/index.ts | 3 ++- .../lib/langchain/content_loaders/security_labs_loader.ts | 5 +++-- .../routes/knowledge_base/get_knowledge_base_status.test.ts | 2 ++ 5 files changed, 12 insertions(+), 5 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index fd599f5798cdc..aad215021da81 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -81,4 +81,5 @@ export const ReadKnowledgeBaseResponse = z.object({ is_setup_in_progress: z.boolean().optional(), pipeline_exists: z.boolean().optional(), security_labs_exists: z.boolean().optional(), + user_data_exists: z.boolean().optional(), }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx index 1ef2db7b26c03..c2158508620f0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx @@ -142,7 +142,8 @@ describe('Assistant', () => { }); describe('persistent storage', () => { - it('should refetchCurrentUserConversations after settings save button click', async () => { + // Settings Modal is not available when `v2KnowledgeBaseEnabled` is true + it.skip('should refetchCurrentUserConversations after settings save button click', async () => { const chatSendSpy = jest.spyOn(all, 'useChatSend'); await renderAssistant(); @@ -188,7 +189,8 @@ describe('Assistant', () => { ); }); - it('should refetchCurrentUserConversations after settings save button click, but do not update convos when refetch returns bad results', async () => { + // Settings Modal is not available when `v2KnowledgeBaseEnabled` is true + it.skip('should refetchCurrentUserConversations after settings save button click, but do not update convos when refetch returns bad results', async () => { jest.mocked(useFetchCurrentUserConversations).mockReturnValue({ data: mockData, isLoading: false, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 2876c63e71127..8a5a0498f85ed 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -506,6 +506,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return !!result.hits?.total; } catch (e) { + this.options.logger.error(`Error checking if user's KB docs exist: ${e.message}`); return false; } }; @@ -521,7 +522,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { ); } - const expectedDocsCount = await getSecurityLabsDocsCount(); + const expectedDocsCount = await getSecurityLabsDocsCount({ logger: this.options.logger }); const esClient = await this.options.elasticsearchClientPromise; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts index bbb44c12019ea..f37e20df2bd98 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts @@ -69,11 +69,12 @@ export const loadSecurityLabs = async ( } }; -export const getSecurityLabsDocsCount = async (): Promise => { +export const getSecurityLabsDocsCount = async ({ logger }: { logger: Logger }): Promise => { try { return (await globby(`${resolve(__dirname, '../../../knowledge_base/security_labs')}/**/*.md`)) - .length; + ?.length; } catch (e) { + logger.error(`Failed to get Security Labs source docs count\n${e}`); return 0; } }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts index 6244599a2af27..b30e5ac3653ad 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts @@ -38,6 +38,7 @@ describe('Get Knowledge Base Status Route', () => { isModelDeployed: jest.fn().mockResolvedValue(true), isSetupInProgress: false, isSecurityLabsDocsLoaded: jest.fn().mockResolvedValue(true), + isUserDataExists: jest.fn().mockResolvedValue(true), }); getKnowledgeBaseStatusRoute(server.router); @@ -58,6 +59,7 @@ describe('Get Knowledge Base Status Route', () => { is_setup_available: true, pipeline_exists: true, security_labs_exists: true, + user_data_exists: true, }); }); }); From 9217d9ace0e75078a70e5fd5e73d0d5c96403926 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 29 Oct 2024 01:49:51 +0000 Subject: [PATCH 36/41] [CI] Auto-commit changed files from 'yarn openapi:generate' --- .../impl/schemas/knowledge_base/crud_kb_route.gen.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index aad215021da81..fd599f5798cdc 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -81,5 +81,4 @@ export const ReadKnowledgeBaseResponse = z.object({ is_setup_in_progress: z.boolean().optional(), pipeline_exists: z.boolean().optional(), security_labs_exists: z.boolean().optional(), - user_data_exists: z.boolean().optional(), }); From dec121d5dfee318344831440e318a61d8e1f8e57 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 29 Oct 2024 09:10:26 +0100 Subject: [PATCH 37/41] fix --- .../impl/schemas/knowledge_base/crud_kb_route.gen.ts | 1 + .../impl/schemas/knowledge_base/crud_kb_route.schema.yaml | 2 ++ .../server/ai_assistant_data_clients/knowledge_base/index.ts | 4 ++-- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index fd599f5798cdc..aad215021da81 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -81,4 +81,5 @@ export const ReadKnowledgeBaseResponse = z.object({ is_setup_in_progress: z.boolean().optional(), pipeline_exists: z.boolean().optional(), security_labs_exists: z.boolean().optional(), + user_data_exists: z.boolean().optional(), }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml index a61e98602ab40..0e0f1e9267916 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml @@ -78,6 +78,8 @@ paths: type: boolean security_labs_exists: type: boolean + user_data_exists: + type: boolean 400: description: Generic Error content: diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 8a5a0498f85ed..7f29ce214cf69 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -544,13 +544,13 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { const existingDocs = (result.hits?.total as SearchTotalHits).value; if (existingDocs !== expectedDocsCount) { - this.options.logger.error( + this.options.logger.debug( `Security Labs docs are not loaded, existing docs: ${existingDocs}, expected docs: ${expectedDocsCount}` ); } return existingDocs === expectedDocsCount; } catch (e) { - this.options.logger.error(`Error checking if Security Labs docs are loaded: ${e.message}`); + this.options.logger.info(`Error checking if Security Labs docs are loaded: ${e.message}`); return false; } }; From f06ada10ec7876ac17c124979e0075e53f033238 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 29 Oct 2024 17:35:11 +0100 Subject: [PATCH 38/41] switch to use assistant elser id --- .../knowledge_base_settings_management/index.tsx | 3 ++- .../server/ai_assistant_data_clients/knowledge_base/index.ts | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index 5cb156941ae2e..bc2d60941679a 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -191,7 +191,8 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d indices.push(entry.index); } }); - return dataViews.getExistingIndices(indices); + + return indices.length ? dataViews.getExistingIndices(indices) : Promise.resolve([]); }, [entries.data]); const { getColumns } = useKnowledgeBaseTable(); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 7f29ce214cf69..d32fe5bccfbb8 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -733,7 +733,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } try { - const elserId = await this.options.getElserId(); + const elserId = this.isV2KnowledgeBaseEnabled + ? ASSISTANT_ELSER_INFERENCE_ID + : await this.options.getElserId(); const userFilter = getKBUserFilter(user); const results = await this.findDocuments({ // Note: This is a magic number to set some upward bound as to not blow the context with too From 9a84bf150bca5cece29d2164b263ecc3604c63e9 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Tue, 29 Oct 2024 23:25:36 +0100 Subject: [PATCH 39/41] revert FF --- .../impl/capabilities/index.ts | 2 +- .../kbn-elastic-assistant/impl/assistant/index.test.tsx | 6 ++---- .../knowledge_base_settings_management/translations.ts | 7 +++++++ .../use_knowledge_base_table.tsx | 4 +++- .../knowledge_base/field_maps_configuration.ts | 2 +- .../ai_assistant_data_clients/knowledge_base/index.ts | 8 ++++---- .../security_solution/common/experimental_features.ts | 2 +- 7 files changed, 19 insertions(+), 12 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 54c24f6ce7b8f..c1c101fd74cd8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -19,6 +19,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures; * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, assistantModelEvaluation: false, }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx index c2158508620f0..1ef2db7b26c03 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx @@ -142,8 +142,7 @@ describe('Assistant', () => { }); describe('persistent storage', () => { - // Settings Modal is not available when `v2KnowledgeBaseEnabled` is true - it.skip('should refetchCurrentUserConversations after settings save button click', async () => { + it('should refetchCurrentUserConversations after settings save button click', async () => { const chatSendSpy = jest.spyOn(all, 'useChatSend'); await renderAssistant(); @@ -189,8 +188,7 @@ describe('Assistant', () => { ); }); - // Settings Modal is not available when `v2KnowledgeBaseEnabled` is true - it.skip('should refetchCurrentUserConversations after settings save button click, but do not update convos when refetch returns bad results', async () => { + it('should refetchCurrentUserConversations after settings save button click, but do not update convos when refetch returns bad results', async () => { jest.mocked(useFetchCurrentUserConversations).mockReturnValue({ data: mockData, isLoading: false, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts index b311f373c214b..98af0eabea6b5 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts @@ -372,3 +372,10 @@ export const MISSING_INDEX_TOOLTIP_CONTENT = i18n.translate( 'The index assigned to this knowledge base entry is unavailable. Check the permissions on the configured index, or that the index has not been deleted. You can update the index to be used for this knowledge entry, or delete the entry entirely.', } ); + +export const SECURITY_LABS_NOT_FULLY_LOADED = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.securityLabsNotFullyLoadedTooltipContent', + { + defaultMessage: 'Security Labs content is not fully loaded. Click to reload.', + } +); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx index c8a80e3e05995..cbdf97f116f7b 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx @@ -199,7 +199,9 @@ export const useKnowledgeBaseTable = () => { `} /> ) : ( - + + + )} ) : entry.type === DocumentEntryType.value ? ( diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index 6bb0a700c1015..348efb5a18f7d 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -6,7 +6,7 @@ */ import { FieldMap } from '@kbn/data-stream-adapter'; -export const ASSISTANT_ELSER_INFERENCE_ID = 'assistant-internal-elser2'; +export const ASSISTANT_ELSER_INFERENCE_ID = 'elastic-security-ai-assistant-elser2'; export const knowledgeBaseFieldMap: FieldMap = { '@timestamp': { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index d32fe5bccfbb8..ba1aa7d17a21b 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -200,7 +200,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { ); } } catch (e) { - this.options.logger.error(`Error checking if ELSER model '${elserId}' is deployed: ${e}`); + this.options.logger.debug(`Error checking if ELSER model '${elserId}' is deployed: ${e}`); // Returns 404 if it doesn't exist return false; } @@ -215,7 +215,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { task_type: 'sparse_embedding', })); } catch (error) { - this.options.logger.error( + this.options.logger.debug( `Error checking if Inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} exists: ${error}` ); return false; @@ -504,9 +504,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { track_total_hits: true, }); - return !!result.hits?.total; + return !!(result.hits?.total as SearchTotalHits).value; } catch (e) { - this.options.logger.error(`Error checking if user's KB docs exist: ${e.message}`); + this.options.logger.debug(`Error checking if user's KB docs exist: ${e.message}`); return false; } }; diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 68ced40ca885f..792b6352912b3 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -111,7 +111,7 @@ export const allowedExperimentalValues = Object.freeze({ /** * Enables new Knowledge Base Entries features, introduced in `8.15.0`. */ - assistantKnowledgeBaseByDefault: true, + assistantKnowledgeBaseByDefault: false, /** * Enables the Managed User section inside the new user details flyout. From b11e196e5e19174b0a6034ec224420e7d7cc939e Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 30 Oct 2024 08:13:26 +0100 Subject: [PATCH 40/41] tests --- .../knowledge_base/index.ts | 4 ++-- .../knowledge_base/get_knowledge_base_status.ts | 12 +++++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index ba1aa7d17a21b..555fc2cabbcb5 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -191,11 +191,11 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { stats.deployment_stats?.allocation_status.state === 'fully_allocated'; const isReadyServerless = (stats: MlTrainedModelStats) => - (stats.deployment_stats?.nodes as unknown as MlTrainedModelDeploymentNodesStats[]).some( + (stats.deployment_stats?.nodes as unknown as MlTrainedModelDeploymentNodesStats[])?.some( (node) => node.routing_state.routing_state === 'started' ); - return getResponse.trained_model_stats.some( + return getResponse.trained_model_stats?.some( (stats) => isReadyESS(stats) || isReadyServerless(stats) ); } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 3c43b57118adc..f278cd469ac0e 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -73,13 +73,19 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter pipeline_exists: pipelineExists, }; - if (v2KnowledgeBaseEnabled && indexExists && isModelDeployed) { - const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); + if (indexExists && isModelDeployed) { + const securityLabsExists = v2KnowledgeBaseEnabled + ? await kbDataClient.isSecurityLabsDocsLoaded() + : true; + const userDataExists = v2KnowledgeBaseEnabled + ? await kbDataClient.isUserDataExists() + : true; + return response.ok({ body: { ...body, security_labs_exists: securityLabsExists, - user_data_exists: await kbDataClient.isUserDataExists(), + user_data_exists: userDataExists, }, }); } From 4412d4de624a84dc8594be349d13c848d186d6f5 Mon Sep 17 00:00:00 2001 From: Patryk Kopycinski Date: Wed, 30 Oct 2024 16:23:44 +0100 Subject: [PATCH 41/41] fix kbretrievaltool with flag off --- .../knowledge_base/helpers.ts | 35 ++++++++++++++----- .../knowledge_base/index.ts | 6 ++++ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index d6cb71c082717..a19b3f0945086 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -37,6 +37,7 @@ export const isModelAlreadyExistsError = (error: Error) => { export const getKBVectorSearchQuery = ({ filter, kbResource, + modelId, query, required, user, @@ -44,6 +45,7 @@ export const getKBVectorSearchQuery = ({ }: { filter?: QueryDslQueryContainer | undefined; kbResource?: string | undefined; + modelId: string; query?: string; required?: boolean | undefined; user: AuthenticatedUser; @@ -112,16 +114,33 @@ export const getKBVectorSearchQuery = ({ ], }; - const semanticTextFilter = query - ? [ - { - semantic: { - field: 'semantic_text', - query, + let semanticTextFilter: + | Array<{ semantic: { field: string; query: string } }> + | Array<{ + text_expansion: { 'vector.tokens': { model_id: string; model_text: string } }; + }> = []; + + if (v2KnowledgeBaseEnabled && query) { + semanticTextFilter = [ + { + semantic: { + field: 'semantic_text', + query, + }, + }, + ]; + } else if (!v2KnowledgeBaseEnabled) { + semanticTextFilter = [ + { + text_expansion: { + 'vector.tokens': { + model_id: modelId, + model_text: query as string, }, }, - ] - : []; + }, + ]; + } return { bool: { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 555fc2cabbcb5..f985095661f3e 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -488,12 +488,14 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } const esClient = await this.options.elasticsearchClientPromise; + const modelId = await this.options.getElserId(); try { const vectorSearchQuery = getKBVectorSearchQuery({ kbResource: USER_RESOURCE, required: false, user, + modelId, v2KnowledgeBaseEnabled: this.options.v2KnowledgeBaseEnabled, }); @@ -525,12 +527,14 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { const expectedDocsCount = await getSecurityLabsDocsCount({ logger: this.options.logger }); const esClient = await this.options.elasticsearchClientPromise; + const modelId = await this.options.getElserId(); try { const vectorSearchQuery = getKBVectorSearchQuery({ kbResource: SECURITY_LABS_RESOURCE, required: false, user, + modelId, v2KnowledgeBaseEnabled: this.options.v2KnowledgeBaseEnabled, }); @@ -577,6 +581,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } const esClient = await this.options.elasticsearchClientPromise; + const modelId = await this.options.getElserId(); const vectorSearchQuery = getKBVectorSearchQuery({ filter, @@ -584,6 +589,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { query, required, user, + modelId, v2KnowledgeBaseEnabled: this.options.v2KnowledgeBaseEnabled, });