diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/tool_prompts.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/tool_prompts.ts index 9c5c7a9691b85..e38300bb55d0e 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/tool_prompts.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/tool_prompts.ts @@ -151,4 +151,12 @@ If no relevant information is found, inform the user you could not locate the re default: 'Call this for Elastic Defend insights.', }, }, + { + promptId: 'IntegrationKnowledgeTool', + promptGroupId, + prompt: { + default: + 'Call this for knowledge from Fleet-installed integrations, which contains information on how to configure and use integrations for data ingestion.', + }, + }, ]; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts index 1b49844db9b91..e3c253930094d 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts @@ -96,6 +96,7 @@ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{ ProductDocumentationTool?: number; CustomTool?: number; EntityRiskScoreTool?: number; + IntegrationKnowledgeTool?: number; }; model?: string; isOssModel?: boolean; @@ -152,6 +153,7 @@ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{ SecurityLabsKnowledgeBaseTool: toolCountSchema, CustomTool: toolCountSchema, EntityRiskScoreTool: toolCountSchema, + IntegrationKnowledgeTool: toolCountSchema, }, }, }, diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts index ef88e24ac7693..f0648d6808b3c 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts @@ -14,6 +14,7 @@ import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_r import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; import { SECURITY_LABS_KNOWLEDGE_BASE_TOOL } from './security_labs/security_labs_tool'; import { ENTITY_RISK_SCORE_TOOL } from './entity_risk_score/entity_risk_score'; +import { INTEGRATION_KNOWLEDGE_TOOL } from './integration_knowledge/integration_knowledge_tool'; // any new tool should also be added to telemetry schema in // x-pack/solutions/security/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts @@ -27,4 +28,5 @@ export const assistantTools = [ PRODUCT_DOCUMENTATION_TOOL, SECURITY_LABS_KNOWLEDGE_BASE_TOOL, ENTITY_RISK_SCORE_TOOL, + INTEGRATION_KNOWLEDGE_TOOL, ]; diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/integration_knowledge/integration_knowledge_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/integration_knowledge/integration_knowledge_tool.test.ts new file mode 100644 index 0000000000000..a348508ad0a2e --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/integration_knowledge/integration_knowledge_tool.test.ts @@ -0,0 +1,277 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; +import { INTEGRATION_KNOWLEDGE_TOOL } from './integration_knowledge_tool'; +import type { + ContentReferencesStore, + HrefContentReference, + KnowledgeBaseEntryContentReference, +} from '@kbn/elastic-assistant-common'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import type { AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; + +const mockSearch = jest.fn(); +const mockAssistantContext = { + core: { + elasticsearch: { + client: { + asInternalUser: { + search: mockSearch, + }, + }, + }, + }, + getServerBasePath: () => '/test-base-path', +}; + +describe('IntegrationKnowledgeTool', () => { + const contentReferencesStore = newContentReferencesStoreMock(); + const defaultArgs = { + assistantContext: mockAssistantContext, + contentReferencesStore, + } as unknown as AssistantToolParams; + + beforeEach(() => { + jest.clearAllMocks(); + // Default to index existing - mock search call with size: 0 for index existence check + mockSearch.mockResolvedValue({ hits: { total: { value: 0 } } }); + }); + + describe('isSupported', () => { + it('returns true when assistantContext is provided', () => { + expect(INTEGRATION_KNOWLEDGE_TOOL.isSupported(defaultArgs)).toBe(true); + }); + + it('returns false when assistantContext is not provided', () => { + const argsWithoutContext = { contentReferencesStore } as AssistantToolParams; + expect(INTEGRATION_KNOWLEDGE_TOOL.isSupported(argsWithoutContext)).toBe(false); + }); + }); + + describe('getTool', () => { + it('returns null when not supported', async () => { + const argsWithoutContext = { contentReferencesStore } as AssistantToolParams; + const result = await INTEGRATION_KNOWLEDGE_TOOL.getTool(argsWithoutContext); + expect(result).toBeNull(); + }); + + it('returns a DynamicStructuredTool when supported and index exists', async () => { + const tool = await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs); + expect(tool).toBeDefined(); + expect(tool?.name).toBe('IntegrationKnowledgeTool'); + expect(mockSearch).toHaveBeenCalledWith({ + index: '.integration_knowledge', + size: 0, + }); + }); + + it('returns null when index does not exist', async () => { + mockSearch.mockRejectedValue(new Error('index_not_found_exception')); + const tool = await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs); + expect(tool).toBeNull(); + }); + + it('returns null when index existence check throws error', async () => { + mockSearch.mockRejectedValue(new Error('Index check failed')); + const tool = await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs); + expect(tool).toBeNull(); + }); + }); + + describe('DynamicStructuredTool', () => { + it('includes href citations for integration packages', async () => { + // First call is for index existence check + mockSearch.mockResolvedValueOnce({ hits: { total: { value: 0 } } }).mockResolvedValueOnce({ + hits: { + hits: [ + { + _id: 'test-id', + _source: { + package_name: 'nginx', + filename: 'README.md', + content: 'This is how to configure nginx integration for web server monitoring.', + version: '1.2.3', + }, + }, + ], + }, + }); + + const tool = (await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs)) as DynamicStructuredTool; + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'exampleContentReferenceId' }); + expect(reference.type).toEqual('Href'); + expect((reference as HrefContentReference).href).toEqual( + '/test-base-path/app/integrations/detail/nginx' + ); + expect((reference as HrefContentReference).label).toEqual( + 'nginx integration (README.md)' + ); + return reference; + } + ); + + const result = await tool.func({ question: 'How do I configure nginx?' }); + + expect(mockSearch).toHaveBeenCalledWith({ + index: '.integration_knowledge', + size: 10, + query: { + semantic: { + field: 'content', + query: 'How do I configure nginx?', + }, + }, + _source: ['package_name', 'filename', 'content', 'version'], + }); + + expect(result).toContain('Citation: {reference(exampleContentReferenceId)}'); + expect(result).toContain('Package: nginx (v1.2.3)'); + expect(result).toContain('File: README.md'); + expect(result).toContain('This is how to configure nginx integration'); + }); + + it('includes knowledge base citations as fallback', async () => { + // First call is for index existence check + mockSearch.mockResolvedValueOnce({ hits: { total: { value: 0 } } }).mockResolvedValueOnce({ + hits: { + hits: [ + { + _id: 'test-id', + _source: { + package_name: 'apache', + filename: 'config.yml', + content: 'Apache integration configuration details.', + }, + }, + ], + }, + }); + + const tool = (await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs)) as DynamicStructuredTool; + + // Mock the href reference creation to throw an error to trigger fallback + (contentReferencesStore.add as jest.Mock) + .mockImplementationOnce(() => { + throw new Error('Reference creation failed'); + }) + .mockImplementationOnce((creator: Parameters[0]) => { + const reference = creator({ id: 'fallbackReferenceId' }); + expect(reference.type).toEqual('KnowledgeBaseEntry'); + expect((reference as KnowledgeBaseEntryContentReference).knowledgeBaseEntryId).toEqual( + 'integrationKnowledge' + ); + expect((reference as KnowledgeBaseEntryContentReference).knowledgeBaseEntryName).toEqual( + 'Integration knowledge for apache' + ); + return reference; + }); + + const result = await tool.func({ question: 'How do I setup apache?' }); + + expect(result).toContain('Citation: {reference(fallbackReferenceId)}'); + expect(result).toContain('Package: apache'); + expect(result).toContain('File: config.yml'); + }); + + it('handles package without version', async () => { + // First call is for index existence check + mockSearch.mockResolvedValueOnce({ hits: { total: { value: 0 } } }).mockResolvedValueOnce({ + hits: { + hits: [ + { + _id: 'test-id', + _source: { + package_name: 'mysql', + filename: 'setup.md', + content: 'MySQL integration setup instructions.', + }, + }, + ], + }, + }); + + const tool = (await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs)) as DynamicStructuredTool; + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + const reference = creator({ id: 'mysqlReferenceId' }); + return reference; + } + ); + + const result = await tool.func({ question: 'MySQL setup' }); + + expect(result).toContain('Package: mysql'); + expect(result).not.toContain('(v'); + }); + + it('returns appropriate message when no results found', async () => { + // First call is for index existence check + mockSearch.mockResolvedValueOnce({ hits: { total: { value: 0 } } }).mockResolvedValueOnce({ + hits: { + hits: [], + }, + }); + + const tool = (await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs)) as DynamicStructuredTool; + + const result = await tool.func({ question: 'nonexistent integration' }); + + expect(result).toBe('[]'); + }); + + it('handles search errors gracefully', async () => { + // First call is for index existence check, second call throws error + mockSearch + .mockResolvedValueOnce({ hits: { total: { value: 0 } } }) + .mockRejectedValueOnce(new Error('Elasticsearch connection failed')); + + const tool = (await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs)) as DynamicStructuredTool; + + const result = await tool.func({ question: 'test question' }); + + expect(result).toBe( + 'Error querying integration knowledge: Elasticsearch connection failed. The integration knowledge base may not be available.' + ); + }); + + it('truncates long results to 20000 characters', async () => { + const longContent = 'A'.repeat(25000); + // First call is for index existence check + mockSearch.mockResolvedValueOnce({ hits: { total: { value: 0 } } }).mockResolvedValueOnce({ + hits: { + hits: [ + { + _id: 'test-id', + _source: { + package_name: 'large-integration', + filename: 'large-doc.md', + content: longContent, + }, + }, + ], + }, + }); + + const tool = (await INTEGRATION_KNOWLEDGE_TOOL.getTool(defaultArgs)) as DynamicStructuredTool; + + (contentReferencesStore.add as jest.Mock).mockImplementation( + (creator: Parameters[0]) => { + return creator({ id: 'largeContentId' }); + } + ); + + const result = await tool.func({ question: 'large integration' }); + + expect(result.length).toBe(20000); + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/integration_knowledge/integration_knowledge_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/integration_knowledge/integration_knowledge_tool.ts new file mode 100644 index 0000000000000..608862a7485b9 --- /dev/null +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/integration_knowledge/integration_knowledge_tool.ts @@ -0,0 +1,146 @@ +/* + * 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 { tool } from '@langchain/core/tools'; + +import { z } from '@kbn/zod'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import type { ContentReference } from '@kbn/elastic-assistant-common'; +import { contentReferenceString } from '@kbn/elastic-assistant-common'; +import { + hrefReference, + knowledgeBaseReference, +} from '@kbn/elastic-assistant-common/impl/content_references/references'; +import { Document } from 'langchain/document'; +import type { Require } from '@kbn/elastic-assistant-plugin/server/types'; +import { APP_UI_ID } from '../../../../common'; + +export type IntegrationKnowledgeToolParams = Require; + +const INTEGRATIONS_BASE_PATH = '/app/integrations/detail'; + +const toolDetails = { + // note: this description is overwritten when `getTool` is called + // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts + // local definitions can be overwritten by security-ai-prompt integration definitions + description: + 'Call this for knowledge from Fleet-installed integrations, which contains information on how to configure and use integrations for data ingestion.', + id: 'integration-knowledge-tool', + name: 'IntegrationKnowledgeTool', +}; + +const INTEGRATION_KNOWLEDGE_INDEX = '.integration_knowledge'; + +export const INTEGRATION_KNOWLEDGE_TOOL: AssistantTool = { + ...toolDetails, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is IntegrationKnowledgeToolParams => { + const { assistantContext } = params; + return assistantContext != null; + }, + async getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { assistantContext, contentReferencesStore } = params as IntegrationKnowledgeToolParams; + + // Check if the .integration_knowledge index exists before registering the tool + // This has to be done with `.search` since `.exists` and `.get` can't be performed + // with the internal system user (lack of permissions) + try { + const indexExists = await assistantContext.core.elasticsearch.client.asInternalUser.search({ + index: INTEGRATION_KNOWLEDGE_INDEX, + size: 0, + }); + if (!indexExists) { + return null; + } + } catch (error) { + // If there's an error checking the index, assume it doesn't exist and don't register the tool + return null; + } + + return tool( + async (input) => { + try { + // Search the .integration_knowledge index using semantic search on the content field + const response = await assistantContext.core.elasticsearch.client.asInternalUser.search({ + index: INTEGRATION_KNOWLEDGE_INDEX, + size: 10, + query: { + semantic: { + field: 'content', + query: input.question, + }, + }, + _source: ['package_name', 'filename', 'content', 'version'], + }); + + const citedDocs = response.hits.hits.map((hit) => { + const source = hit._source as { + package_name: string; + filename: string; + content: string; + version?: string; + }; + + let reference: ContentReference | undefined; + try { + // Create a reference to the integration details page + const packageUrl = `${assistantContext.getServerBasePath()}${INTEGRATIONS_BASE_PATH}/${ + source.package_name + }`; + const title = `${source.package_name} integration (${source.filename})`; + + reference = contentReferencesStore.add((p) => hrefReference(p.id, packageUrl, title)); + } catch (_error) { + reference = contentReferencesStore.add((p) => + knowledgeBaseReference( + p.id, + `Integration knowledge for ${source.package_name}`, + 'integrationKnowledge' + ) + ); + } + + return new Document({ + id: hit._id, + pageContent: `${contentReferenceString(reference)}\n\nPackage: ${ + source.package_name + }${source.version ? ` (v${source.version})` : ''}\nFile: ${source.filename}\n${ + source.content + }`, + metadata: { + package_name: source.package_name, + package_version: source.version, + filename: source.filename, + }, + }); + }); + + // TODO: Token pruning + const result = JSON.stringify(citedDocs).substring(0, 20000); + + return result; + } catch (error) { + return `Error querying integration knowledge: ${error.message}. The integration knowledge base may not be available.`; + } + }, + { + name: toolDetails.name, + description: params.description || toolDetails.description, + schema: z.object({ + question: z + .string() + .describe( + 'Key terms to retrieve Fleet-installed integration knowledge for, like specific integration names, configuration questions, or data ingestion topics.' + ), + }), + tags: ['integration', 'knowledge-base', 'fleet'], + } + ); + }, +}; diff --git a/x-pack/solutions/security/test/security_solution_api_integration/test_suites/genai/evaluations/trial_license_complete_tier/evaluations.ts b/x-pack/solutions/security/test/security_solution_api_integration/test_suites/genai/evaluations/trial_license_complete_tier/evaluations.ts index 61156952b97c1..f5484b7a24269 100644 --- a/x-pack/solutions/security/test/security_solution_api_integration/test_suites/genai/evaluations/trial_license_complete_tier/evaluations.ts +++ b/x-pack/solutions/security/test/security_solution_api_integration/test_suites/genai/evaluations/trial_license_complete_tier/evaluations.ts @@ -66,6 +66,21 @@ export default ({ getService }: FtrProviderContext) => { const route = routeWithNamespace(`/api/saved_objects/epm-packages/security_ai_prompts`); await supertest.delete(route).set('kbn-xsrf', 'foo'); } + + // Ensure .integration_knowledge index exists to ensure that integration knowledge tool + // is registered but doesn't intefere with evals + const INTEGRATION_KNOWLEDGE_INDEX = '.integration_knowledge'; + try { + const indexExists = await es.indices.exists({ + index: INTEGRATION_KNOWLEDGE_INDEX, + }); + if (!indexExists) { + await es.indices.create({ index: INTEGRATION_KNOWLEDGE_INDEX }); + } + } catch (e) { + // Log errors but don't fail evals + log.error(`Error creating ${INTEGRATION_KNOWLEDGE_INDEX} index: ${e}`); + } }); after(async () => {