diff --git a/oas_docs/output/kibana.serverless.yaml b/oas_docs/output/kibana.serverless.yaml index e4c63ba1b5fd2..27a25d53ccc3a 100644 --- a/oas_docs/output/kibana.serverless.yaml +++ b/oas_docs/output/kibana.serverless.yaml @@ -36273,14 +36273,10 @@ paths: properties: elser_exists: type: boolean - index_exists: - type: boolean is_setup_available: type: boolean is_setup_in_progress: type: boolean - pipeline_exists: - type: boolean product_documentation_status: type: string security_labs_exists: diff --git a/oas_docs/output/kibana.yaml b/oas_docs/output/kibana.yaml index b4fbd8d90a6ce..de5eb0693c55f 100644 --- a/oas_docs/output/kibana.yaml +++ b/oas_docs/output/kibana.yaml @@ -20366,14 +20366,10 @@ paths: properties: elser_exists: type: boolean - index_exists: - type: boolean is_setup_available: type: boolean is_setup_in_progress: type: boolean - pipeline_exists: - type: boolean product_documentation_status: type: string security_labs_exists: diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml index f8ce62bd45111..dba5691ce5879 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml @@ -437,14 +437,10 @@ paths: properties: elser_exists: type: boolean - index_exists: - type: boolean is_setup_available: type: boolean is_setup_in_progress: type: boolean - pipeline_exists: - type: boolean product_documentation_status: type: string security_labs_exists: diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml index a66191ea04538..59a1906b227dc 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml @@ -437,14 +437,10 @@ paths: properties: elser_exists: type: boolean - index_exists: - type: boolean is_setup_available: type: boolean is_setup_in_progress: type: boolean - pipeline_exists: - type: boolean product_documentation_status: type: string security_labs_exists: diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index b031bc8d37185..f22089412d3a7 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -67,10 +67,8 @@ export type ReadKnowledgeBaseRequestParamsInput = z.input; export const ReadKnowledgeBaseResponse = z.object({ elser_exists: z.boolean().optional(), - index_exists: z.boolean().optional(), is_setup_available: z.boolean().optional(), is_setup_in_progress: z.boolean().optional(), - pipeline_exists: z.boolean().optional(), security_labs_exists: z.boolean().optional(), user_data_exists: z.boolean().optional(), product_documentation_status: z.string().optional(), diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml index 8b9e3bf741652..9320d175e9205 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml @@ -75,14 +75,10 @@ paths: properties: elser_exists: type: boolean - index_exists: - type: boolean is_setup_available: type: boolean is_setup_in_progress: type: boolean - pipeline_exists: - type: boolean security_labs_exists: type: boolean user_data_exists: diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx index a7d73c8b10a62..0a7469c2e1fde 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx @@ -33,8 +33,6 @@ jest.mock('@tanstack/react-query', () => ({ const statusResponse = { elser_exists: true, - index_exists: true, - pipeline_exists: true, security_labs_exists: true, }; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index 5745e17a3a151..51b0298c710d0 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx @@ -100,8 +100,6 @@ export const useInvalidateKnowledgeBaseStatus = () => { */ 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 || diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index c240d5ac6b60b..359f82dcf3185 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -62,11 +62,7 @@ export const useChatSend = ({ const { isLoading, sendMessage, abortStream } = useSendMessage(); const { clearConversation, removeLastMessage } = useConversation(); const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled }); - const isSetupComplete = - kbStatus?.elser_exists && - kbStatus?.index_exists && - kbStatus?.pipeline_exists && - kbStatus?.security_labs_exists; + const isSetupComplete = kbStatus?.elser_exists && kbStatus?.security_labs_exists; // Handles sending latest user prompt to API const handleSendMessage = useCallback( diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx index 14267ffb60868..9230bab09dc7f 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx @@ -70,8 +70,6 @@ jest.mock('../assistant/api/knowledge_base/use_knowledge_base_status', () => ({ return { data: { elser_exists: true, - index_exists: true, - pipeline_exists: true, }, isLoading: false, isFetching: false, @@ -88,8 +86,6 @@ describe('Knowledge base settings', () => { return { data: { elser_exists: true, - index_exists: false, - pipeline_exists: false, is_setup_available: true, }, isLoading: false, @@ -111,8 +107,6 @@ describe('Knowledge base settings', () => { return { data: { elser_exists: false, - index_exists: false, - pipeline_exists: false, }, isLoading: false, isFetching: false, diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx index 49830d337e7cb..100f885c13f09 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx @@ -60,11 +60,7 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const isElserEnabled = kbStatus?.elser_exists ?? false; const isSecurityLabsEnabled = kbStatus?.security_labs_exists ?? false; const isKnowledgeBaseSetup = - (isElserEnabled && - kbStatus?.index_exists && - kbStatus?.pipeline_exists && - (isSecurityLabsEnabled || kbStatus?.user_data_exists)) ?? - false; + (isElserEnabled && (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/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx index a48010f088c42..451f52b7f4d81 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx @@ -49,6 +49,7 @@ export const DocumentEntryEditor: React.FC = React.memo( setEntry((prevEntry) => ({ ...prevEntry, users: value === i18n.SHARING_GLOBAL_OPTION_LABEL ? [] : privateUsers, + global: value === i18n.SHARING_GLOBAL_OPTION_LABEL ? true : false, })), [privateUsers, setEntry] ); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx index 222df5c87c855..c4edaae964ad6 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.test.tsx @@ -27,6 +27,7 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import { useKnowledgeBaseIndices } from '../../assistant/api/knowledge_base/use_knowledge_base_indices'; import { Router } from '@kbn/shared-ux-router'; import { createMemoryHistory, History } from 'history'; +import { userProfileServiceMock } from '@kbn/core-user-profile-browser-mocks'; const mockContext = { basePromptContexts: MOCK_QUICK_PROMPTS, @@ -39,6 +40,7 @@ const mockContext = { isAssistantEnabled: true, hasManageGlobalKnowledgeBase: true, }, + userProfileService: userProfileServiceMock.createStart(), }; jest.mock('../../assistant_context'); jest.mock('../../assistant/api/knowledge_base/entries/use_create_knowledge_base_entry'); @@ -157,8 +159,6 @@ describe('KnowledgeBaseSettingsManagement', () => { data: { elser_exists: true, security_labs_exists: true, - index_exists: true, - pipeline_exists: true, }, isFetched: true, }); @@ -202,8 +202,6 @@ describe('KnowledgeBaseSettingsManagement', () => { data: { elser_exists: false, security_labs_exists: false, - index_exists: false, - pipeline_exists: false, }, isFetched: true, }); @@ -458,9 +456,11 @@ describe('KnowledgeBaseSettingsManagement', () => { }); const updatedName = 'New Entry Name'; - await waitFor(() => { + + await waitFor(async () => { const nameInput = screen.getByTestId('entryNameInput'); - fireEvent.change(nameInput, { target: { value: updatedName } }); + await userEvent.clear(nameInput); + await userEvent.type(nameInput, updatedName); }); await waitFor(() => { @@ -476,7 +476,7 @@ describe('KnowledgeBaseSettingsManagement', () => { }); expect(mockCreateEntry).toHaveBeenCalledTimes(0); expect(mockUpdateEntry).toHaveBeenCalledWith([{ ...mockData[0], name: updatedName }]); - }, 100000000); + }); it('does not create a duplicate index entry when switching sharing option twice', async () => { (useFlyoutModalVisibility as jest.Mock).mockReturnValue({ @@ -553,7 +553,11 @@ describe('KnowledgeBaseSettingsManagement', () => { expect(mockCreateEntry).toHaveBeenCalledTimes(1); }); expect(mockUpdateEntry).toHaveBeenCalledTimes(0); - expect(mockCreateEntry).toHaveBeenCalledWith({ ...mockData[3], users: undefined }); + expect(mockCreateEntry).toHaveBeenCalledWith({ + ...mockData[3], + global: false, + users: undefined, + }); }); it('does not show duplicate entry modal on new document entry creation', async () => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx index 7f55cfc5538c7..1f1a1b54fd729 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx @@ -57,6 +57,7 @@ export const IndexEntryEditor: React.FC = React.memo( setEntry((prevEntry) => ({ ...prevEntry, users: value === i18n.SHARING_GLOBAL_OPTION_LABEL ? [] : privateUsers, + global: value === i18n.SHARING_GLOBAL_OPTION_LABEL ? true : false, })), [privateUsers, setEntry] ); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx index f11cb4473c74b..6957a8a5db028 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx @@ -22,9 +22,8 @@ import { IndexEntryType, KnowledgeBaseEntryResponse, } from '@kbn/elastic-assistant-common'; - -import useAsync from 'react-use/lib/useAsync'; import { UserProfileAvatarData } from '@kbn/user-profile-components'; +import { useQuery } from '@tanstack/react-query'; import { useAssistantContext } from '../../..'; import * as i18n from './translations'; import { BadgesColumn } from '../../assistant/common/components/assistant_settings_management/badges'; @@ -32,26 +31,37 @@ import { useInlineActions } from '../../assistant/common/components/assistant_se import { isSystemEntry } from './helpers'; import { SetupKnowledgeBaseButton } from '../setup_knowledge_base_button'; -const AuthorColumn = ({ entry }: { entry: KnowledgeBaseEntryResponse }) => { +const useUserProfile = ({ username, enabled = true }: { username: string; enabled: boolean }) => { const { userProfileService } = useAssistantContext(); - const userProfile = useAsync(async () => { - if (isSystemEntry(entry) || entry.createdBy === 'unknown') { - return; - } + return useQuery({ + queryKey: ['userProfile', username], + queryFn: async () => { + const data = await userProfileService?.bulkGet<{ avatar: UserProfileAvatarData }>({ + uids: new Set([username]), + dataPath: 'avatar', + }); - const profile = await userProfileService?.bulkGet<{ avatar: UserProfileAvatarData }>({ - uids: new Set([entry.createdBy]), - dataPath: 'avatar', - }); - return { username: profile?.[0].user.username, avatar: profile?.[0].data.avatar }; - }, [entry.createdBy]); + return data ?? null; + }, + select: (profile) => { + return { + username: profile?.[0]?.user.username ?? 'Unknown', + avatar: profile?.[0]?.data.avatar, + }; + }, + enabled: !!(enabled && username?.length), + }); +}; - const userName = useMemo( - () => userProfile?.value?.username ?? 'Unknown', - [userProfile?.value?.username] - ); - const userAvatar = userProfile?.value?.avatar; +const AuthorColumn = ({ entry }: { entry: KnowledgeBaseEntryResponse }) => { + const { data: userProfile } = useUserProfile({ + username: entry.createdBy, + enabled: !(isSystemEntry(entry) || entry.createdBy === 'unknown'), + }); + + const userName = useMemo(() => userProfile?.username ?? 'Unknown', [userProfile?.username]); + const userAvatar = userProfile?.avatar; const badgeItem = isSystemEntry(entry) ? 'Elastic' : userName; const userImage = isSystemEntry(entry) ? ( = React.memo(({ display } const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts }); const isSetupInProgress = kbStatus?.is_setup_in_progress || isSettingUpKB; - const isSetupComplete = - kbStatus?.elser_exists && - kbStatus?.index_exists && - kbStatus?.pipeline_exists && - kbStatus?.security_labs_exists; + const isSetupComplete = kbStatus?.elser_exists && kbStatus?.security_labs_exists; const onInstallKnowledgeBase = useCallback(() => { setupKB(); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json b/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json index 42fc0eb1a5b53..d6f3eaef65560 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/tsconfig.json @@ -41,5 +41,6 @@ "@kbn/spaces-plugin", "@kbn/shared-ux-router", "@kbn/inference-endpoint-ui-common", + "@kbn/core-user-profile-browser-mocks", ] } diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index 16fe8ac9a22a1..f8fd88b113bf4 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -6,7 +6,6 @@ */ import { FieldMap } from '@kbn/data-stream-adapter'; -export const ELSER_MODEL_2 = ['.elser_model_2', '.elser_model_2_linux-x86_64']; export const ASSISTANT_ELSER_INFERENCE_ID = 'elastic-security-ai-assistant-elser2'; export const ELASTICSEARCH_ELSER_INFERENCE_ID = '.elser-2-elasticsearch'; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts index f34b523efcad4..aa8fcb0c1922e 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts @@ -63,12 +63,12 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { kibanaVersion: '8.8.0', ml, getElserId: getElserId.mockResolvedValue('elser-id'), + modelIdOverride: false, getIsKBSetupInProgress: mockGetIsKBSetupInProgress.mockReturnValue(false), getProductDocumentationStatus: jest.fn().mockResolvedValue('installed'), ingestPipelineResourceName: 'something', setIsKBSetupInProgress: jest.fn().mockImplementation(() => {}), manageGlobalKnowledgeBaseAIAssistant: true, - assistantDefaultInferenceEndpoint: false, trainedModelsProvider: trainedModelsProviderMock, }; esClientMock.search.mockReturnValue( @@ -332,7 +332,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { { fully_defined: false, model_id: '', tags: [], input: { field_names: ['content'] } }, ], }); - mockLoadSecurityLabs.mockRejectedValue(new Error('Installation error')); + (getMlNodeCount as jest.Mock).mockRejectedValue(new Error('Installation error')); const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); await expect(client.setupKnowledgeBase({})).rejects.toThrow( diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index afbc972add33a..44a7daa301d45 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -69,7 +69,6 @@ import { import { ASSISTANT_ELSER_INFERENCE_ID, ELASTICSEARCH_ELSER_INFERENCE_ID, - ELSER_MODEL_2, } from './field_maps_configuration'; import { BulkOperationError } from '../../lib/data_stream/documents_data_writer'; import { AUDIT_OUTCOME, KnowledgeBaseAuditAction, knowledgeBaseAuditEvent } from './audit_events'; @@ -91,8 +90,8 @@ export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientPara ingestPipelineResourceName: string; setIsKBSetupInProgress: (spaceId: string, isInProgress: boolean) => void; manageGlobalKnowledgeBaseAIAssistant: boolean; - assistantDefaultInferenceEndpoint: boolean; trainedModelsProvider: ReturnType; + modelIdOverride: boolean; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { constructor(public readonly options: KnowledgeBaseDataClientParams) { @@ -164,8 +163,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }; public getInferenceEndpointId = async () => { - const elserId = await this.options.getElserId(); - if (!this.options.assistantDefaultInferenceEndpoint || !ELSER_MODEL_2.includes(elserId)) { + // Don't use default enpdpoint for pt_tiny_elser + if (this.options.modelIdOverride) { return ASSISTANT_ELSER_INFERENCE_ID; } const esClient = await this.options.elasticsearchClientPromise; @@ -180,9 +179,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return ASSISTANT_ELSER_INFERENCE_ID; } } catch (error) { - this.options.logger.debug( - `Error checking if Inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} exists: ${error}` - ); + /* empty */ } // Fallback to the dedicated inference endpoint @@ -234,7 +231,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { ?.some((stats) => isReadyESS(stats) || isReadyServerless(stats)); } catch (error) { this.options.logger.debug( - `Error checking if Inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} exists: ${error}` + `Error checking if Inference endpoint ${inferenceId} exists: ${error}` ); return false; } @@ -274,7 +271,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { // it's being used in the mapping so we need to force delete force: true, }); - this.options.logger.debug(`Deleted existing inference endpoint for ELSER model '${elserId}'`); + this.options.logger.debug( + `Deleted existing inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} for ELSER model '${elserId}'` + ); } catch (error) { this.options.logger.error( `Error deleting inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} for ELSER model '${elserId}':\n${error}` @@ -362,37 +361,39 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return; } - this.options.logger.debug('Checking if ML nodes are available...'); - const mlNodesCount = await getMlNodeCount({ asInternalUser: esClient } as IScopedClusterClient); + try { + this.options.logger.debug('Checking if ML nodes are available...'); + const mlNodesCount = await getMlNodeCount({ + asInternalUser: esClient, + } as IScopedClusterClient); - if (mlNodesCount.count === 0 && mlNodesCount.lazyNodeCount === 0) { - throw new Error('No ML nodes available'); - } + if (mlNodesCount.count === 0 && mlNodesCount.lazyNodeCount === 0) { + throw new Error('No ML nodes available'); + } - this.options.logger.debug('Starting Knowledge Base setup...'); - this.options.setIsKBSetupInProgress(this.spaceId, true); - const elserId = await this.options.getElserId(); + this.options.logger.debug('Starting Knowledge Base setup...'); + this.options.setIsKBSetupInProgress(this.spaceId, true); + const elserId = await this.options.getElserId(); - // Delete legacy ESQL knowledge base docs if they exist, and silence the error if they do not - try { - const legacyESQL = await esClient.deleteByQuery({ - index: this.indexTemplateAndPattern.alias, - query: { - bool: { - must: [{ terms: { 'metadata.kbResource': ['esql', 'unknown'] } }], + // Delete legacy ESQL knowledge base docs if they exist, and silence the error if they do not + try { + const legacyESQL = await esClient.deleteByQuery({ + index: this.indexTemplateAndPattern.alias, + query: { + bool: { + must: [{ terms: { 'metadata.kbResource': ['esql', 'unknown'] } }], + }, }, - }, - }); - if (legacyESQL?.total != null && legacyESQL?.total > 0) { - this.options.logger.info( - `Removed ${legacyESQL?.total} ESQL knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.` - ); + }); + if (legacyESQL?.total != null && legacyESQL?.total > 0) { + this.options.logger.info( + `Removed ${legacyESQL?.total} ESQL knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.` + ); + } + } catch (e) { + this.options.logger.info('No legacy ESQL or Security Labs knowledge base docs to delete'); } - } catch (e) { - this.options.logger.info('No legacy ESQL or Security Labs knowledge base docs to delete'); - } - try { /* #1 Check if ELSER model is downloaded #2 Check if inference endpoint is deployed @@ -408,7 +409,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { (await this.isModelInstalled()) ? Promise.resolve() : Promise.reject(new Error('Model not installed')), - { minTimeout: 10000, maxTimeout: 10000, retries: 10 } + { minTimeout: 30000, maxTimeout: 30000, retries: 10 } ); this.options.logger.debug(`ELSER model '${elserId}' successfully installed!`); } else { @@ -452,17 +453,35 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } this.options.logger.debug(`Loading Security Labs KB docs...`); - await loadSecurityLabs(this, this.options.logger); + void loadSecurityLabs(this, this.options.logger); } else { this.options.logger.debug(`Security Labs Knowledge Base docs already loaded!`); } } + + const inferenceId = await this.getInferenceEndpointId(); + + if ( + inferenceId !== ASSISTANT_ELSER_INFERENCE_ID && + (await this.isInferenceEndpointExists(ASSISTANT_ELSER_INFERENCE_ID)) + ) { + try { + await this.deleteInferenceEndpoint(); + } catch (error) { + this.options.logger.debug( + `Error deleting inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID}` + ); + } + } + + // If loading security labs, we need to wait for the docs to be loaded + if (ignoreSecurityLabs) { + this.options.setIsKBSetupInProgress(this.spaceId, false); + } } catch (e) { this.options.setIsKBSetupInProgress(this.spaceId, 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(this.spaceId, false); } }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts index b94a46e81a937..ce951edbafdf9 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts @@ -112,6 +112,7 @@ describe('AI Assistant Service', () => { ); clusterClient.indices.getAlias.mockImplementation(async () => GetAliasResponse); clusterClient.indices.getDataStream.mockImplementation(async () => GetDataStreamResponse); + clusterClient.indices.simulateTemplate.mockImplementation(async () => SimulateTemplateResponse); ml = mlPluginMock.createSetupContract() as unknown as MlPluginSetup; // Missing SharedServices mock, so manually mocking trainedModelsProvider ml.trainedModelsProvider = jest.fn().mockImplementation(() => ({ getELSER: jest.fn().mockImplementation(() => '.elser_model_2'), diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 4eb101a3e857a..2e232b36210b0 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -14,10 +14,10 @@ import { Subject } from 'rxjs'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; import { ProductDocBaseStartContract } from '@kbn/product-doc-base-plugin/server'; import { - IndicesGetFieldMappingResponse, IndicesIndexSettings, + IndicesSimulateTemplateResponse, } from '@elastic/elasticsearch/lib/api/types'; -import { omit } from 'lodash'; +import { omit, some } from 'lodash'; import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status'; import { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers'; import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration'; @@ -95,6 +95,7 @@ export class AIAssistantService { private initialized: boolean; private isInitializing: boolean = false; private getElserId: GetElser; + private modelIdOverride: boolean = false; private conversationsDataStream: DataStreamSpacesAdapter; private knowledgeBaseDataStream: DataStreamSpacesAdapter; private promptsDataStream: DataStreamSpacesAdapter; @@ -107,8 +108,6 @@ export class AIAssistantService { private hasInitializedV2KnowledgeBase: boolean = false; private productDocManager?: ProductDocBaseStartContract['management']; private isProductDocumentationInProgress: boolean = false; - // Temporary 'feature flag' to determine if we should initialize the new knowledge base mappings - private assistantDefaultInferenceEndpoint: boolean = false; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -222,6 +221,76 @@ export class AIAssistantService { return newDataStream; }; + private async rolloverDataStream( + initialInferenceEndpointId: string, + targetInferenceEndpointId: string + ): Promise { + const esClient = await this.options.elasticsearchClientPromise; + + const currentDataStream = this.createDataStream({ + resource: 'knowledgeBase', + kibanaVersion: this.options.kibanaVersion, + fieldMap: { + ...omit(knowledgeBaseFieldMap, 'semantic_text'), + semantic_text: { + type: 'semantic_text', + array: false, + required: false, + inference_id: initialInferenceEndpointId, + search_inference_id: targetInferenceEndpointId, + }, + }, + }); + + // Add `search_inference_id` to the existing mappings + await currentDataStream.install({ + esClient, + logger: this.options.logger, + pluginStop$: this.options.pluginStop$, + }); + + // Migrate data stream mapping to the default inference_id + const newDS = this.createDataStream({ + resource: 'knowledgeBase', + kibanaVersion: this.options.kibanaVersion, + fieldMap: { + ...omit(knowledgeBaseFieldMap, ['semantic_text', 'vector', 'vector.tokens']), + semantic_text: { + type: 'semantic_text', + array: false, + required: false, + ...(targetInferenceEndpointId !== ELASTICSEARCH_ELSER_INFERENCE_ID + ? { inference_id: targetInferenceEndpointId } + : {}), + }, + }, + settings: { + // force new semantic_text field behavior + 'index.mapping.semantic_text.use_legacy_format': false, + }, + writeIndexOnly: true, + }); + + // We need to first install the templates and then rollover the indices + await newDS.installTemplates({ + esClient, + logger: this.options.logger, + pluginStop$: this.options.pluginStop$, + }); + + const indexNames = ( + await esClient.indices.getDataStream({ name: newDS.name }) + ).data_streams.map((ds) => ds.name); + + try { + await Promise.all( + indexNames.map((indexName) => esClient.indices.rollover({ alias: indexName })) + ); + } catch (e) { + /* empty */ + } + } + private async initializeResources(): Promise { this.isInitializing = true; try { @@ -243,119 +312,79 @@ export class AIAssistantService { pluginStop$: this.options.pluginStop$, }); - if (this.assistantDefaultInferenceEndpoint) { - const knowledgeBaseDataStreamExists = ( - await esClient.indices.getDataStream({ - name: this.knowledgeBaseDataStream.name, - }) - )?.data_streams?.length; - - // update component template for semantic_text field - // rollover - let mappings: IndicesGetFieldMappingResponse = {}; - try { - mappings = await esClient.indices.getFieldMapping({ - index: '.kibana-elastic-ai-assistant-knowledge-base-default', - fields: ['semantic_text'], - }); - } catch (error) { - /* empty */ - } - - const isUsingDedicatedInferenceEndpoint = - ( - Object.values(mappings)[0]?.mappings?.semantic_text?.mapping?.semantic_text as { - inference_id: string; - } - )?.inference_id === ASSISTANT_ELSER_INFERENCE_ID; - - if (knowledgeBaseDataStreamExists && isUsingDedicatedInferenceEndpoint) { - const currentDataStream = this.createDataStream({ - resource: 'knowledgeBase', - kibanaVersion: this.options.kibanaVersion, - fieldMap: { - ...omit(knowledgeBaseFieldMap, 'semantic_text'), - semantic_text: { - type: 'semantic_text', - array: false, - required: false, - inference_id: ASSISTANT_ELSER_INFERENCE_ID, - search_inference_id: ELASTICSEARCH_ELSER_INFERENCE_ID, - }, - }, - }); - - // Add `search_inference_id` to the existing mappings - await currentDataStream.install({ - esClient, - logger: this.options.logger, - pluginStop$: this.options.pluginStop$, - }); - - // Migrate data stream mapping to the default inference_id - const newDS = this.createDataStream({ - resource: 'knowledgeBase', - kibanaVersion: this.options.kibanaVersion, - fieldMap: { - ...omit(knowledgeBaseFieldMap, 'semantic_text'), - semantic_text: { - type: 'semantic_text', - array: false, - required: false, - }, - }, - settings: { - // force new semantic_text field behavior - 'index.mapping.semantic_text.use_legacy_format': false, - }, - writeIndexOnly: true, - }); + const knowledgeBaseDataSteams = ( + await esClient.indices.getDataStream({ + name: this.knowledgeBaseDataStream.name, + }) + )?.data_streams; + + let mappings: IndicesSimulateTemplateResponse[] = []; + try { + mappings = await Promise.all( + knowledgeBaseDataSteams.map((ds) => + esClient.indices.simulateTemplate({ + name: ds.template, + }) + ) + ); + } catch (error) { + /* empty */ + } - // We need to first install the templates and then rollover the indices - await newDS.installTemplates({ - esClient, - logger: this.options.logger, - pluginStop$: this.options.pluginStop$, - }); + const isUsingDedicatedInferenceEndpoint = some( + mappings, + (value) => + (value?.template?.mappings?.properties?.semantic_text as { inference_id: string }) + ?.inference_id === ASSISTANT_ELSER_INFERENCE_ID + ); - const indexNames = ( - await esClient.indices.getDataStream({ name: newDS.name }) - ).data_streams.map((ds) => ds.name); + // Used only for testing purposes + if (this.modelIdOverride && !isUsingDedicatedInferenceEndpoint) { + await this.rolloverDataStream( + ELASTICSEARCH_ELSER_INFERENCE_ID, + ASSISTANT_ELSER_INFERENCE_ID + ); + } else if (isUsingDedicatedInferenceEndpoint) { + await this.rolloverDataStream( + ASSISTANT_ELSER_INFERENCE_ID, + ELASTICSEARCH_ELSER_INFERENCE_ID + ); - try { - await Promise.all( - indexNames.map((indexName) => esClient.indices.rollover({ alias: indexName })) - ); - } catch (e) { - /* empty */ - } - } else { - // We need to make sure that the data stream is created with the correct mappings - this.knowledgeBaseDataStream = this.createDataStream({ - resource: 'knowledgeBase', - kibanaVersion: this.options.kibanaVersion, - fieldMap: { - ...omit(knowledgeBaseFieldMap, 'semantic_text'), - semantic_text: { - type: 'semantic_text', - array: false, - required: false, - }, - }, - settings: { - // force new semantic_text field behavior - 'index.mapping.semantic_text.use_legacy_format': false, - }, - writeIndexOnly: true, - }); - await this.knowledgeBaseDataStream.install({ - esClient, - logger: this.options.logger, - pluginStop$: this.options.pluginStop$, + // Delete the old inference endpoint + const elserId = await this.getElserId(); + try { + await esClient.inference.delete({ + inference_id: ASSISTANT_ELSER_INFERENCE_ID, + // it's being used in the mapping so we need to force delete + force: true, }); + this.options.logger.debug( + `Deleted existing inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} for ELSER model '${elserId}'` + ); + } catch (error) { + this.options.logger.error( + `Error deleting inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} for ELSER model '${elserId}':\n${error}` + ); } } else { - // Legacy path + // We need to make sure that the data stream is created with the correct mappings + this.knowledgeBaseDataStream = this.createDataStream({ + resource: 'knowledgeBase', + kibanaVersion: this.options.kibanaVersion, + fieldMap: { + ...omit(knowledgeBaseFieldMap, ['semantic_text', 'vector', 'vector.tokens']), + semantic_text: { + type: 'semantic_text', + array: false, + required: false, + }, + }, + settings: { + // force new semantic_text field behavior + 'index.mapping.semantic_text.use_legacy_format': false, + }, + writeIndexOnly: true, + }); await this.knowledgeBaseDataStream.install({ esClient, logger: this.options.logger, @@ -522,6 +551,7 @@ export class AIAssistantService { if (opts?.modelIdOverride != null) { const modelIdOverride = opts.modelIdOverride; this.getElserId = async () => modelIdOverride; + this.modelIdOverride = true; } // If a V2 KnowledgeBase has never been initialized or a modelIdOverride is provided, we need to reinitialize all persistence resources to make sure @@ -550,10 +580,10 @@ export class AIAssistantService { getProductDocumentationStatus: this.getProductDocumentationStatus.bind(this), kibanaVersion: this.options.kibanaVersion, ml: this.options.ml, + modelIdOverride: !!opts.modelIdOverride, setIsKBSetupInProgress: this.setIsKBSetupInProgress.bind(this), spaceId: opts.spaceId, manageGlobalKnowledgeBaseAIAssistant: opts.manageGlobalKnowledgeBaseAIAssistant ?? false, - assistantDefaultInferenceEndpoint: this.assistantDefaultInferenceEndpoint, trainedModelsProvider: opts.trainedModelsProvider, }); } diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts index f37e20df2bd98..856be9e61d17c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/content_loaders/security_labs_loader.ts @@ -62,6 +62,12 @@ export const loadSecurityLabs = async ( logger.info(`Loaded ${response?.length ?? 0} Security Labs docs into the Knowledge Base`); + try { + kbDataClient.options.setIsKBSetupInProgress(kbDataClient.spaceId, false); + } catch (e) { + /* empty */ + } + return response.length > 0; } catch (e) { logger.error(`Failed to load Security Labs docs into the Knowledge Base\n${e}`); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts index c846a1bfbbf60..ba63d06fc6105 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts @@ -55,10 +55,8 @@ describe('Get Knowledge Base Status Route', () => { expect(response.status).toEqual(200); expect(response.body).toEqual({ elser_exists: true, - index_exists: true, is_setup_in_progress: false, is_setup_available: true, - pipeline_exists: true, security_labs_exists: true, user_data_exists: true, product_documentation_status: 'installed', diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 6258db2b71d92..a078f8f530521 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -54,8 +54,6 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter return response.custom({ body: { success: false }, statusCode: 500 }); } - const indexExists = true; // Installed at startup, always true - const pipelineExists = true; // Installed at startup, always true const setupAvailable = await kbDataClient.isSetupAvailable(); const isInferenceEndpointExists = await kbDataClient.isInferenceEndpointExists(); const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded(); @@ -66,13 +64,11 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter return response.ok({ body: { elser_exists: isInferenceEndpointExists, - index_exists: indexExists, is_setup_in_progress: kbDataClient.isSetupInProgress, is_setup_available: setupAvailable, security_labs_exists: securityLabsExists, // If user data exists, we should have at least one document in the Security Labs index user_data_exists: userDataExists || !!loadedSecurityLabsDocsCount, - pipeline_exists: pipelineExists, product_documentation_status: productDocumentationStatus, }, }); diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/mocks/entries.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/mocks/entries.ts index c1d8e5594e853..0db7ba0a05255 100644 --- a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/mocks/entries.ts +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/mocks/entries.ts @@ -21,7 +21,6 @@ export const documentEntry: DocumentEntryCreateFields = { namespace: 'default', text: 'This is a sample document entry', global: false, - users: undefined, }; export const globalDocumentEntry: DocumentEntryCreateFields = { @@ -39,6 +38,5 @@ export const indexEntry: IndexEntryCreateFields = { field: 'sample-field', description: 'This is a sample index entry', queryDescription: 'Use sample-field to search in sample-index', - users: undefined, global: false, }; diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/remove_server_generated_properties.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/remove_server_generated_properties.ts index 1dd9e66b98198..24df23807ac0c 100644 --- a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/remove_server_generated_properties.ts +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/remove_server_generated_properties.ts @@ -4,7 +4,7 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import { omit, pickBy } from 'lodash'; +import { map, omit, pickBy } from 'lodash'; import { KnowledgeBaseEntryCreateProps } from '@kbn/elastic-assistant-common'; const serverGeneratedProperties = [ @@ -29,8 +29,10 @@ export type EntryWithoutServerGeneratedProperties = Omit< export const removeServerGeneratedProperties = ( entry: KnowledgeBaseEntryCreateProps ): EntryWithoutServerGeneratedProperties => { - const removedProperties = omit(entry, serverGeneratedProperties); - + const removedProperties = { + ...omit(entry, serverGeneratedProperties), + users: map(entry.users, (user) => omit(user, 'id')), + }; // We're only removing undefined values, so this cast correctly narrows the type return pickBy(removedProperties, (value) => value !== undefined) as KnowledgeBaseEntryCreateProps; };