diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts index a186b3a19051c..6f68d8f20de32 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -24,6 +24,7 @@ export const createMockClients = () => { clusterClient: core.elasticsearch.client, elasticAssistant: { actions: actionsClientMock.create(), + getRegisteredTools: jest.fn(), logger: loggingSystemMock.createLogger(), }, savedObjectsClient: core.savedObjects.client, @@ -72,6 +73,7 @@ const createElasticAssistantRequestContextMock = ( ): jest.Mocked => { return { actions: clients.elasticAssistant.actions as unknown as ActionsPluginStart, + getRegisteredTools: jest.fn(), logger: clients.elasticAssistant.logger, }; }; diff --git a/x-pack/plugins/elastic_assistant/server/index.ts b/x-pack/plugins/elastic_assistant/server/index.ts index 0cccc1ba27a7a..1775fc60528e8 100755 --- a/x-pack/plugins/elastic_assistant/server/index.ts +++ b/x-pack/plugins/elastic_assistant/server/index.ts @@ -13,6 +13,10 @@ export async function plugin(initializerContext: PluginInitializerContext) { } export type { - ElasticAssistantPluginSetup as EcsDataQualityDashboardPluginSetup, - ElasticAssistantPluginStart as EcsDataQualityDashboardPluginStart, + ElasticAssistantPluginSetup, + ElasticAssistantPluginStart, + ElasticAssistantPluginSetupDependencies, + ElasticAssistantPluginStartDependencies, + AssistantTool, + AssistantToolParams, } from './types'; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index 52a066e54de1e..d6868925cc667 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -16,7 +16,7 @@ import { KNOWLEDGE_BASE_INDEX_PATTERN } from '../../../routes/knowledge_base/con import type { AgentExecutorParams, AgentExecutorResponse } from '../executors/types'; import { withAssistantSpan } from '../tracers/with_assistant_span'; import { APMTracer } from '../tracers/apm_tracer'; -import { getApplicableTools } from '../tools'; +import { AssistantToolParams } from '../../../types'; export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor'; @@ -31,6 +31,7 @@ export const callAgentExecutor = async ({ allow, allowReplacement, assistantLangChain, + assistantTools = [], connectorId, elserId, esClient, @@ -71,7 +72,8 @@ export const callAgentExecutor = async ({ // Create a chain that uses the ELSER backed ElasticsearchStore, override k=10 for esql query generation for now const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever(10)); - const tools: Tool[] = getApplicableTools({ + // Fetch any applicable tools that the source plugin may have registered + const assistantToolParams: AssistantToolParams = { allow, allowReplacement, alertsIndexPattern, @@ -83,7 +85,8 @@ export const callAgentExecutor = async ({ replacements, request, size, - }); + }; + const tools: Tool[] = assistantTools.flatMap((tool) => tool.getTool(assistantToolParams) ?? []); logger.debug(`applicable tools: ${JSON.stringify(tools.map((t) => t.name).join(', '), null, 2)}`); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index bd6a38da5cdc4..e7824e2822f8a 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -12,6 +12,7 @@ import { Logger } from '@kbn/logging'; import { KibanaRequest } from '@kbn/core-http-server'; import type { LangChainTracer } from 'langchain/callbacks'; import { RequestBody, ResponseBody } from '../types'; +import type { AssistantTool } from '../../../types'; export interface AgentExecutorParams { alertsIndexPattern?: string; @@ -19,6 +20,7 @@ export interface AgentExecutorParams { allow?: string[]; allowReplacement?: string[]; assistantLangChain: boolean; + assistantTools?: AssistantTool[]; connectorId: string; esClient: ElasticsearchClient; kbResource: string | undefined; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.test.ts deleted file mode 100644 index cffe31cbcfd39..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.test.ts +++ /dev/null @@ -1,105 +0,0 @@ -/* - * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import type { KibanaRequest } from '@kbn/core-http-server'; -import { DynamicTool } from 'langchain/tools'; -import { omit } from 'lodash/fp'; - -import { getAlertCountsTool } from './get_alert_counts_tool'; -import type { RequestBody } from '../../types'; - -describe('getAlertCountsTool', () => { - const alertsIndexPattern = 'alerts-index'; - const esClient = { - search: jest.fn().mockResolvedValue({}), - } as unknown as ElasticsearchClient; - const replacements = { key: 'value' }; - const request = { - body: { - assistantLangChain: false, - alertsIndexPattern: '.alerts-security.alerts-default', - allow: ['@timestamp', 'cloud.availability_zone', 'user.name'], - allowReplacement: ['user.name'], - replacements, - size: 20, - }, - } as unknown as KibanaRequest; - - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => { - const tool: DynamicTool = getAlertCountsTool({ - alertsIndexPattern, - esClient, - replacements, - request, - }) as DynamicTool; - - await tool.func(''); - - expect(esClient.search).toHaveBeenCalledWith({ - aggs: { statusBySeverity: { terms: { field: 'kibana.alert.severity' } } }, - index: ['alerts-index'], - query: { - bool: { - filter: [ - { - bool: { - filter: [{ match_phrase: { 'kibana.alert.workflow_status': 'open' } }], - must_not: [{ exists: { field: 'kibana.alert.building_block_type' } }], - }, - }, - { range: { '@timestamp': { gte: 'now/d', lte: 'now/d' } } }, - ], - }, - }, - size: 0, - }); - }); - - it('returns null when the request is missing required anonymization parameters', () => { - const requestWithMissingParams = omit('body.allow', request) as unknown as KibanaRequest< - unknown, - unknown, - RequestBody - >; - - const tool = getAlertCountsTool({ - alertsIndexPattern, - esClient, - replacements, - request: requestWithMissingParams, - }); - - expect(tool).toBeNull(); - }); - - it('returns null when the alertsIndexPattern is undefined', () => { - const tool = getAlertCountsTool({ - // alertsIndexPattern is undefined - esClient, - replacements, - request, - }); - - expect(tool).toBeNull(); - }); - - it('returns a tool instance with the expected tags', () => { - const tool = getAlertCountsTool({ - alertsIndexPattern, - esClient, - replacements, - request, - }) as DynamicTool; - - expect(tool.tags).toEqual(['alerts', 'alerts-count']); - }); -}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts deleted file mode 100644 index 9c5ec2555d2e5..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts +++ /dev/null @@ -1,47 +0,0 @@ -/* - * 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 { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { KibanaRequest } from '@kbn/core-http-server'; -import { DynamicTool, Tool } from 'langchain/tools'; - -import { getAlertsCountQuery } from './get_alert_counts_query'; -import { requestHasRequiredAnonymizationParams } from '../../helpers'; -import type { RequestBody } from '../../types'; - -export const ALERT_COUNTS_TOOL_DESCRIPTION = - 'Call this for the counts of last 24 hours of open alerts in the environment, grouped by their severity'; - -export const getAlertCountsTool = ({ - alertsIndexPattern, - esClient, - replacements, - request, -}: { - alertsIndexPattern?: string; - esClient: ElasticsearchClient; - replacements?: Record; - request: KibanaRequest; -}): Tool | null => { - if (!requestHasRequiredAnonymizationParams(request) || alertsIndexPattern == null) { - return null; - } - - return new DynamicTool({ - name: 'alert-counts', - description: ALERT_COUNTS_TOOL_DESCRIPTION, - func: async () => { - const query = getAlertsCountQuery(alertsIndexPattern); - - const result = await esClient.search(query); - - return JSON.stringify(result); - }, - tags: ['alerts', 'alerts-count'], - }); -}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts deleted file mode 100644 index ccd97b7deb088..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts +++ /dev/null @@ -1,55 +0,0 @@ -/* - * 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 { RetrievalQAChain } from 'langchain/chains'; -import { DynamicTool } from 'langchain/tools'; - -import { getEsqlLanguageKnowledgeBaseTool } from './get_esql_language_knowledge_base_tool'; - -const chain = {} as RetrievalQAChain; - -describe('getEsqlLanguageKnowledgeBaseTool', () => { - it('returns null if assistantLangChain is false', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: false, - chain, - modelExists: true, - }); - - expect(tool).toBeNull(); - }); - - it('returns null if modelExists is false (the ELSER model is not installed)', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: true, - chain, - modelExists: false, // <-- ELSER model is not installed - }); - - expect(tool).toBeNull(); - }); - - it('should return a Tool instance if assistantLangChain and modelExists are true', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: true, - modelExists: true, - chain, - }); - - expect(tool?.name).toEqual('ESQLKnowledgeBaseTool'); - }); - - it('should return a tool with the expected tags', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: true, - chain, - modelExists: true, - }) as DynamicTool; - - expect(tool.tags).toEqual(['esql', 'query-generation', 'knowledge-base']); - }); -}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts deleted file mode 100644 index f49551261a79d..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts +++ /dev/null @@ -1,29 +0,0 @@ -/* - * 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 { RetrievalQAChain } from 'langchain/chains'; -import { ChainTool, Tool } from 'langchain/tools'; - -export const getEsqlLanguageKnowledgeBaseTool = ({ - assistantLangChain, - modelExists, - chain, -}: { - assistantLangChain: boolean; - chain: RetrievalQAChain; - /** true when the ELSER model is installed */ - modelExists: boolean; -}): Tool | null => - assistantLangChain && modelExists - ? new ChainTool({ - name: 'ESQLKnowledgeBaseTool', - description: - 'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language.', - chain, - tags: ['esql', 'query-generation', 'knowledge-base'], - }) - : null; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.test.ts deleted file mode 100644 index cb0e79c0558d4..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.test.ts +++ /dev/null @@ -1,58 +0,0 @@ -/* - * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { KibanaRequest } from '@kbn/core-http-server'; -import { RetrievalQAChain } from 'langchain/chains'; - -import { RequestBody } from '../types'; -import { getApplicableTools } from '.'; - -describe('getApplicableTools', () => { - const alertsIndexPattern = 'alerts-index'; - const esClient = { - search: jest.fn().mockResolvedValue({}), - } as unknown as ElasticsearchClient; - const modelExists = true; // the ELSER model is installed - const onNewReplacements = jest.fn(); - const replacements = { key: 'value' }; - const request = { - body: { - assistantLangChain: true, - alertsIndexPattern: '.alerts-security.alerts-default', - allow: ['@timestamp', 'cloud.availability_zone', 'user.name'], - allowReplacement: ['user.name'], - replacements, - size: 20, - }, - } as unknown as KibanaRequest; - const chain = {} as unknown as RetrievalQAChain; - - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('should return an array of applicable tools', () => { - const tools = getApplicableTools({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - assistantLangChain: request.body.assistantLangChain, - chain, - esClient, - modelExists, - onNewReplacements, - replacements, - request, - size: request.body.size, - }); - - const minExpectedTools = 3; // 3 tools are currently implemented - - expect(tools.length).toBeGreaterThanOrEqual(minExpectedTools); - }); -}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts deleted file mode 100644 index 048bdfabce437..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { KibanaRequest } from '@kbn/core-http-server'; -import { RetrievalQAChain } from 'langchain/chains'; -import { Tool } from 'langchain/tools'; - -import { getAlertCountsTool } from './alert_counts/get_alert_counts_tool'; -import { getEsqlLanguageKnowledgeBaseTool } from './esql_language_knowledge_base/get_esql_language_knowledge_base_tool'; -import { getOpenAndAcknowledgedAlertsTool } from './open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool'; -import type { RequestBody } from '../types'; - -export interface GetApplicableTools { - alertsIndexPattern?: string; - allow?: string[]; - allowReplacement?: string[]; - assistantLangChain: boolean; - chain: RetrievalQAChain; - esClient: ElasticsearchClient; - modelExists: boolean; - onNewReplacements?: (newReplacements: Record) => void; - replacements?: Record; - request: KibanaRequest; - size?: number; -} - -export const getApplicableTools = ({ - alertsIndexPattern, - allow, - allowReplacement, - assistantLangChain, - chain, - esClient, - modelExists, - onNewReplacements, - replacements, - request, - size, -}: GetApplicableTools): Tool[] => - [ - getEsqlLanguageKnowledgeBaseTool({ assistantLangChain, chain, modelExists }) ?? [], - getAlertCountsTool({ - alertsIndexPattern, - esClient, - replacements, - request, - }) ?? [], - getOpenAndAcknowledgedAlertsTool({ - alertsIndexPattern, - allow, - allowReplacement, - esClient, - onNewReplacements, - replacements, - request, - size, - }) ?? [], - ].flat(); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool.test.ts deleted file mode 100644 index 8c769ef284df3..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool.test.ts +++ /dev/null @@ -1,218 +0,0 @@ -/* - * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import type { KibanaRequest } from '@kbn/core-http-server'; -import { DynamicTool } from 'langchain/tools'; -import { omit } from 'lodash/fp'; - -import { getOpenAndAcknowledgedAlertsTool } from './get_open_and_acknowledged_alerts_tool'; -import { mockAlertsFieldsApi } from '../../../../__mocks__/alerts'; -import type { RequestBody } from '../../types'; -import { MAX_SIZE } from './helpers'; - -describe('getOpenAndAcknowledgedAlertsTool', () => { - const alertsIndexPattern = 'alerts-index'; - const esClient = { - search: jest.fn().mockResolvedValue(mockAlertsFieldsApi), - } as unknown as ElasticsearchClient; - const replacements = { key: 'value' }; - const request = { - body: { - assistantLangChain: false, - alertsIndexPattern: '.alerts-security.alerts-default', - allow: ['@timestamp', 'cloud.availability_zone', 'user.name'], - allowReplacement: ['user.name'], - replacements, - size: 20, - }, - } as unknown as KibanaRequest; - - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => { - const tool: DynamicTool = getOpenAndAcknowledgedAlertsTool({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - size: request.body.size, - }) as DynamicTool; - - await tool.func(''); - - expect(esClient.search).toHaveBeenCalledWith({ - allow_no_indices: true, - body: { - _source: false, - fields: [ - { - field: '@timestamp', - include_unmapped: true, - }, - { - field: 'cloud.availability_zone', - include_unmapped: true, - }, - { - field: 'user.name', - include_unmapped: true, - }, - ], - query: { - bool: { - filter: [ - { - bool: { - filter: [ - { - bool: { - should: [ - { - match_phrase: { - 'kibana.alert.workflow_status': 'open', - }, - }, - { - match_phrase: { - 'kibana.alert.workflow_status': 'acknowledged', - }, - }, - ], - minimum_should_match: 1, - }, - }, - { - range: { - '@timestamp': { - format: 'strict_date_optional_time', - gte: 'now-1d/d', - lte: 'now/d', - }, - }, - }, - ], - must: [], - must_not: [ - { - exists: { - field: 'kibana.alert.building_block_type', - }, - }, - ], - should: [], - }, - }, - ], - }, - }, - runtime_mappings: {}, - size: 20, - sort: [ - { - 'kibana.alert.risk_score': { - order: 'desc', - }, - }, - { - '@timestamp': { - order: 'desc', - }, - }, - ], - }, - ignore_unavailable: true, - index: ['alerts-index'], - }); - }); - - it('returns null when the request is missing required anonymization parameters', () => { - const requestWithMissingParams = omit('body.allow', request) as unknown as KibanaRequest< - unknown, - unknown, - RequestBody - >; - - const tool = getOpenAndAcknowledgedAlertsTool({ - alertsIndexPattern, - allow: requestWithMissingParams.body.allow, - allowReplacement: requestWithMissingParams.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request: requestWithMissingParams, - size: requestWithMissingParams.body.size, - }); - - expect(tool).toBeNull(); - }); - - it('returns null when alertsIndexPattern is undefined', () => { - const tool = getOpenAndAcknowledgedAlertsTool({ - // alertsIndexPattern is undefined - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - size: request.body.size, - }); - - expect(tool).toBeNull(); - }); - - it('returns null when size is undefined', () => { - const tool = getOpenAndAcknowledgedAlertsTool({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - // size is undefined - }); - - expect(tool).toBeNull(); - }); - - it('returns null when size out of range', () => { - const tool = getOpenAndAcknowledgedAlertsTool({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - size: MAX_SIZE + 1, // <-- size is out of range - }); - - expect(tool).toBeNull(); - }); - - it('returns a tool instance with the expected tags', () => { - const tool = getOpenAndAcknowledgedAlertsTool({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - size: request.body.size, - }) as DynamicTool; - - expect(tool.tags).toEqual(['alerts', 'open-alerts']); - }); -}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool.ts deleted file mode 100644 index c4748230f86ca..0000000000000 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_tool.ts +++ /dev/null @@ -1,90 +0,0 @@ -/* - * 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 { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { KibanaRequest } from '@kbn/core-http-server'; -import { getAnonymizedValue, transformRawData } from '@kbn/elastic-assistant-common'; -import { DynamicTool, Tool } from 'langchain/tools'; -import { requestHasRequiredAnonymizationParams } from '../../helpers'; -import { RequestBody } from '../../types'; - -import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query'; -import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers'; - -export const OPEN_ALERTS_TOOL_DESCRIPTION = - 'Call this for knowledge about the latest n open alerts (sorted by `kibana.alert.risk_score`) in the environment, or when answering questions about open alerts'; - -/** - * Returns a tool for querying open and acknowledged alerts, or null if the - * request doesn't have all the required parameters. - */ -export const getOpenAndAcknowledgedAlertsTool = ({ - alertsIndexPattern, - allow, - allowReplacement, - esClient, - onNewReplacements, - replacements, - request, - size, -}: { - alertsIndexPattern?: string; - allow?: string[]; - allowReplacement?: string[]; - esClient: ElasticsearchClient; - onNewReplacements?: (newReplacements: Record) => void; - replacements?: Record; - request: KibanaRequest; - size?: number; -}): Tool | null => { - if ( - !requestHasRequiredAnonymizationParams(request) || - alertsIndexPattern == null || - size == null || - sizeIsOutOfRange(size) - ) { - return null; - } - - return new DynamicTool({ - name: 'open-alerts', - description: OPEN_ALERTS_TOOL_DESCRIPTION, - func: async () => { - const query = getOpenAndAcknowledgedAlertsQuery({ - alertsIndexPattern, - allow: allow ?? [], - size, - }); - - const result = await esClient.search(query); - - // Accumulate replacements locally so we can, for example use the same - // replacement for a hostname when we see it in multiple alerts: - let localReplacements = { ...replacements }; - const localOnNewReplacements = (newReplacements: Record) => { - localReplacements = { ...localReplacements, ...newReplacements }; // update the local state - - onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements - }; - - return JSON.stringify( - result.hits?.hits?.map((x) => - transformRawData({ - allow: allow ?? [], - allowReplacement: allowReplacement ?? [], - currentReplacements: localReplacements, // <-- the latest local replacements - getAnonymizedValue, - onNewReplacements: localOnNewReplacements, // <-- the local callback - rawData: getRawDataOrDefault(x.fields), - }) - ) - ); - }, - tags: ['alerts', 'open-alerts'], - }); -}; diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index a0df339695885..bd06165e57284 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -18,12 +18,14 @@ import { import { once } from 'lodash'; import { + AssistantTool, ElasticAssistantPluginSetup, ElasticAssistantPluginSetupDependencies, ElasticAssistantPluginStart, ElasticAssistantPluginStartDependencies, ElasticAssistantRequestHandlerContext, GetElser, + PLUGIN_ID, } from './types'; import { deleteKnowledgeBaseRoute, @@ -32,6 +34,13 @@ import { postEvaluateRoute, postKnowledgeBaseRoute, } from './routes'; +import { appContextService, GetRegisteredTools } from './services/app_context'; + +interface CreateRouteHandlerContextParams { + core: CoreSetup; + logger: Logger; + getRegisteredTools: GetRegisteredTools; +} export class ElasticAssistantPlugin implements @@ -48,15 +57,20 @@ export class ElasticAssistantPlugin this.logger = initializerContext.logger.get(); } - private createRouteHandlerContext = ( - core: CoreSetup, - logger: Logger - ): IContextProvider => { + private createRouteHandlerContext = ({ + core, + logger, + getRegisteredTools, + }: CreateRouteHandlerContextParams): IContextProvider< + ElasticAssistantRequestHandlerContext, + typeof PLUGIN_ID + > => { return async function elasticAssistantRouteHandlerContext(context, request) { const [_, pluginsStart] = await core.getStartServices(); return { actions: pluginsStart.actions, + getRegisteredTools, logger, }; }; @@ -65,16 +79,15 @@ export class ElasticAssistantPlugin public setup(core: CoreSetup, plugins: ElasticAssistantPluginSetupDependencies) { this.logger.debug('elasticAssistant: Setup'); const router = core.http.createRouter(); - - core.http.registerRouteHandlerContext< - ElasticAssistantRequestHandlerContext, - 'elasticAssistant' - >( - 'elasticAssistant', - this.createRouteHandlerContext( - core as CoreSetup, - this.logger - ) + core.http.registerRouteHandlerContext( + PLUGIN_ID, + this.createRouteHandlerContext({ + core: core as CoreSetup, + logger: this.logger, + getRegisteredTools: (pluginName: string) => { + return appContextService.getRegisteredTools(pluginName); + }, + }) ); const getElserId: GetElser = once( @@ -94,16 +107,45 @@ export class ElasticAssistantPlugin postEvaluateRoute(router, getElserId); return { actions: plugins.actions, + getRegisteredTools: (pluginName: string) => { + return appContextService.getRegisteredTools(pluginName); + }, }; } public start(core: CoreStart, plugins: ElasticAssistantPluginStartDependencies) { this.logger.debug('elasticAssistant: Started'); + appContextService.start({ logger: this.logger }); return { + /** + * Actions plugin start contract + */ actions: plugins.actions, + + /** + * Get the registered tools for a given plugin name. + * @param pluginName + */ + getRegisteredTools: (pluginName: string) => { + return appContextService.getRegisteredTools(pluginName); + }, + + /** + * Register tools to be used by the Elastic Assistant for a given plugin. Use the plugin name that + * corresponds to your application as defined in the `x-kbn-context` header of requests made from your + * application. + * + * @param pluginName + * @param tools + */ + registerTools: (pluginName: string, tools: AssistantTool[]) => { + return appContextService.registerTools(pluginName, tools); + }, }; } - public stop() {} + public stop() { + appContextService.stop(); + } } diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index 39e3233f63744..f5ec4c4555a76 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -100,6 +100,11 @@ export const postEvaluateRoute = ( throwIfSystemAction: false, }); + // Fetch any tools registered by the request's originating plugin + const assistantTools = (await context.elasticAssistant).getRegisteredTools( + 'securitySolution' + ); + // Get a scoped esClient for passing to the agents for retrieval, and // writing results to the output index const esClient = (await context.core).elasticsearch.client.asCurrentUser; @@ -142,6 +147,7 @@ export const postEvaluateRoute = ( AGENT_EXECUTOR_MAP[agentName]({ actions, assistantLangChain: true, + assistantTools, connectorId, esClient, elserId, diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts new file mode 100644 index 0000000000000..384e1f8865736 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts @@ -0,0 +1,67 @@ +/* + * 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 { KibanaRequest } from '@kbn/core-http-server'; +import type { RequestBody } from '../lib/langchain/types'; + +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from './helpers'; + +describe('getPluginNameFromRequest', () => { + const contextRequestHeaderEncoded = encodeURIComponent( + JSON.stringify({ + type: 'application', + name: 'superSolution', + url: '/kbn/app/super/rules/id/163fa5a4-d72a-45fa-8142-8edc298ecd17/alerts', + page: 'app', + id: 'new', + }) + ); + + const request = { + headers: { + 'x-kbn-context': contextRequestHeaderEncoded, + }, + } as unknown as KibanaRequest; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('extracts plugin name from "x-kbn-context" request header', async () => { + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); + expect(pluginName).toEqual('superSolution'); + }); + + it('fails to extract plugin name from undefined "x-kbn-context" request header, falls back to default provided', async () => { + const invalidRequest = { + headers: { + 'x-kbn-context': undefined, + }, + } as unknown as KibanaRequest; + const pluginName = getPluginNameFromRequest({ + request: invalidRequest, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); + expect(pluginName).toEqual(DEFAULT_PLUGIN_NAME); + }); + + it('fails to extract plugin name from malformed "x-kbn-context" invalidRequest header, falls back to default provided', async () => { + const invalidRequest = { + headers: { + 'x-kbn-context': 'asdfku', + }, + } as unknown as KibanaRequest; + const pluginName = getPluginNameFromRequest({ + request: invalidRequest, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); + expect(pluginName).toEqual(DEFAULT_PLUGIN_NAME); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts new file mode 100644 index 0000000000000..99d4493c16cca --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -0,0 +1,53 @@ +/* + * 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 { KibanaRequest } from '@kbn/core-http-server'; +import { Logger } from '@kbn/core/server'; +import { RequestBody } from '../lib/langchain/types'; + +interface GetPluginNameFromRequestParams { + request: KibanaRequest; + defaultPluginName: string; + logger?: Logger; +} + +export const DEFAULT_PLUGIN_NAME = 'securitySolutionUI'; + +/** + * Attempts to extract the plugin name the request originated from using the request headers. + * + * Note from Kibana Core: This is not a 100% fit solution, though, because plugins can run in the background, + * or even use other plugins’ helpers (ie, APM can use the infra helpers to call a third plugin) + * + * Should suffice for our purposes here with where the Elastic Assistant is currently used, but if needing a + * dedicated solution, the core folks said to reach out. + * + * @param logger optional logger to log any errors + * @param defaultPluginName default plugin name to use if unable to determine from request + * @param request Kibana Request + * + * @returns plugin name + */ +export const getPluginNameFromRequest = ({ + logger, + defaultPluginName, + request, +}: GetPluginNameFromRequestParams): string => { + try { + const contextHeader = request.headers['x-kbn-context']; + if (contextHeader != null) { + return JSON.parse( + decodeURIComponent(Array.isArray(contextHeader) ? contextHeader[0] : contextHeader) + )?.name; + } + } catch (err) { + logger?.error( + `Error determining source plugin for selecting tools, using ${defaultPluginName}.` + ); + } + return defaultPluginName; +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts index 537ead452c8ea..0f7d964eea480 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts @@ -56,6 +56,7 @@ jest.mock('../lib/langchain/execute_custom_llm_chain', () => ({ const mockContext = { elasticAssistant: { actions: jest.fn(), + getRegisteredTools: jest.fn(() => []), logger: loggingSystemMock.createLogger(), }, core: { diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index ed68f3526a112..9c1d8601da532 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -23,6 +23,7 @@ import { import { ElasticAssistantRequestHandlerContext, GetElser } from '../types'; import { ESQL_RESOURCE } from './knowledge_base/constants'; import { callAgentExecutor } from '../lib/langchain/execute_custom_llm_chain'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from './helpers'; export const postActionsConnectorExecuteRoute = ( router: IRouter, @@ -58,6 +59,14 @@ export const postActionsConnectorExecuteRoute = ( // TODO: Add `traceId` to actions request when calling via langchain logger.debug('Executing via langchain, assistantLangChain: true'); + // Fetch any tools registered by the request's originating plugin + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const assistantTools = (await context.elasticAssistant).getRegisteredTools(pluginName); + // get a scoped esClient for assistant memory const esClient = (await context.core).elasticsearch.client.asCurrentUser; @@ -79,6 +88,7 @@ export const postActionsConnectorExecuteRoute = ( allowReplacement: request.body.allowReplacement, actions, assistantLangChain: request.body.assistantLangChain, + assistantTools, connectorId, elserId, esClient, diff --git a/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts new file mode 100644 index 0000000000000..621995d3452be --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts @@ -0,0 +1,87 @@ +/* + * 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 { appContextService, ElasticAssistantAppContext } from './app_context'; +import { loggerMock } from '@kbn/logging-mocks'; +import { AssistantTool } from '../types'; + +// Mock Logger +const mockLogger = loggerMock.create(); + +// Mock ElasticAssistantAppContext +const mockAppContext: ElasticAssistantAppContext = { + logger: mockLogger, +}; + +describe('AppContextService', () => { + const toolOne: AssistantTool = { + id: 'tool-one', + name: 'ToolOne', + description: 'Description 1', + sourceRegister: 'Source1', + isSupported: jest.fn(), + getTool: jest.fn(), + }; + const toolTwo: AssistantTool = { + id: 'tool-two', + name: 'ToolTwo', + description: 'Description 2', + sourceRegister: 'Source2', + isSupported: jest.fn(), + getTool: jest.fn(), + }; + + beforeEach(() => { + appContextService.stop(); + jest.clearAllMocks(); + }); + + describe('starting and stopping', () => { + it('should clear registered tools when stopped ', () => { + appContextService.start(mockAppContext); + appContextService.registerTools('super', [toolOne]); + appContextService.stop(); + + expect(appContextService.getRegisteredTools('super').length).toBe(0); + }); + }); + + describe('registering tools', () => { + it('should register and get tools for a single plugin', () => { + const pluginName = 'pluginName'; + + appContextService.start(mockAppContext); + appContextService.registerTools(pluginName, [toolOne, toolTwo]); + + // Check if getRegisteredTools returns the correct tools + const retrievedTools = appContextService.getRegisteredTools(pluginName); + expect(retrievedTools).toEqual([toolOne, toolTwo]); + }); + + it('should register and get tools for multiple plugins', () => { + const pluginOne = 'plugin1'; + const pluginTwo = 'plugin2'; + + appContextService.start(mockAppContext); + appContextService.registerTools(pluginOne, [toolOne]); + appContextService.registerTools(pluginTwo, [toolTwo]); + + expect(appContextService.getRegisteredTools(pluginOne)).toEqual([toolOne]); + expect(appContextService.getRegisteredTools(pluginTwo)).toEqual([toolTwo]); + }); + + it('should not add the same tool twice', () => { + const pluginName = 'pluginName'; + + appContextService.start(mockAppContext); + appContextService.registerTools(pluginName, [toolOne]); + appContextService.registerTools(pluginName, [toolOne]); + + expect(appContextService.getRegisteredTools(pluginName).length).toEqual(1); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/services/app_context.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.ts new file mode 100644 index 0000000000000..bd7a7c0cc3203 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -0,0 +1,69 @@ +/* + * 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 { Logger } from '@kbn/core/server'; +import type { AssistantTool } from '../types'; + +export type PluginName = string; +export type RegisteredToolsStorage = Map>; +export type GetRegisteredTools = (pluginName: string) => AssistantTool[]; +export interface ElasticAssistantAppContext { + logger: Logger; +} + +/** + * Service for managing context specific to the Elastic Assistant + * + * Inspired by `AppContextService` impl from fleet plugin: x-pack/plugins/fleet/server/services/app_context.ts + */ +class AppContextService { + private logger: Logger | undefined; + private registeredTools: RegisteredToolsStorage = new Map>(); + + public start(appContext: ElasticAssistantAppContext) { + this.logger = appContext.logger; + } + + public stop() { + this.registeredTools.clear(); + } + + /** + * Register tools to be used by the Elastic Assistant + * + * @param pluginName + * @param tools + */ + public registerTools(pluginName: string, tools: AssistantTool[]) { + this.logger?.debug('AppContextService:registerTools'); + this.logger?.debug(`pluginName: ${pluginName}`); + this.logger?.debug(`tools: ${tools.map((tool) => tool.name).join(', ')}`); + + if (!this.registeredTools.has(pluginName)) { + this.logger?.debug('plugin has no tools, making new set'); + this.registeredTools.set(pluginName, new Set()); + } + tools.forEach((tool) => this.registeredTools.get(pluginName)?.add(tool)); + } + + /** + * Get the registered tools + * + * @param pluginName + */ + public getRegisteredTools(pluginName: string): AssistantTool[] { + const tools = Array.from(this.registeredTools?.get(pluginName) ?? new Set()); + + this.logger?.debug('AppContextService:getRegisteredTools'); + this.logger?.debug(`pluginName: ${pluginName}`); + this.logger?.debug(`tools: ${tools.map((tool) => tool.name).join(', ')}`); + + return tools; + } +} + +export const appContextService = new AppContextService(); diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index ed9081c084420..5be8a35a275b9 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -16,6 +16,13 @@ import type { SavedObjectsClientContract, } from '@kbn/core/server'; import { type MlPluginSetup } from '@kbn/ml-plugin/server'; +import { Tool } from 'langchain/dist/tools/base'; +import { RetrievalQAChain } from 'langchain/chains'; +import { ElasticsearchClient } from '@kbn/core/server'; +import { RequestBody } from './lib/langchain/types'; +import type { GetRegisteredTools } from './services/app_context'; + +export const PLUGIN_ID = 'elasticAssistant' as const; /** The plugin setup interface */ export interface ElasticAssistantPluginSetup { @@ -25,6 +32,17 @@ export interface ElasticAssistantPluginSetup { /** The plugin start interface */ export interface ElasticAssistantPluginStart { actions: ActionsPluginStart; + /** + * Register tools to be used by the elastic assistant + * @param pluginName Name of the plugin the tool should be registered to + * @param tools AssistantTools to be registered with for the given plugin + */ + registerTools: (pluginName: string, tools: AssistantTool[]) => void; + /** + * Get the registered tools + * @param pluginName Name of the plugin to get the tools for + */ + getRegisteredTools: GetRegisteredTools; } export interface ElasticAssistantPluginSetupDependencies { @@ -37,6 +55,7 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { actions: ActionsPluginStart; + getRegisteredTools: GetRegisteredTools; logger: Logger; } @@ -51,3 +70,30 @@ export type GetElser = ( request: KibanaRequest, savedObjectsClient: SavedObjectsClientContract ) => Promise | never; + +/** + * Interfaces for registering tools to be used by the elastic assistant + */ + +export interface AssistantTool { + id: string; + name: string; + description: string; + sourceRegister: string; + isSupported: (params: AssistantToolParams) => boolean; + getTool: (params: AssistantToolParams) => Tool | null; +} + +export interface AssistantToolParams { + alertsIndexPattern?: string; + allow?: string[]; + allowReplacement?: string[]; + assistantLangChain: boolean; + chain: RetrievalQAChain; + esClient: ElasticsearchClient; + modelExists: boolean; + onNewReplacements?: (newReplacements: Record) => void; + replacements?: Record; + request: KibanaRequest; + size?: number; +} diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 5cad4d4b52141..53616fc2dc2b0 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -21,7 +21,6 @@ "@kbn/securitysolution-io-ts-utils", "@kbn/actions-plugin", "@kbn/elastic-assistant", - "@kbn/elastic-assistant-common", "@kbn/logging-mocks", "@kbn/core-elasticsearch-server-mocks", "@kbn/core-logging-server-mocks", diff --git a/x-pack/plugins/security_solution/server/assistant/jest.config.js b/x-pack/plugins/security_solution/server/assistant/jest.config.js new file mode 100644 index 0000000000000..787dae7ce8d68 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/jest.config.js @@ -0,0 +1,19 @@ +/* + * 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. + */ + +module.exports = { + preset: '@kbn/test', + rootDir: '../../../../..', + roots: ['/x-pack/plugins/security_solution/server/assistant'], + coverageDirectory: + '/target/kibana-coverage/jest/x-pack/plugins/security_solution/server/assistant', + coverageReporters: ['text', 'html'], + collectCoverageFrom: [ + '/x-pack/plugins/security_solution/server/assistant/**/*.{ts,tsx}', + ], + moduleNameMapper: require('../__mocks__/module_name_map'), +}; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts new file mode 100644 index 0000000000000..715d656944c86 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts @@ -0,0 +1,162 @@ +/* + * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { DynamicTool } from 'langchain/tools'; +import { omit } from 'lodash/fp'; + +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; +import { ALERT_COUNTS_TOOL } from './alert_counts_tool'; +import type { RetrievalQAChain } from 'langchain/chains'; + +describe('AlertCountsTool', () => { + const alertsIndexPattern = 'alerts-index'; + const esClient = { + search: jest.fn().mockResolvedValue({}), + } as unknown as ElasticsearchClient; + const replacements = { key: 'value' }; + const request = { + body: { + assistantLangChain: false, + alertsIndexPattern: '.alerts-security.alerts-default', + allow: ['@timestamp', 'cloud.availability_zone', 'user.name'], + allowReplacement: ['user.name'], + replacements, + size: 20, + }, + } as unknown as KibanaRequest; + const assistantLangChain = true; + const chain = {} as unknown as RetrievalQAChain; + const modelExists = true; + const rest = { + assistantLangChain, + chain, + modelExists, + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('isSupported', () => { + it('returns false when the alertsIndexPattern is undefined', () => { + const params = { + esClient, + request, + ...rest, + }; + + expect(ALERT_COUNTS_TOOL.isSupported(params)).toBe(false); + }); + + it('returns false when the request is missing required anonymization parameters', () => { + const requestMissingAnonymizationParams = { + body: { + assistantLangChain: false, + alertsIndexPattern: '.alerts-security.alerts-default', + size: 20, + }, + } as unknown as KibanaRequest; + const params = { + esClient, + request: requestMissingAnonymizationParams, + ...rest, + }; + + expect(ALERT_COUNTS_TOOL.isSupported(params)).toBe(false); + }); + + it('returns true if alertsIndexPattern is defined and request includes required anonymization parameters', () => { + const params = { + alertsIndexPattern, + esClient, + request, + ...rest, + }; + + expect(ALERT_COUNTS_TOOL.isSupported(params)).toBe(true); + }); + }); + + describe('getTool', () => { + it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => { + const tool: DynamicTool = ALERT_COUNTS_TOOL.getTool({ + alertsIndexPattern, + esClient, + replacements, + request, + ...rest, + }) as DynamicTool; + + await tool.func(''); + + expect(esClient.search).toHaveBeenCalledWith({ + aggs: { statusBySeverity: { terms: { field: 'kibana.alert.severity' } } }, + index: ['alerts-index'], + query: { + bool: { + filter: [ + { + bool: { + filter: [{ match_phrase: { 'kibana.alert.workflow_status': 'open' } }], + must_not: [{ exists: { field: 'kibana.alert.building_block_type' } }], + }, + }, + { range: { '@timestamp': { gte: 'now/d', lte: 'now/d' } } }, + ], + }, + }, + size: 0, + }); + }); + + it('returns null when the request is missing required anonymization parameters', () => { + const requestWithMissingParams = omit('body.allow', request) as unknown as KibanaRequest< + unknown, + unknown, + RequestBody + >; + + const tool = ALERT_COUNTS_TOOL.getTool({ + alertsIndexPattern, + esClient, + replacements, + request: requestWithMissingParams, + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('returns null when the alertsIndexPattern is undefined', () => { + const tool = ALERT_COUNTS_TOOL.getTool({ + // alertsIndexPattern is undefined + esClient, + replacements, + request, + + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('returns a tool instance with the expected tags', () => { + const tool = ALERT_COUNTS_TOOL.getTool({ + alertsIndexPattern, + esClient, + replacements, + request, + + ...rest, + }) as DynamicTool; + + expect(tool.tags).toEqual(['alerts', 'alerts-count']); + }); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts new file mode 100644 index 0000000000000..e1c265c3dc239 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts @@ -0,0 +1,47 @@ +/* + * 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 { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; +import { DynamicTool } from 'langchain/tools'; + +import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { getAlertsCountQuery } from './get_alert_counts_query'; +import { APP_UI_ID } from '../../../../common'; + +export interface AlertCountsToolParams extends AssistantToolParams { + alertsIndexPattern: string; +} +export const ALERT_COUNTS_TOOL_DESCRIPTION = + 'Call this for the counts of last 24 hours of open alerts in the environment, grouped by their severity'; + +export const ALERT_COUNTS_TOOL: AssistantTool = { + id: 'alert-counts-tool', + name: 'AlertCountsTool', + description: ALERT_COUNTS_TOOL_DESCRIPTION, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is AlertCountsToolParams => { + const { request, alertsIndexPattern } = params; + return requestHasRequiredAnonymizationParams(request) && alertsIndexPattern != null; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + const { alertsIndexPattern, esClient } = params as AlertCountsToolParams; + return new DynamicTool({ + name: 'AlertCountsTool', + description: ALERT_COUNTS_TOOL_DESCRIPTION, + func: async () => { + const query = getAlertsCountQuery(alertsIndexPattern); + + const result = await esClient.search(query); + + return JSON.stringify(result); + }, + tags: ['alerts', 'alerts-count'], + }); + }, +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_query.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.test.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_query.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.test.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_query.ts rename to x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts new file mode 100644 index 0000000000000..19e3bf7e7b40a --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts @@ -0,0 +1,109 @@ +/* + * 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 { RetrievalQAChain } from 'langchain/chains'; +import type { DynamicTool } from 'langchain/tools'; +import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base_tool'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; + +describe('EsqlLanguageKnowledgeBaseTool', () => { + const chain = {} as RetrievalQAChain; + const esClient = { + search: jest.fn().mockResolvedValue({}), + } as unknown as ElasticsearchClient; + const request = { + body: { + assistantLangChain: false, + alertsIndexPattern: '.alerts-security.alerts-default', + allow: ['@timestamp', 'cloud.availability_zone', 'user.name'], + allowReplacement: ['user.name'], + replacements: { key: 'value' }, + size: 20, + }, + } as unknown as KibanaRequest; + const rest = { + chain, + esClient, + request, + }; + + describe('isSupported', () => { + it('returns false if assistantLangChain is false', () => { + const params = { + assistantLangChain: false, + modelExists: true, + ...rest, + }; + + expect(ESQL_KNOWLEDGE_BASE_TOOL.isSupported(params)).toBe(false); + }); + + it('returns false if modelExists is false (the ELSER model is not installed)', () => { + const params = { + assistantLangChain: true, + modelExists: false, // <-- ELSER model is not installed + ...rest, + }; + + expect(ESQL_KNOWLEDGE_BASE_TOOL.isSupported(params)).toBe(false); + }); + + it('returns true if assistantLangChain and modelExists are true', () => { + const params = { + assistantLangChain: true, + modelExists: true, + ...rest, + }; + + expect(ESQL_KNOWLEDGE_BASE_TOOL.isSupported(params)).toBe(true); + }); + }); + + describe('getTool', () => { + it('returns null if assistantLangChain is false', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: false, + modelExists: true, + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('returns null if modelExists is false (the ELSER model is not installed)', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: true, + modelExists: false, // <-- ELSER model is not installed + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('should return a Tool instance if assistantLangChain and modelExists are true', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: true, + modelExists: true, + ...rest, + }); + + expect(tool?.name).toEqual('ESQLKnowledgeBaseTool'); + }); + + it('should return a tool with the expected tags', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: true, + modelExists: true, + ...rest, + }) as DynamicTool; + + expect(tool.tags).toEqual(['esql', 'query-generation', 'knowledge-base']); + }); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts new file mode 100644 index 0000000000000..3dc7dfd8d976c --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts @@ -0,0 +1,35 @@ +/* + * 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 { ChainTool } from 'langchain/tools'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { APP_UI_ID } from '../../../../common'; + +export type EsqlKnowledgeBaseToolParams = AssistantToolParams; + +export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { + id: 'esql-knowledge-base-tool', + name: 'ESQLKnowledgeBaseTool', + description: + 'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language.', + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => { + const { assistantLangChain, modelExists } = params; + return assistantLangChain && modelExists; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + const { chain } = params as EsqlKnowledgeBaseToolParams; + return new ChainTool({ + name: 'ESQLKnowledgeBaseTool', + description: + 'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language.', + chain, + tags: ['esql', 'query-generation', 'knowledge-base'], + }); + }, +}; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts new file mode 100644 index 0000000000000..047c84ceddf3b --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts @@ -0,0 +1,22 @@ +/* + * 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 { getAssistantTools } from '.'; + +describe('getAssistantTools', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('should return an array of applicable tools', () => { + const tools = getAssistantTools(); + + const minExpectedTools = 3; // 3 tools are currently implemented + + expect(tools.length).toBeGreaterThanOrEqual(minExpectedTools); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts new file mode 100644 index 0000000000000..790e674a4b390 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -0,0 +1,17 @@ +/* + * 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 { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; +import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/esql_language_knowledge_base_tool'; +import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool'; +import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; + +export const getAssistantTools = (): AssistantTool[] => [ + ALERT_COUNTS_TOOL, + ESQL_KNOWLEDGE_BASE_TOOL, + OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL, +]; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.test.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.test.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/helpers.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/helpers.test.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/helpers.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/helpers.test.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/helpers.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/helpers.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_and_acknowledged_alerts/helpers.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/helpers.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts new file mode 100644 index 0000000000000..c485c001704eb --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts @@ -0,0 +1,298 @@ +/* + * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { DynamicTool } from 'langchain/tools'; +import { omit } from 'lodash/fp'; + +import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts_tool'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; +import { MAX_SIZE } from './helpers'; +import type { RetrievalQAChain } from 'langchain/chains'; +import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; + +describe('OpenAndAcknowledgedAlertsTool', () => { + const alertsIndexPattern = 'alerts-index'; + const esClient = { + search: jest.fn().mockResolvedValue(mockAlertsFieldsApi), + } as unknown as ElasticsearchClient; + const replacements = { key: 'value' }; + const request = { + body: { + assistantLangChain: false, + alertsIndexPattern: '.alerts-security.alerts-default', + allow: ['@timestamp', 'cloud.availability_zone', 'user.name'], + allowReplacement: ['user.name'], + replacements, + size: 20, + }, + } as unknown as KibanaRequest; + const assistantLangChain = true; + const chain = {} as unknown as RetrievalQAChain; + const modelExists = true; + const rest = { + assistantLangChain, + esClient, + chain, + modelExists, + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('isSupported', () => { + it('returns false when alertsIndexPattern is undefined', () => { + const params = { + request, + size: 20, + ...rest, + }; + + expect(OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.isSupported(params)).toBe(false); + }); + + it('returns false when the request is missing required anonymization parameters', () => { + const requestMissingAnonymizationParams = { + body: { + assistantLangChain: false, + alertsIndexPattern: '.alerts-security.alerts-default', + size: 20, + }, + } as unknown as KibanaRequest; + const params = { + request: requestMissingAnonymizationParams, + ...rest, + }; + + expect(OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.isSupported(params)).toBe(false); + }); + + it('returns false when size is undefined', () => { + const params = { + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + ...rest, + }; + + expect(OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.isSupported(params)).toBe(false); + }); + + it('returns false when size is out of range', () => { + const params = { + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + size: MAX_SIZE + 1, // <-- size is out of range + + ...rest, + }; + + expect(OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.isSupported(params)).toBe(false); + }); + + it('returns true when anonymization fields, alertsIndexPattern, and size within reange is provided', () => { + const params = { + alertsIndexPattern, + size: 20, + request, + ...rest, + }; + + expect(OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.isSupported(params)).toBe(true); + }); + }); + describe('getTool', () => { + it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => { + const tool: DynamicTool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + size: request.body.size, + ...rest, + }) as DynamicTool; + + await tool.func(''); + + expect(esClient.search).toHaveBeenCalledWith({ + allow_no_indices: true, + body: { + _source: false, + fields: [ + { + field: '@timestamp', + include_unmapped: true, + }, + { + field: 'cloud.availability_zone', + include_unmapped: true, + }, + { + field: 'user.name', + include_unmapped: true, + }, + ], + query: { + bool: { + filter: [ + { + bool: { + filter: [ + { + bool: { + should: [ + { + match_phrase: { + 'kibana.alert.workflow_status': 'open', + }, + }, + { + match_phrase: { + 'kibana.alert.workflow_status': 'acknowledged', + }, + }, + ], + minimum_should_match: 1, + }, + }, + { + range: { + '@timestamp': { + format: 'strict_date_optional_time', + gte: 'now-1d/d', + lte: 'now/d', + }, + }, + }, + ], + must: [], + must_not: [ + { + exists: { + field: 'kibana.alert.building_block_type', + }, + }, + ], + should: [], + }, + }, + ], + }, + }, + runtime_mappings: {}, + size: 20, + sort: [ + { + 'kibana.alert.risk_score': { + order: 'desc', + }, + }, + { + '@timestamp': { + order: 'desc', + }, + }, + ], + }, + ignore_unavailable: true, + index: ['alerts-index'], + }); + }); + + it('returns null when the request is missing required anonymization parameters', () => { + const requestWithMissingParams = omit('body.allow', request) as unknown as KibanaRequest< + unknown, + unknown, + RequestBody + >; + + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: requestWithMissingParams.body.allow, + allowReplacement: requestWithMissingParams.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request: requestWithMissingParams, + size: requestWithMissingParams.body.size, + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('returns null when alertsIndexPattern is undefined', () => { + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + // alertsIndexPattern is undefined + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + size: request.body.size, + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('returns null when size is undefined', () => { + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + ...rest, + // size is undefined + }); + + expect(tool).toBeNull(); + }); + + it('returns null when size out of range', () => { + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + size: MAX_SIZE + 1, // <-- size is out of range + ...rest, + }); + + expect(tool).toBeNull(); + }); + + it('returns a tool instance with the expected tags', () => { + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + onNewReplacements: jest.fn(), + replacements, + request, + size: request.body.size, + ...rest, + }) as DynamicTool; + + expect(tool.tags).toEqual(['alerts', 'open-and-acknowledged-alerts']); + }); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts new file mode 100644 index 0000000000000..210fdd42c1555 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts @@ -0,0 +1,93 @@ +/* + * 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 { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; +import { getAnonymizedValue, transformRawData } from '@kbn/elastic-assistant-common'; +import { DynamicTool } from 'langchain/tools'; +import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; + +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query'; +import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers'; +import { APP_UI_ID } from '../../../../common'; + +export interface OpenAndAcknowledgedAlertsToolParams extends AssistantToolParams { + alertsIndexPattern: string; + size: number; +} + +export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION = + 'Call this for knowledge about the latest n open and acknowledged alerts (sorted by `kibana.alert.risk_score`) in the environment, or when answering questions about open alerts'; + +/** + * Returns a tool for querying open and acknowledged alerts, or null if the + * request doesn't have all the required parameters. + */ +export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = { + id: 'open-and-acknowledged-alerts-tool', + name: 'OpenAndAcknowledgedAlertsTool', + description: OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is OpenAndAcknowledgedAlertsToolParams => { + const { alertsIndexPattern, request, size } = params; + return ( + requestHasRequiredAnonymizationParams(request) && + alertsIndexPattern != null && + size != null && + !sizeIsOutOfRange(size) + ); + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { + alertsIndexPattern, + allow, + allowReplacement, + esClient, + onNewReplacements, + replacements, + size, + } = params as OpenAndAcknowledgedAlertsToolParams; + return new DynamicTool({ + name: 'OpenAndAcknowledgedAlertsTool', + description: OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION, + func: async () => { + const query = getOpenAndAcknowledgedAlertsQuery({ + alertsIndexPattern, + allow: allow ?? [], + size, + }); + + const result = await esClient.search(query); + + // Accumulate replacements locally so we can, for example use the same + // replacement for a hostname when we see it in multiple alerts: + let localReplacements = { ...replacements }; + const localOnNewReplacements = (newReplacements: Record) => { + localReplacements = { ...localReplacements, ...newReplacements }; // update the local state + + onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements + }; + + return JSON.stringify( + result.hits?.hits?.map((x) => + transformRawData({ + allow: allow ?? [], + allowReplacement: allowReplacement ?? [], + currentReplacements: localReplacements, // <-- the latest local replacements + getAnonymizedValue, + onNewReplacements: localOnNewReplacements, // <-- the local callback + rawData: getRawDataOrDefault(x.fields), + }) + ) + ); + }, + tags: ['alerts', 'open-and-acknowledged-alerts'], + }); + }, +}; diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 4841ec9fed73d..227a7ea7e1439 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -46,7 +46,7 @@ import { AppClientFactory } from './client'; import type { ConfigType } from './config'; import { createConfig } from './config'; import { initUiSettings } from './ui_settings'; -import { APP_ID, DEFAULT_ALERTS_INDEX, SERVER_APP_ID } from '../common/constants'; +import { APP_ID, APP_UI_ID, DEFAULT_ALERTS_INDEX, SERVER_APP_ID } from '../common/constants'; import { registerEndpointRoutes } from './endpoint/routes/metadata'; import { registerPolicyRoutes } from './endpoint/routes/policy'; import { registerActionRoutes } from './endpoint/routes/actions'; @@ -114,6 +114,7 @@ import { allRiskScoreIndexPattern, } from '../common/entity_analytics/risk_engine'; import { isEndpointPackageV2 } from '../common/endpoint/utils/package_v2'; +import { getAssistantTools } from './assistant/tools'; export type { SetupPlugins, StartPlugins, PluginSetup, PluginStart } from './plugin_contract'; @@ -509,6 +510,9 @@ export class Plugin implements ISecuritySolutionPlugin { this.licensing$ = plugins.licensing.license$; + // Assistant Tool and Feature Registration + plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools()); + if (this.lists && plugins.taskManager && plugins.fleet) { // Exceptions, Artifacts and Manifests start const taskManager = plugins.taskManager; diff --git a/x-pack/plugins/security_solution/server/plugin_contract.ts b/x-pack/plugins/security_solution/server/plugin_contract.ts index b59d9aac37efe..8370e405c6807 100644 --- a/x-pack/plugins/security_solution/server/plugin_contract.ts +++ b/x-pack/plugins/security_solution/server/plugin_contract.ts @@ -41,6 +41,7 @@ import type { CloudExperimentsPluginStart } from '@kbn/cloud-experiments-plugin/ import type { SharePluginStart } from '@kbn/share-plugin/server'; import type { GuidedOnboardingPluginSetup } from '@kbn/guided-onboarding-plugin/server'; import type { PluginSetup as UnifiedSearchServerPluginSetup } from '@kbn/unified-search-plugin/server'; +import type { ElasticAssistantPluginStart } from '@kbn/elastic-assistant-plugin/server'; import type { AppFeaturesService } from './lib/app_features_service/app_features_service'; import type { ExperimentalFeatures } from '../common'; @@ -72,6 +73,7 @@ export interface SecuritySolutionPluginStartDependencies { cloudExperiments?: CloudExperimentsPluginStart; data: DataPluginStart; dataViews: DataViewsPluginStart; + elasticAssistant: ElasticAssistantPluginStart; eventLog: IEventLogClientService; fleet?: FleetPluginStart; licensing: LicensingPluginStart; diff --git a/x-pack/plugins/security_solution/tsconfig.json b/x-pack/plugins/security_solution/tsconfig.json index 4400f4b8c31e3..2158fe97996dc 100644 --- a/x-pack/plugins/security_solution/tsconfig.json +++ b/x-pack/plugins/security_solution/tsconfig.json @@ -84,6 +84,7 @@ "@kbn/kibana-react-plugin", "@kbn/ecs-data-quality-dashboard", "@kbn/elastic-assistant", + "@kbn/elastic-assistant-plugin", "@kbn/data-views-plugin", "@kbn/datemath", "@kbn/ui-theme", @@ -183,6 +184,7 @@ "@kbn/zod-helpers", "@kbn/core-http-common", "@kbn/search-errors", - "@kbn/stack-connectors-plugin" + "@kbn/stack-connectors-plugin", + "@kbn/elastic-assistant-common" ] }