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 bb1ad4a9d15e2..13db8a8e7d1f7 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 @@ -28,9 +28,11 @@ import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/imp import { KnowledgeBaseResource } from '@kbn/elastic-assistant-common'; import { createTrainedModelsProviderMock } from '@kbn/ml-plugin/server/shared_services/providers/__mocks__/trained_models'; import { ASSISTANT_ELSER_INFERENCE_ID } from './field_maps_configuration'; +import { ensureIntegrationKnowledgeIndexEntry } from '../../ai_assistant_service/integration_knowledge_helper'; jest.mock('@kbn/ml-plugin/server/lib/node_utils'); jest.mock('../../lib/langchain/content_loaders/security_labs_loader'); +jest.mock('../../ai_assistant_service/integration_knowledge_helper'); jest.mock('p-retry'); const date = '2023-03-28T22:27:28.159Z'; let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>; @@ -42,6 +44,11 @@ const mockedPRetry = pRetry as jest.MockedFunction; mockedPRetry.mockResolvedValue({}); const telemetry = coreMock.createSetup().analytics; +const mockEnsureIntegrationKnowledgeIndexEntry = + ensureIntegrationKnowledgeIndexEntry as jest.MockedFunction< + typeof ensureIntegrationKnowledgeIndexEntry + >; + describe('AIAssistantKnowledgeBaseDataClient', () => { let mockOptions: KnowledgeBaseDataClientParams; let ml: MlPluginSetup; @@ -54,6 +61,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { jest.clearAllMocks(); logger = loggingSystemMock.createLogger(); mockLoadSecurityLabs.mockClear(); + mockEnsureIntegrationKnowledgeIndexEntry.mockResolvedValue(true); ml = mlPluginMock.createSetupContract() as unknown as MlPluginSetup; // Missing SharedServices mock, so manually mocking trainedModelsProvider mockOptions = { logger, @@ -71,6 +79,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { setIsKBSetupInProgress: jest.fn().mockImplementation(() => {}), manageGlobalKnowledgeBaseAIAssistant: true, getTrainedModelsProvider: () => trainedModelsProviderMock, + telemetry, }; esClientMock.search.mockReturnValue( // @ts-expect-error not full response interface @@ -301,6 +310,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); await client.setupKnowledgeBase({}); + expect(mockEnsureIntegrationKnowledgeIndexEntry).toHaveBeenCalled(); expect(loadSecurityLabs).toHaveBeenCalled(); }); 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 dfc0b2a5e1025..6b825974eb0ac 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 @@ -74,6 +74,7 @@ import { ASSISTANT_ELSER_INFERENCE_ID } from './field_maps_configuration'; import type { BulkOperationError } from '../../lib/data_stream/documents_data_writer'; import { AUDIT_OUTCOME, KnowledgeBaseAuditAction, knowledgeBaseAuditEvent } from './audit_events'; import { findDocuments } from '../find'; +import { ensureIntegrationKnowledgeIndexEntry } from '../../ai_assistant_service/integration_knowledge_helper'; /** * Params for when creating KbDataClient in Request Context Factory. Useful if needing to modify @@ -94,6 +95,7 @@ export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientPara manageGlobalKnowledgeBaseAIAssistant: boolean; getTrainedModelsProvider: () => ReturnType; elserInferenceId?: string; + telemetry: AnalyticsServiceSetup; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { constructor(public readonly options: KnowledgeBaseDataClientParams) { @@ -208,6 +210,19 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return; } + // Ensure integration knowledge index entry exists first + try { + await ensureIntegrationKnowledgeIndexEntry( + this, + this.options.logger.get('integrationKnowledge'), + this.options.telemetry + ); + } catch (error) { + this.options.logger.error( + `Failed to ensure integration knowledge index entry: ${error.message}` + ); + } + try { this.options.logger.debug('Checking if ML nodes are available...'); const mlNodesCount = await getMlNodeCount({ 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 19a8cd075c6d4..8c75ef1bc764a 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 @@ -9,6 +9,7 @@ import { elasticsearchServiceMock, loggingSystemMock, savedObjectsClientMock, + analyticsServiceMock, } from '@kbn/core/server/mocks'; import type { IndicesDataStream, @@ -141,6 +142,7 @@ describe('AI Assistant Service', () => { update: jest.fn(), uninstall: jest.fn(), }), + telemetry: analyticsServiceMock.createAnalyticsServiceSetup(), }; }); 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 6f2e5826e5978..af5046314e15a 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,6 +14,7 @@ import type { ElasticsearchClient, KibanaRequest, SavedObjectsClientContract, + AnalyticsServiceSetup, } from '@kbn/core/server'; import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; @@ -87,6 +88,7 @@ export interface AIAssistantServiceOpts { taskManager: TaskManagerSetupContract; pluginStop$: Subject; productDocManager: Promise; + telemetry: AnalyticsServiceSetup; } export interface CreateAIAssistantClientParams { @@ -608,6 +610,7 @@ export class AIAssistantService { spaceId: opts.spaceId, manageGlobalKnowledgeBaseAIAssistant: opts.manageGlobalKnowledgeBaseAIAssistant ?? false, getTrainedModelsProvider: opts.getTrainedModelsProvider, + telemetry: this.options.telemetry, }); } diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/integration_knowledge_helper.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/integration_knowledge_helper.test.ts new file mode 100644 index 0000000000000..7b36621128dae --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/integration_knowledge_helper.test.ts @@ -0,0 +1,293 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { AnalyticsServiceSetup, Logger } from '@kbn/core/server'; +import { loggingSystemMock } from '@kbn/core/server/mocks'; +import { IndexEntryType } from '@kbn/elastic-assistant-common'; +import type { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; +import { + checkIntegrationKnowledgeIndexEntryExists, + ensureIntegrationKnowledgeIndexEntry, +} from './integration_knowledge_helper'; + +describe('Integration Knowledge Helper', () => { + let mockKbDataClient: jest.Mocked; + let mockLogger: jest.Mocked; + let mockTelemetry: jest.Mocked; + + beforeEach(() => { + mockKbDataClient = { + findDocuments: jest.fn(), + createKnowledgeBaseEntry: jest.fn(), + } as unknown as jest.Mocked; + + mockLogger = loggingSystemMock.createLogger(); + + mockTelemetry = { + reportEvent: jest.fn(), + } as unknown as jest.Mocked; + + jest.clearAllMocks(); + }); + + describe('checkIntegrationKnowledgeIndexEntryExists', () => { + it('should return true when integration knowledge index entry exists', async () => { + const mockResults = { + total: 1, + data: [ + { + id: 'test-id', + type: IndexEntryType.value, + index: '.integration_knowledge', + field: 'content', + name: 'Integration Knowledge', + }, + ], + }; + + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockResults); + + const result = await checkIntegrationKnowledgeIndexEntryExists({ + kbDataClient: mockKbDataClient, + logger: mockLogger, + }); + + expect(result).toBe(true); + expect(mockKbDataClient.findDocuments).toHaveBeenCalledWith({ + page: 1, + perPage: 1, + filter: 'type:index AND index:.integration_knowledge', + }); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Integration knowledge index entry exists: true' + ); + }); + + it('should return false when integration knowledge index entry does not exist', async () => { + const mockResults = { + total: 0, + data: [], + }; + + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockResults); + + const result = await checkIntegrationKnowledgeIndexEntryExists({ + kbDataClient: mockKbDataClient, + logger: mockLogger, + }); + + expect(result).toBe(false); + expect(mockKbDataClient.findDocuments).toHaveBeenCalledWith({ + page: 1, + perPage: 1, + filter: 'type:index AND index:.integration_knowledge', + }); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Integration knowledge index entry exists: false' + ); + }); + + it('should return false and log error when findDocuments throws an error', async () => { + const errorMessage = 'Database connection failed'; + (mockKbDataClient.findDocuments as jest.Mock).mockRejectedValue(new Error(errorMessage)); + + const result = await checkIntegrationKnowledgeIndexEntryExists({ + kbDataClient: mockKbDataClient, + logger: mockLogger, + }); + + expect(result).toBe(false); + expect(mockLogger.error).toHaveBeenCalledWith( + `Error checking integration knowledge index entry: ${errorMessage}` + ); + }); + }); + + describe('ensureIntegrationKnowledgeIndexEntry', () => { + it('should return true when integration knowledge index entry already exists', async () => { + const mockResults = { + total: 1, + data: [ + { + id: 'existing-id', + type: IndexEntryType.value, + index: '.integration_knowledge', + }, + ], + }; + + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockResults); + + const result = await ensureIntegrationKnowledgeIndexEntry( + mockKbDataClient, + mockLogger, + mockTelemetry + ); + + expect(result).toBe(true); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Checking if integration knowledge index entry exists...' + ); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Integration knowledge index entry already exists' + ); + expect(mockKbDataClient.createKnowledgeBaseEntry).not.toHaveBeenCalled(); + }); + + it('should create integration knowledge index entry when it does not exist', async () => { + const mockFindResults = { + total: 0, + data: [], + }; + + const mockCreatedEntry = { + id: 'new-entry-id', + type: IndexEntryType.value, + index: '.integration_knowledge', + field: 'content', + name: 'Integration Knowledge', + description: + 'Integration knowledge base containing semantic information about integrations installed via Fleet. Use this tool to search for information about integrations, integration configurations, troubleshooting guides, and best practices', + queryDescription: + 'Key terms to retrieve relevant integration details, like integration name, configuration values the user is having issues with, and/or any other general keywords', + global: true, + users: [], + }; + + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockFindResults); + (mockKbDataClient.createKnowledgeBaseEntry as jest.Mock).mockResolvedValue(mockCreatedEntry); + + const result = await ensureIntegrationKnowledgeIndexEntry( + mockKbDataClient, + mockLogger, + mockTelemetry + ); + + expect(result).toBe(true); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Checking if integration knowledge index entry exists...' + ); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Creating integration knowledge index entry...' + ); + expect(mockLogger.info).toHaveBeenCalledWith( + 'Integration knowledge index entry created successfully' + ); + + expect(mockKbDataClient.createKnowledgeBaseEntry).toHaveBeenCalledWith({ + knowledgeBaseEntry: { + type: IndexEntryType.value, + index: '.integration_knowledge', + field: 'content', + name: 'Integration Knowledge', + description: + 'Integration knowledge base containing semantic information about integrations installed via Fleet. Use this tool to search for information about integrations, integration configurations, troubleshooting guides, and best practices', + queryDescription: + 'Key terms to retrieve relevant integration details, like integration name, configuration values the user is having issues with, and/or any other general keywords', + global: true, + users: [], + }, + telemetry: mockTelemetry, + }); + }); + + it('should return false and log error when createKnowledgeBaseEntry returns null', async () => { + const mockFindResults = { + total: 0, + data: [], + }; + + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockFindResults); + (mockKbDataClient.createKnowledgeBaseEntry as jest.Mock).mockResolvedValue(null); + + const result = await ensureIntegrationKnowledgeIndexEntry( + mockKbDataClient, + mockLogger, + mockTelemetry + ); + + expect(result).toBe(false); + expect(mockLogger.error).toHaveBeenCalledWith( + 'Failed to create integration knowledge index entry' + ); + }); + + it('should return false and log error when createKnowledgeBaseEntry throws an error', async () => { + const mockFindResults = { + total: 0, + data: [], + }; + + const errorMessage = 'Failed to create entry'; + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockFindResults); + (mockKbDataClient.createKnowledgeBaseEntry as jest.Mock).mockRejectedValue( + new Error(errorMessage) + ); + + const result = await ensureIntegrationKnowledgeIndexEntry( + mockKbDataClient, + mockLogger, + mockTelemetry + ); + + expect(result).toBe(false); + expect(mockLogger.error).toHaveBeenCalledWith( + `Error ensuring integration knowledge index entry: ${errorMessage}` + ); + }); + + it('should return false and log error when checkIntegrationKnowledgeIndexEntryExists throws an error', async () => { + const errorMessage = 'Check failed'; + (mockKbDataClient.findDocuments as jest.Mock).mockRejectedValue(new Error(errorMessage)); + + const result = await ensureIntegrationKnowledgeIndexEntry( + mockKbDataClient, + mockLogger, + mockTelemetry + ); + + expect(result).toBe(false); + expect(mockLogger.error).toHaveBeenCalledWith( + `Error checking integration knowledge index entry: ${errorMessage}` + ); + }); + + it('should use correct knowledge base entry configuration', async () => { + const mockFindResults = { + total: 0, + data: [], + }; + + const mockCreatedEntry = { id: 'test-id' }; + (mockKbDataClient.findDocuments as jest.Mock).mockResolvedValue(mockFindResults); + (mockKbDataClient.createKnowledgeBaseEntry as jest.Mock).mockResolvedValue(mockCreatedEntry); + + await ensureIntegrationKnowledgeIndexEntry(mockKbDataClient, mockLogger, mockTelemetry); + + const createCall = (mockKbDataClient.createKnowledgeBaseEntry as jest.Mock).mock.calls[0][0]; + const knowledgeBaseEntry = createCall.knowledgeBaseEntry; + + expect(knowledgeBaseEntry.type).toBe(IndexEntryType.value); + expect(knowledgeBaseEntry.name).toBe('Integration Knowledge'); + expect(knowledgeBaseEntry.global).toBe(true); + expect(knowledgeBaseEntry.users).toEqual([]); + + // Type guard to check if it's an index entry + if (knowledgeBaseEntry.type === IndexEntryType.value) { + expect(knowledgeBaseEntry.index).toBe('.integration_knowledge'); + expect(knowledgeBaseEntry.field).toBe('content'); + expect(knowledgeBaseEntry.description).toContain( + 'Integration knowledge base containing semantic information' + ); + expect(knowledgeBaseEntry.queryDescription).toContain( + 'Key terms to retrieve relevant integration details' + ); + } + expect(createCall.telemetry).toBe(mockTelemetry); + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/integration_knowledge_helper.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/integration_knowledge_helper.ts new file mode 100644 index 0000000000000..d8280b829c9de --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/integration_knowledge_helper.ts @@ -0,0 +1,92 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { AnalyticsServiceSetup, Logger } from '@kbn/core/server'; +import { IndexEntryType } from '@kbn/elastic-assistant-common'; +import type { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; +import type { EsIndexEntry } from '../ai_assistant_data_clients/knowledge_base/types'; + +const INTEGRATION_KNOWLEDGE_INDEX_NAME = '.integration_knowledge'; + +/** + * Checks if the integration knowledge index entry already exists + */ +export const checkIntegrationKnowledgeIndexEntryExists = async ({ + kbDataClient, + logger, +}: { + kbDataClient: AIAssistantKnowledgeBaseDataClient; + logger: Logger; +}): Promise => { + try { + const results = await kbDataClient.findDocuments({ + page: 1, + perPage: 1, + filter: `type:index AND index:${INTEGRATION_KNOWLEDGE_INDEX_NAME}`, + }); + + const exists = results.total > 0; + logger.debug(`Integration knowledge index entry exists: ${exists}`); + return exists; + } catch (error) { + logger.error(`Error checking integration knowledge index entry: ${error.message}`); + return false; + } +}; + +/** + * Ensures the integration knowledge index entry exists during Knowledge Base setup. + * Similar to loadSecurityLabs() but for Index Entries rather than Document Entries. + */ +export const ensureIntegrationKnowledgeIndexEntry = async ( + kbDataClient: AIAssistantKnowledgeBaseDataClient, + logger: Logger, + telemetry: AnalyticsServiceSetup +): Promise => { + try { + logger.debug('Checking if integration knowledge index entry exists...'); + + const entryExists = await checkIntegrationKnowledgeIndexEntryExists({ + kbDataClient, + logger, + }); + + if (!entryExists) { + logger.debug('Creating integration knowledge index entry...'); + + const entry = await kbDataClient.createKnowledgeBaseEntry({ + knowledgeBaseEntry: { + type: IndexEntryType.value, + index: INTEGRATION_KNOWLEDGE_INDEX_NAME, + field: 'content', + name: 'Integration Knowledge', + description: + 'Integration knowledge base containing semantic information about integrations installed via Fleet. Use this tool to search for information about integrations, integration configurations, troubleshooting guides, and best practices', + queryDescription: + 'Key terms to retrieve relevant integration details, like integration name, configuration values the user is having issues with, and/or any other general keywords', + global: true, + users: [], + }, + telemetry, + }); + + if (entry) { + logger.info('Integration knowledge index entry created successfully'); + return true; + } else { + logger.error('Failed to create integration knowledge index entry'); + return false; + } + } else { + logger.debug('Integration knowledge index entry already exists'); + return true; + } + } catch (error) { + logger.error(`Error ensuring integration knowledge index entry: ${error.message}`); + return false; + } +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts index b10e4b1e105d2..ea06f857938d5 100755 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts @@ -106,6 +106,7 @@ export class ElasticAssistantPlugin .getStartServices() .then(([_, { productDocBase }]) => productDocBase.management), pluginStop$: this.pluginStop$, + telemetry: core.analytics, }); const requestContextFactory = new RequestContextFactory({