From 9aeed7cd82b76a2d7a117b97a2ecd2553610f6d6 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 29 Nov 2023 17:35:14 -0700 Subject: [PATCH 01/13] Support registration of arbitrary tools on plugin start --- .../server/__mocks__/request_context.ts | 2 + .../plugins/elastic_assistant/server/index.ts | 6 +- .../execute_custom_llm_chain/index.ts | 2 + .../server/lib/langchain/executors/types.ts | 2 + .../elastic_assistant/server/plugin.ts | 55 +++++++++++---- .../routes/post_actions_connector_execute.ts | 7 ++ .../server/services/app_context.ts | 70 +++++++++++++++++++ .../plugins/elastic_assistant/server/types.ts | 15 ++++ .../get_security_assistant_tools.tsx | 24 +++++++ .../security_solution/server/plugin.ts | 4 ++ .../server/plugin_contract.ts | 2 + 11 files changed, 172 insertions(+), 17 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/services/app_context.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx 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 a186b3a19051ce..6f68d8f20de32d 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 0cccc1ba27a7a3..721f5da706e957 100755 --- a/x-pack/plugins/elastic_assistant/server/index.ts +++ b/x-pack/plugins/elastic_assistant/server/index.ts @@ -13,6 +13,8 @@ export async function plugin(initializerContext: PluginInitializerContext) { } export type { - ElasticAssistantPluginSetup as EcsDataQualityDashboardPluginSetup, - ElasticAssistantPluginStart as EcsDataQualityDashboardPluginStart, + ElasticAssistantPluginSetup, + ElasticAssistantPluginStart, + ElasticAssistantPluginSetupDependencies, + ElasticAssistantPluginStartDependencies, } 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 3a27226dc804de..6865dfe4cba5e8 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 @@ -31,6 +31,7 @@ export const callAgentExecutor = async ({ langChainMessages, llmType, logger, + registeredTools = [], request, elserId, kbResource, @@ -77,6 +78,7 @@ export const callAgentExecutor = async ({ chain, tags: ['esql', 'query-generation', 'knowledge-base'], }), + ...registeredTools, ]; const executor = await initializeAgentExecutorWithOptions(tools, llm, { 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 e4d564f677e148..43b378c7bb5af3 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 @@ -11,6 +11,7 @@ import { BaseMessage } from 'langchain/schema'; import { Logger } from '@kbn/logging'; import { KibanaRequest } from '@kbn/core-http-server'; import type { LangChainTracer } from 'langchain/callbacks'; +import type { Tool } from 'langchain/tools'; import { RequestBody, ResponseBody } from '../types'; export interface AgentExecutorParams { @@ -21,6 +22,7 @@ export interface AgentExecutorParams { langChainMessages: BaseMessage[]; llmType?: string; logger: Logger; + registeredTools?: Tool[]; request: KibanaRequest; elserId?: string; traceOptions?: TraceOptions; diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 827b428c97803d..ab1d5c0d23169d 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -17,6 +17,7 @@ import { } from '@kbn/core/server'; import { once } from 'lodash'; +import type { Tool } from 'langchain/tools'; import { ElasticAssistantPluginSetup, ElasticAssistantPluginSetupDependencies, @@ -24,6 +25,7 @@ import { 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,28 @@ 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: plugins.actions, + getRegisteredTools: (pluginName: string) => { + return appContextService.getRegisteredTools(pluginName); + }, + registerTool: (pluginName: string, tool: Tool) => { + return appContextService.registerTools(pluginName, tool); + }, }; } - public stop() {} + public stop() { + appContextService.stop(); + } } 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 299d8ade24a3f2..4a413aabac6111 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 @@ -54,6 +54,12 @@ 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 registeredTools = (await context.elasticAssistant).getRegisteredTools( + 'securitySolution' + ); + logger.debug(`Registered tools: ${registeredTools.map((tool) => tool.name).join(', ')}`); + // get a scoped esClient for assistant memory const esClient = (await context.core).elasticsearch.client.asCurrentUser; @@ -70,6 +76,7 @@ export const postActionsConnectorExecuteRoute = ( esClient, langChainMessages, logger, + registeredTools, request, elserId, kbResource: ESQL_RESOURCE, 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 00000000000000..f25fe66c106515 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -0,0 +1,70 @@ +/* + * 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 { Tool } from 'langchain/tools'; + +export type PluginName = string; +export type RegisteredToolsStorage = Map>; +export type RegisterTools = (pluginName: string, tool: Tool) => void; +export type GetRegisteredTools = (pluginName: string) => Tool[]; + +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 a tool to be used by the Elastic Assistant + * + * @param pluginName + * @param tool + */ + public registerTools(pluginName: string, tool: Tool) { + this.logger?.debug('AppContextService:registerTools'); + this.logger?.debug(`pluginName: ${pluginName}`); + this.logger?.debug(`tool: ${tool.name}`); + + if (!this.registeredTools.has(pluginName)) { + this.logger?.debug('plugin has no tools, making new set'); + this.registeredTools.set(pluginName, new Set()); + } + this.registeredTools.get(pluginName)?.add(tool); + } + + /** + * Get the registered tools + * @param pluginName + */ + public getRegisteredTools(pluginName: string): Tool[] { + 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 ed9081c0844208..b0adfa673d98d6 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -16,6 +16,9 @@ import type { SavedObjectsClientContract, } from '@kbn/core/server'; import { type MlPluginSetup } from '@kbn/ml-plugin/server'; +import { Tool } from 'langchain/dist/tools/base'; + +export const PLUGIN_ID = 'elasticAssistant' as const; /** The plugin setup interface */ export interface ElasticAssistantPluginSetup { @@ -25,6 +28,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 tool The tool to register + */ + registerTool: (pluginName: string, tool: Tool) => void; + /** + * Get the registered tools + * @param pluginName Name of the plugin to get the tools for + */ + getRegisteredTools: (pluginName: string) => Tool[]; } export interface ElasticAssistantPluginSetupDependencies { @@ -37,6 +51,7 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { actions: ActionsPluginStart; + getRegisteredTools: (pluginName: string) => Tool[]; logger: Logger; } diff --git a/x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx b/x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx new file mode 100644 index 00000000000000..5c5adb1a5dc1f1 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx @@ -0,0 +1,24 @@ +/* + * 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 { Tool } from 'langchain/tools'; +import { DynamicTool } from 'langchain/tools'; + +/** + * This function returns a list of tools that will be available in the Elastic AI Assistant for Security Solution users. + */ +export const getSecurityAssistantTools = (): Tool => { + return new DynamicTool({ + name: 'ListFavoriteFoods', + verbose: true, + description: + 'Call this function to get a list of my favorite foods. Input is an empty string, output is a string containing a comma separated list of favorite foods.', + func: async (input: string): Promise => { + return 'pizza, pizza, pizza!'; + }, + }); +}; diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index c3e20042d224f1..a60b0ab94e7810 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -111,6 +111,7 @@ import { registerRiskScoringTask } from './lib/entity_analytics/risk_engine/task import { registerProtectionUpdatesNoteRoutes } from './endpoint/routes/protection_updates_note'; import { latestRiskScoreIndexPattern, allRiskScoreIndexPattern } from '../common/risk_engine'; import { isEndpointPackageV2 } from '../common/endpoint/utils/package_v2'; +import { getSecurityAssistantTools } from './assistant/get_security_assistant_tools'; export type { SetupPlugins, StartPlugins, PluginSetup, PluginStart } from './plugin_contract'; @@ -506,6 +507,9 @@ export class Plugin implements ISecuritySolutionPlugin { this.licensing$ = plugins.licensing.license$; + // Assistant Tool and Feature Registration + plugins.elasticAssistant.registerTool('securitySolution', getSecurityAssistantTools()); + 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 b59d9aac37efef..8370e405c6807c 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; From 5df33d50ed9ba777c6b9824369cfe67f6bf3a42f Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Thu, 30 Nov 2023 00:49:09 +0000 Subject: [PATCH 02/13] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/security_solution/tsconfig.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/x-pack/plugins/security_solution/tsconfig.json b/x-pack/plugins/security_solution/tsconfig.json index 55a36a020fa08e..70940c2b8a5541 100644 --- a/x-pack/plugins/security_solution/tsconfig.json +++ b/x-pack/plugins/security_solution/tsconfig.json @@ -183,6 +183,7 @@ "@kbn/unified-doc-viewer-plugin", "@kbn/shared-ux-error-boundary", "@kbn/zod-helpers", - "@kbn/core-http-common" + "@kbn/core-http-common", + "@kbn/elastic-assistant-plugin" ] } From 38a39411de273c631577c1e773722c07221f8030 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 29 Nov 2023 21:26:44 -0700 Subject: [PATCH 03/13] Update mocks --- .../server/routes/post_actions_connector_execute.test.ts | 1 + 1 file changed, 1 insertion(+) 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 507246670833c0..6420f6e74a68f7 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: { From 52d57cf7c28510448d414cd982e6ab70899439e1 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 8 Dec 2023 10:51:36 -0700 Subject: [PATCH 04/13] Move new RAG on alerts tools to Security Solution plugin --- .../plugins/elastic_assistant/server/index.ts | 2 ++ .../execute_custom_llm_chain/index.ts | 33 +++++++++---------- .../server/lib/langchain/executors/types.ts | 4 +-- .../elastic_assistant/server/plugin.ts | 6 ++-- .../routes/post_actions_connector_execute.ts | 7 ++-- .../server/services/app_context.ts | 27 +++++++-------- .../plugins/elastic_assistant/server/types.ts | 28 ++++++++++++++-- .../get_security_assistant_tools.tsx | 24 -------------- .../get_alert_counts_query.test.ts | 0 .../alert_counts/get_alert_counts_query.ts | 0 .../get_alert_counts_tool.test.ts | 6 ++-- .../alert_counts/get_alert_counts_tool.ts | 11 ++++--- ..._esql_language_knowledge_base_tool.test.ts | 0 .../get_esql_language_knowledge_base_tool.ts | 5 +-- .../server/assistant}/tools/index.test.ts | 8 ++--- .../server/assistant}/tools/index.ts | 28 ++++------------ .../open_alerts/get_open_alerts_query.test.ts | 0 .../open_alerts/get_open_alerts_query.ts | 0 .../open_alerts/get_open_alerts_tool.test.ts | 8 ++--- .../tools/open_alerts/get_open_alerts_tool.ts | 11 ++++--- .../tools/open_alerts/helpers.test.ts | 0 .../assistant}/tools/open_alerts/helpers.ts | 0 .../security_solution/server/plugin.ts | 4 +-- 23 files changed, 99 insertions(+), 113 deletions(-) delete mode 100644 x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/alert_counts/get_alert_counts_query.test.ts (100%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/alert_counts/get_alert_counts_query.ts (100%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/alert_counts/get_alert_counts_tool.test.ts (92%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/alert_counts/get_alert_counts_tool.ts (74%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts (100%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts (85%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/index.test.ts (85%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/index.ts (59%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/open_alerts/get_open_alerts_query.test.ts (100%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/open_alerts/get_open_alerts_query.ts (100%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/open_alerts/get_open_alerts_tool.test.ts (94%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/open_alerts/get_open_alerts_tool.ts (86%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/open_alerts/helpers.test.ts (100%) rename x-pack/plugins/{elastic_assistant/server/lib/langchain => security_solution/server/assistant}/tools/open_alerts/helpers.ts (100%) diff --git a/x-pack/plugins/elastic_assistant/server/index.ts b/x-pack/plugins/elastic_assistant/server/index.ts index 721f5da706e957..bd55705d546eed 100755 --- a/x-pack/plugins/elastic_assistant/server/index.ts +++ b/x-pack/plugins/elastic_assistant/server/index.ts @@ -17,4 +17,6 @@ export type { ElasticAssistantPluginStart, ElasticAssistantPluginSetupDependencies, ElasticAssistantPluginStartDependencies, + GetApplicableTools, + GetApplicableToolsParams, } 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 f01c01613addc4..5baa65636dce0b 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,6 @@ 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'; export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor'; @@ -39,7 +38,7 @@ export const callAgentExecutor = async ({ llmType, logger, onNewReplacements, - registeredTools = [], + getApplicableTools = (params: GetRegisteredToolParams) => [], replacements, request, size, @@ -72,22 +71,20 @@ 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({ - allow, - allowReplacement, - alertsIndexPattern, - assistantLangChain, - chain, - esClient, - modelExists, - onNewReplacements, - replacements, - request, - size, - }), - ...registeredTools, - ]; + // Fetch any applicable tools that the source plugin may have registered + const tools: Tool[] = getApplicableTools({ + allow, + allowReplacement, + alertsIndexPattern, + assistantLangChain, + chain, + esClient, + modelExists, + onNewReplacements, + replacements, + request, + size, + }); 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 533efab8e22307..bd19352d9c346a 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 @@ -11,8 +11,8 @@ import { BaseMessage } from 'langchain/schema'; import { Logger } from '@kbn/logging'; import { KibanaRequest } from '@kbn/core-http-server'; import type { LangChainTracer } from 'langchain/callbacks'; -import type { Tool } from 'langchain/tools'; import { RequestBody, ResponseBody } from '../types'; +import type { GetApplicableTools } from '../../../types'; export interface AgentExecutorParams { alertsIndexPattern?: string; @@ -22,12 +22,12 @@ export interface AgentExecutorParams { assistantLangChain: boolean; connectorId: string; esClient: ElasticsearchClient; + getApplicableTools?: GetApplicableTools; kbResource: string | undefined; langChainMessages: BaseMessage[]; llmType?: string; logger: Logger; onNewReplacements?: (newReplacements: Record) => void; - registeredTools?: Tool[]; replacements?: Record; request: KibanaRequest; size?: number; diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 72bd2b65d106c9..b9574e0cd12d32 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -17,13 +17,13 @@ import { } from '@kbn/core/server'; import { once } from 'lodash'; -import type { Tool } from 'langchain/tools'; import { ElasticAssistantPluginSetup, ElasticAssistantPluginSetupDependencies, ElasticAssistantPluginStart, ElasticAssistantPluginStartDependencies, ElasticAssistantRequestHandlerContext, + GetApplicableTools, GetElser, PLUGIN_ID, } from './types'; @@ -122,8 +122,8 @@ export class ElasticAssistantPlugin getRegisteredTools: (pluginName: string) => { return appContextService.getRegisteredTools(pluginName); }, - registerTool: (pluginName: string, tool: Tool) => { - return appContextService.registerTools(pluginName, tool); + registerTools: (pluginName: string, getApplicableTools: GetApplicableTools) => { + return appContextService.registerTools(pluginName, getApplicableTools); }, }; } 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 671b805b3888c2..e4cf642ba15ebf 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 @@ -58,11 +58,10 @@ 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 registeredTools = (await context.elasticAssistant).getRegisteredTools( + // Fetch the function for fetching any tools registered by the request's originating plugin + const getApplicableTools = (await context.elasticAssistant).getRegisteredTools( 'securitySolution' ); - logger.debug(`Registered tools: ${registeredTools.map((tool) => tool.name).join(', ')}`); // get a scoped esClient for assistant memory const esClient = (await context.core).elasticsearch.client.asCurrentUser; @@ -88,11 +87,11 @@ export const postActionsConnectorExecuteRoute = ( connectorId, elserId, esClient, + getApplicableTools, kbResource: ESQL_RESOURCE, langChainMessages, logger, onNewReplacements, - registeredTools, request, replacements: request.body.replacements, size: request.body.size, diff --git a/x-pack/plugins/elastic_assistant/server/services/app_context.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.ts index f25fe66c106515..0f10aa793d735f 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -7,10 +7,11 @@ import type { Logger } from '@kbn/core/server'; import type { Tool } from 'langchain/tools'; +import type { GetApplicableTools, GetApplicableToolsParams } from '../types'; export type PluginName = string; -export type RegisteredToolsStorage = Map>; -export type RegisterTools = (pluginName: string, tool: Tool) => void; +export type RegisteredToolsStorage = Map; +export type RegisterTools = (pluginName: string, getApplicableTools: GetApplicableTools) => void; export type GetRegisteredTools = (pluginName: string) => Tool[]; export interface ElasticAssistantAppContext { @@ -35,35 +36,35 @@ class AppContextService { } /** - * Register a tool to be used by the Elastic Assistant + * Register tools to be used by the Elastic Assistant * * @param pluginName - * @param tool + * @param getApplicableTools */ - public registerTools(pluginName: string, tool: Tool) { + public registerTools(pluginName: string, getApplicableTools: GetApplicableTools) { this.logger?.debug('AppContextService:registerTools'); this.logger?.debug(`pluginName: ${pluginName}`); - this.logger?.debug(`tool: ${tool.name}`); if (!this.registeredTools.has(pluginName)) { - this.logger?.debug('plugin has no tools, making new set'); - this.registeredTools.set(pluginName, new Set()); + this.logger?.debug('plugin has no tools, setting "getApplicableTools"'); + } else { + this.logger?.debug('plugin already has tools, overriding "getApplicableTools"'); } - this.registeredTools.get(pluginName)?.add(tool); + this.registeredTools.set(pluginName, getApplicableTools); } /** * Get the registered tools * @param pluginName */ - public getRegisteredTools(pluginName: string): Tool[] { - const tools = Array.from(this.registeredTools?.get(pluginName) ?? new Set()); + public getRegisteredTools(pluginName: string): GetApplicableTools { + const getApplicableTools = + this.registeredTools?.get(pluginName) ?? ((params: GetApplicableToolsParams) => []); this.logger?.debug('AppContextService:getRegisteredTools'); this.logger?.debug(`pluginName: ${pluginName}`); - this.logger?.debug(`tools: ${tools.map((tool) => tool.name).join(', ')}`); - return tools; + return getApplicableTools; } } diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index b0adfa673d98d6..6cb58ecf5ad466 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -17,9 +17,28 @@ import type { } 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'; export const PLUGIN_ID = 'elasticAssistant' as const; +export interface GetApplicableToolsParams { + 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 type GetApplicableTools = (params: GetApplicableToolsParams) => Tool[]; + /** The plugin setup interface */ export interface ElasticAssistantPluginSetup { actions: ActionsPluginSetup; @@ -31,9 +50,9 @@ export interface ElasticAssistantPluginStart { /** * Register tools to be used by the elastic assistant * @param pluginName Name of the plugin the tool should be registered to - * @param tool The tool to register + * @param getApplicableTools Function that returns the tools for the specified plugin */ - registerTool: (pluginName: string, tool: Tool) => void; + registerTools: (pluginName: string, getApplicableTools: GetApplicableTools) => void; /** * Get the registered tools * @param pluginName Name of the plugin to get the tools for @@ -51,7 +70,10 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { actions: ActionsPluginStart; - getRegisteredTools: (pluginName: string) => Tool[]; + getRegisteredTools: ( + pluginName: string, + getApplicableToolsParams: GetApplicableToolsParams + ) => Tool[]; logger: Logger; } diff --git a/x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx b/x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx deleted file mode 100644 index 5c5adb1a5dc1f1..00000000000000 --- a/x-pack/plugins/security_solution/server/assistant/get_security_assistant_tools.tsx +++ /dev/null @@ -1,24 +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 { Tool } from 'langchain/tools'; -import { DynamicTool } from 'langchain/tools'; - -/** - * This function returns a list of tools that will be available in the Elastic AI Assistant for Security Solution users. - */ -export const getSecurityAssistantTools = (): Tool => { - return new DynamicTool({ - name: 'ListFavoriteFoods', - verbose: true, - description: - 'Call this function to get a list of my favorite foods. Input is an empty string, output is a string containing a comma separated list of favorite foods.', - func: async (input: string): Promise => { - return 'pizza, pizza, pizza!'; - }, - }); -}; 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/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts similarity index 92% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts index cffe31cbcfd397..343d4df4a79374 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts @@ -5,13 +5,13 @@ * 2.0. */ -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; -import { DynamicTool } from 'langchain/tools'; +import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; import { getAlertCountsTool } from './get_alert_counts_tool'; -import type { RequestBody } from '../../types'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; describe('getAlertCountsTool', () => { const alertsIndexPattern = 'alerts-index'; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts similarity index 74% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts index 9c5ec2555d2e5c..62b13e10fab76c 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/alert_counts/get_alert_counts_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts @@ -6,13 +6,14 @@ */ 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 type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { Tool } from 'langchain/tools'; +import { DynamicTool } from 'langchain/tools'; +import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; 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'; 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/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts 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/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts similarity index 85% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts index f49551261a79d5..eaf73dc0785603 100644 --- 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/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts @@ -5,8 +5,9 @@ * 2.0. */ -import { RetrievalQAChain } from 'langchain/chains'; -import { ChainTool, Tool } from 'langchain/tools'; +import type { RetrievalQAChain } from 'langchain/chains'; +import type { Tool } from 'langchain/tools'; +import { ChainTool } from 'langchain/tools'; export const getEsqlLanguageKnowledgeBaseTool = ({ assistantLangChain, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts similarity index 85% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/index.test.ts index cb0e79c0558d4b..daaa5ec3edcd3c 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts @@ -5,11 +5,11 @@ * 2.0. */ -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { KibanaRequest } from '@kbn/core-http-server'; -import { RetrievalQAChain } from 'langchain/chains'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { RetrievalQAChain } from 'langchain/chains'; -import { RequestBody } from '../types'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; import { getApplicableTools } from '.'; describe('getApplicableTools', () => { diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts similarity index 59% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts rename to x-pack/plugins/security_solution/server/assistant/tools/index.ts index edc9c264b636ad..b532a8183fa66d 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -5,31 +5,17 @@ * 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 type { Tool } from 'langchain/tools'; +import type { + GetApplicableTools, + GetApplicableToolsParams, +} from '@kbn/elastic-assistant-plugin/server'; import { getAlertCountsTool } from './alert_counts/get_alert_counts_tool'; import { getEsqlLanguageKnowledgeBaseTool } from './esql_language_knowledge_base/get_esql_language_knowledge_base_tool'; import { getOpenAlertsTool } from './open_alerts/get_open_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 = ({ +export const getApplicableTools: GetApplicableTools = ({ alertsIndexPattern, allow, allowReplacement, @@ -41,7 +27,7 @@ export const getApplicableTools = ({ replacements, request, size, -}: GetApplicableTools): Tool[] => +}: GetApplicableToolsParams): Tool[] => [ getEsqlLanguageKnowledgeBaseTool({ assistantLangChain, chain, modelExists }) ?? [], getAlertCountsTool({ diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_query.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.test.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_query.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.test.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_query.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts similarity index 94% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts index 8c996db2d63b45..65403feebcf127 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts @@ -5,14 +5,14 @@ * 2.0. */ -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; -import { DynamicTool } from 'langchain/tools'; +import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; import { getOpenAlertsTool } from './get_open_alerts_tool'; -import { mockAlertsFieldsApi } from '../../../../__mocks__/alerts'; -import type { RequestBody } from '../../types'; +import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; import { MAX_SIZE } from './helpers'; describe('getOpenAlertsTool', () => { diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts similarity index 86% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts index 755bfa7f9dc3a4..31837e8bc960ed 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/get_open_alerts_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts @@ -6,12 +6,13 @@ */ import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { KibanaRequest } from '@kbn/core-http-server'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { 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 type { Tool } from 'langchain/tools'; +import { DynamicTool } from 'langchain/tools'; +import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; +import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; import { getOpenAlertsQuery } from './get_open_alerts_query'; import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers'; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/helpers.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.test.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/helpers.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.test.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/helpers.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tools/open_alerts/helpers.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.ts diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 8ec2b2e743d1b9..e12fba23589d8d 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -111,7 +111,7 @@ import { registerRiskScoringTask } from './lib/entity_analytics/risk_score/tasks import { registerProtectionUpdatesNoteRoutes } from './endpoint/routes/protection_updates_note'; import { latestRiskScoreIndexPattern, allRiskScoreIndexPattern } from '../common/risk_engine'; import { isEndpointPackageV2 } from '../common/endpoint/utils/package_v2'; -import { getSecurityAssistantTools } from './assistant/get_security_assistant_tools'; +import { getApplicableTools } from './assistant/tools'; export type { SetupPlugins, StartPlugins, PluginSetup, PluginStart } from './plugin_contract'; @@ -508,7 +508,7 @@ export class Plugin implements ISecuritySolutionPlugin { this.licensing$ = plugins.licensing.license$; // Assistant Tool and Feature Registration - plugins.elasticAssistant.registerTool('securitySolution', getSecurityAssistantTools()); + plugins.elasticAssistant.registerTools(APP_ID, getApplicableTools); if (this.lists && plugins.taskManager && plugins.fleet) { // Exceptions, Artifacts and Manifests start From 68e1ceba82c4ecf151efa9f4e97c53f97b6679a3 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 12 Dec 2023 23:07:30 -0700 Subject: [PATCH 05/13] Adds tool metadata and separates params validation from tool creation --- .../plugins/elastic_assistant/server/index.ts | 4 +- .../execute_custom_llm_chain/index.ts | 8 +- .../server/lib/langchain/executors/types.ts | 4 +- .../elastic_assistant/server/plugin.ts | 6 +- .../server/routes/evaluate/post_evaluate.ts | 6 + .../routes/post_actions_connector_execute.ts | 6 +- .../server/services/app_context.ts | 32 +- .../plugins/elastic_assistant/server/types.ts | 55 +-- .../get_alert_counts_tool.test.ts | 131 ++++--- .../alert_counts/get_alert_counts_tool.ts | 57 ++- ..._esql_language_knowledge_base_tool.test.ts | 94 +++-- .../get_esql_language_knowledge_base_tool.ts | 47 +-- .../server/assistant/tools/index.test.ts | 42 +-- .../server/assistant/tools/index.ts | 50 +-- .../open_alerts/get_open_alerts_tool.test.ts | 353 +++++++++--------- .../tools/open_alerts/get_open_alerts_tool.ts | 131 +++---- .../security_solution/server/plugin.ts | 4 +- 17 files changed, 521 insertions(+), 509 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/index.ts b/x-pack/plugins/elastic_assistant/server/index.ts index bd55705d546eed..1775fc60528e82 100755 --- a/x-pack/plugins/elastic_assistant/server/index.ts +++ b/x-pack/plugins/elastic_assistant/server/index.ts @@ -17,6 +17,6 @@ export type { ElasticAssistantPluginStart, ElasticAssistantPluginSetupDependencies, ElasticAssistantPluginStartDependencies, - GetApplicableTools, - GetApplicableToolsParams, + 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 5baa65636dce0b..d6868925cc667b 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,6 +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 { AssistantToolParams } from '../../../types'; export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor'; @@ -30,6 +31,7 @@ export const callAgentExecutor = async ({ allow, allowReplacement, assistantLangChain, + assistantTools = [], connectorId, elserId, esClient, @@ -38,7 +40,6 @@ export const callAgentExecutor = async ({ llmType, logger, onNewReplacements, - getApplicableTools = (params: GetRegisteredToolParams) => [], replacements, request, size, @@ -72,7 +73,7 @@ export const callAgentExecutor = async ({ const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever(10)); // Fetch any applicable tools that the source plugin may have registered - const tools: Tool[] = getApplicableTools({ + const assistantToolParams: AssistantToolParams = { allow, allowReplacement, alertsIndexPattern, @@ -84,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 bd19352d9c346a..e7824e2822f8a3 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,7 +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 { GetApplicableTools } from '../../../types'; +import type { AssistantTool } from '../../../types'; export interface AgentExecutorParams { alertsIndexPattern?: string; @@ -20,9 +20,9 @@ export interface AgentExecutorParams { allow?: string[]; allowReplacement?: string[]; assistantLangChain: boolean; + assistantTools?: AssistantTool[]; connectorId: string; esClient: ElasticsearchClient; - getApplicableTools?: GetApplicableTools; kbResource: string | undefined; langChainMessages: BaseMessage[]; llmType?: string; diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index b9574e0cd12d32..2919e7b7f80128 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -18,12 +18,12 @@ import { import { once } from 'lodash'; import { + AssistantTool, ElasticAssistantPluginSetup, ElasticAssistantPluginSetupDependencies, ElasticAssistantPluginStart, ElasticAssistantPluginStartDependencies, ElasticAssistantRequestHandlerContext, - GetApplicableTools, GetElser, PLUGIN_ID, } from './types'; @@ -122,8 +122,8 @@ export class ElasticAssistantPlugin getRegisteredTools: (pluginName: string) => { return appContextService.getRegisteredTools(pluginName); }, - registerTools: (pluginName: string, getApplicableTools: GetApplicableTools) => { - return appContextService.registerTools(pluginName, getApplicableTools); + registerTools: (pluginName: string, tools: AssistantTool[]) => { + return appContextService.registerTools(pluginName, tools); }, }; } 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 39e3233f637444..f5ec4c4555a766 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/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index e4cf642ba15ebf..aede10c8aa263d 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 @@ -58,8 +58,8 @@ export const postActionsConnectorExecuteRoute = ( // TODO: Add `traceId` to actions request when calling via langchain logger.debug('Executing via langchain, assistantLangChain: true'); - // Fetch the function for fetching any tools registered by the request's originating plugin - const getApplicableTools = (await context.elasticAssistant).getRegisteredTools( + // Fetch any tools registered by the request's originating plugin + const assistantTools = (await context.elasticAssistant).getRegisteredTools( 'securitySolution' ); @@ -84,10 +84,10 @@ export const postActionsConnectorExecuteRoute = ( allowReplacement: request.body.allowReplacement, actions, assistantLangChain: request.body.assistantLangChain, + assistantTools, connectorId, elserId, esClient, - getApplicableTools, kbResource: ESQL_RESOURCE, langChainMessages, logger, diff --git a/x-pack/plugins/elastic_assistant/server/services/app_context.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.ts index 0f10aa793d735f..bd7a7c0cc3203b 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -6,14 +6,11 @@ */ import type { Logger } from '@kbn/core/server'; -import type { Tool } from 'langchain/tools'; -import type { GetApplicableTools, GetApplicableToolsParams } from '../types'; +import type { AssistantTool } from '../types'; export type PluginName = string; -export type RegisteredToolsStorage = Map; -export type RegisterTools = (pluginName: string, getApplicableTools: GetApplicableTools) => void; -export type GetRegisteredTools = (pluginName: string) => Tool[]; - +export type RegisteredToolsStorage = Map>; +export type GetRegisteredTools = (pluginName: string) => AssistantTool[]; export interface ElasticAssistantAppContext { logger: Logger; } @@ -25,7 +22,7 @@ export interface ElasticAssistantAppContext { */ class AppContextService { private logger: Logger | undefined; - private registeredTools: RegisteredToolsStorage = new Map(); + private registeredTools: RegisteredToolsStorage = new Map>(); public start(appContext: ElasticAssistantAppContext) { this.logger = appContext.logger; @@ -39,32 +36,33 @@ class AppContextService { * Register tools to be used by the Elastic Assistant * * @param pluginName - * @param getApplicableTools + * @param tools */ - public registerTools(pluginName: string, getApplicableTools: GetApplicableTools) { + 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, setting "getApplicableTools"'); - } else { - this.logger?.debug('plugin already has tools, overriding "getApplicableTools"'); + this.logger?.debug('plugin has no tools, making new set'); + this.registeredTools.set(pluginName, new Set()); } - this.registeredTools.set(pluginName, getApplicableTools); + tools.forEach((tool) => this.registeredTools.get(pluginName)?.add(tool)); } /** * Get the registered tools + * * @param pluginName */ - public getRegisteredTools(pluginName: string): GetApplicableTools { - const getApplicableTools = - this.registeredTools?.get(pluginName) ?? ((params: GetApplicableToolsParams) => []); + 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 getApplicableTools; + return tools; } } diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 6cb58ecf5ad466..5be8a35a275b97 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -20,25 +20,10 @@ 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; -export interface GetApplicableToolsParams { - 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 type GetApplicableTools = (params: GetApplicableToolsParams) => Tool[]; - /** The plugin setup interface */ export interface ElasticAssistantPluginSetup { actions: ActionsPluginSetup; @@ -50,14 +35,14 @@ export interface ElasticAssistantPluginStart { /** * Register tools to be used by the elastic assistant * @param pluginName Name of the plugin the tool should be registered to - * @param getApplicableTools Function that returns the tools for the specified plugin + * @param tools AssistantTools to be registered with for the given plugin */ - registerTools: (pluginName: string, getApplicableTools: GetApplicableTools) => void; + registerTools: (pluginName: string, tools: AssistantTool[]) => void; /** * Get the registered tools * @param pluginName Name of the plugin to get the tools for */ - getRegisteredTools: (pluginName: string) => Tool[]; + getRegisteredTools: GetRegisteredTools; } export interface ElasticAssistantPluginSetupDependencies { @@ -70,10 +55,7 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { actions: ActionsPluginStart; - getRegisteredTools: ( - pluginName: string, - getApplicableToolsParams: GetApplicableToolsParams - ) => Tool[]; + getRegisteredTools: GetRegisteredTools; logger: Logger; } @@ -88,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/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts index 343d4df4a79374..37239362480956 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts @@ -10,10 +10,11 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; -import { getAlertCountsTool } from './get_alert_counts_tool'; import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; +import { ALERT_COUNTS_TOOL } from './get_alert_counts_tool'; +import type { RetrievalQAChain } from 'langchain/chains'; -describe('getAlertCountsTool', () => { +describe('AlertCountsTool', () => { const alertsIndexPattern = 'alerts-index'; const esClient = { search: jest.fn().mockResolvedValue({}), @@ -29,77 +30,93 @@ describe('getAlertCountsTool', () => { 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(); }); - 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' } }], + 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' } } }, - ], + { range: { '@timestamp': { gte: 'now/d', lte: 'now/d' } } }, + ], + }, }, - }, - size: 0, + 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 - >; + 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, + const tool = ALERT_COUNTS_TOOL.getTool({ + alertsIndexPattern, + esClient, + replacements, + request: requestWithMissingParams, + ...rest, + }); + + expect(tool).toBeNull(); }); - expect(tool).toBeNull(); - }); + it('returns null when the alertsIndexPattern is undefined', () => { + const tool = ALERT_COUNTS_TOOL.getTool({ + // alertsIndexPattern is undefined + esClient, + replacements, + request, - it('returns null when the alertsIndexPattern is undefined', () => { - const tool = getAlertCountsTool({ - // alertsIndexPattern is undefined - esClient, - replacements, - request, + ...rest, + }); + + expect(tool).toBeNull(); }); - expect(tool).toBeNull(); - }); + it('returns a tool instance with the expected tags', () => { + const tool = ALERT_COUNTS_TOOL.getTool({ + alertsIndexPattern, + esClient, + replacements, + request, - it('returns a tool instance with the expected tags', () => { - const tool = getAlertCountsTool({ - alertsIndexPattern, - esClient, - replacements, - request, - }) as DynamicTool; + ...rest, + }) as DynamicTool; - expect(tool.tags).toEqual(['alerts', 'alerts-count']); + expect(tool.tags).toEqual(['alerts', 'alerts-count']); + }); }); }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts index 62b13e10fab76c..441795dc472ec7 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts @@ -6,43 +6,42 @@ */ import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; -import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import type { KibanaRequest } from '@kbn/core-http-server'; -import type { Tool } from 'langchain/tools'; import { DynamicTool } from 'langchain/tools'; import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; -import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import { getAlertsCountQuery } from './get_alert_counts_query'; +import { APP_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 getAlertCountsTool = ({ - alertsIndexPattern, - esClient, - replacements, - request, -}: { - alertsIndexPattern?: string; - esClient: ElasticsearchClient; - replacements?: Record; - request: KibanaRequest; -}): Tool | null => { - if (!requestHasRequiredAnonymizationParams(request) || alertsIndexPattern == null) { - return null; - } +export const ALERT_COUNTS_TOOL: AssistantTool = { + id: 'alert-counts-tool', + name: 'AlertCountsTool', + description: ALERT_COUNTS_TOOL_DESCRIPTION, + sourceRegister: APP_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: 'alert-counts', + description: ALERT_COUNTS_TOOL_DESCRIPTION, + func: async () => { + const query = getAlertsCountQuery(alertsIndexPattern); - return new DynamicTool({ - name: 'alert-counts', - description: ALERT_COUNTS_TOOL_DESCRIPTION, - func: async () => { - const query = getAlertsCountQuery(alertsIndexPattern); + const result = await esClient.search(query); - const result = await esClient.search(query); - - return JSON.stringify(result); - }, - tags: ['alerts', 'alerts-count'], - }); + return JSON.stringify(result); + }, + tags: ['alerts', 'alerts-count'], + }); + }, }; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts index ccd97b7deb0883..c45a57ddf75a8d 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts @@ -5,51 +5,77 @@ * 2.0. */ -import { RetrievalQAChain } from 'langchain/chains'; -import { DynamicTool } from 'langchain/tools'; - -import { getEsqlLanguageKnowledgeBaseTool } from './get_esql_language_knowledge_base_tool'; +import type { RetrievalQAChain } from 'langchain/chains'; +import type { DynamicTool } from 'langchain/tools'; +import { ESQL_KNOWLEDGE_BASE_TOOL } from './get_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'; const chain = {} as RetrievalQAChain; -describe('getEsqlLanguageKnowledgeBaseTool', () => { - it('returns null if assistantLangChain is false', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: false, - chain, - modelExists: true, - }); +describe('EsqlLanguageKnowledgeBaseTool', () => { + describe('getTool', () => { + 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 = { + esClient, + request, + }; - expect(tool).toBeNull(); - }); + it('returns null if assistantLangChain is false', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: false, + chain, + modelExists: true, + ...rest, + }); - 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(); }); - 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, + chain, + modelExists: false, // <-- ELSER model is not installed + ...rest, + }); - it('should return a Tool instance if assistantLangChain and modelExists are true', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: true, - modelExists: true, - chain, + expect(tool).toBeNull(); }); - expect(tool?.name).toEqual('ESQLKnowledgeBaseTool'); - }); + it('should return a Tool instance if assistantLangChain and modelExists are true', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: true, + chain, + modelExists: true, + ...rest, + }); + + expect(tool?.name).toEqual('ESQLKnowledgeBaseTool'); + }); - it('should return a tool with the expected tags', () => { - const tool = getEsqlLanguageKnowledgeBaseTool({ - assistantLangChain: true, - chain, - modelExists: true, - }) as DynamicTool; + it('should return a tool with the expected tags', () => { + const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ + assistantLangChain: true, + chain, + modelExists: true, + ...rest, + }) as DynamicTool; - expect(tool.tags).toEqual(['esql', 'query-generation', 'knowledge-base']); + expect(tool.tags).toEqual(['esql', 'query-generation', 'knowledge-base']); + }); }); }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts index eaf73dc0785603..70891b661d2aff 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts @@ -5,26 +5,31 @@ * 2.0. */ -import type { RetrievalQAChain } from 'langchain/chains'; -import type { Tool } from 'langchain/tools'; import { ChainTool } from 'langchain/tools'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { APP_ID } from '../../../../common'; -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; +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_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 index daaa5ec3edcd3c..047c84ceddf3b2 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.test.ts @@ -5,51 +5,15 @@ * 2.0. */ -import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import type { KibanaRequest } from '@kbn/core-http-server'; -import type { RetrievalQAChain } from 'langchain/chains'; - -import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/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; +import { getAssistantTools } from '.'; +describe('getAssistantTools', () => { 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 tools = getAssistantTools(); const minExpectedTools = 3; // 3 tools are currently implemented diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index b532a8183fa66d..44b7a5642328c1 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -5,45 +5,13 @@ * 2.0. */ -import type { Tool } from 'langchain/tools'; +import type { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; +import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/get_esql_language_knowledge_base_tool'; +import { OPEN_ALERTS_TOOL } from './open_alerts/get_open_alerts_tool'; +import { ALERT_COUNTS_TOOL } from './alert_counts/get_alert_counts_tool'; -import type { - GetApplicableTools, - GetApplicableToolsParams, -} from '@kbn/elastic-assistant-plugin/server'; -import { getAlertCountsTool } from './alert_counts/get_alert_counts_tool'; -import { getEsqlLanguageKnowledgeBaseTool } from './esql_language_knowledge_base/get_esql_language_knowledge_base_tool'; -import { getOpenAlertsTool } from './open_alerts/get_open_alerts_tool'; - -export const getApplicableTools: GetApplicableTools = ({ - alertsIndexPattern, - allow, - allowReplacement, - assistantLangChain, - chain, - esClient, - modelExists, - onNewReplacements, - replacements, - request, - size, -}: GetApplicableToolsParams): Tool[] => - [ - getEsqlLanguageKnowledgeBaseTool({ assistantLangChain, chain, modelExists }) ?? [], - getAlertCountsTool({ - alertsIndexPattern, - esClient, - replacements, - request, - }) ?? [], - getOpenAlertsTool({ - alertsIndexPattern, - allow, - allowReplacement, - esClient, - onNewReplacements, - replacements, - request, - size, - }) ?? [], - ].flat(); +export const getAssistantTools = (): AssistantTool[] => [ + ALERT_COUNTS_TOOL, + ESQL_KNOWLEDGE_BASE_TOOL, + OPEN_ALERTS_TOOL, +]; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts index 65403feebcf127..9a17b0d12e9498 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts @@ -10,197 +10,214 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; -import { getOpenAlertsTool } from './get_open_alerts_tool'; +import { OPEN_ALERTS_TOOL } from './get_open_alerts_tool'; import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; import { MAX_SIZE } from './helpers'; +import type { RetrievalQAChain } from 'langchain/chains'; describe('getOpenAlertsTool', () => { - 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 = getOpenAlertsTool({ - 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, + describe('getTool', () => { + const alertsIndexPattern = 'alerts-index'; + const esClient = { + search: jest.fn().mockResolvedValue(mockAlertsFieldsApi), + } as unknown as ElasticsearchClient; + const replacements = { key: 'value' }; + const request = { 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: [ - { - match_phrase: { - 'kibana.alert.workflow_status': 'open', + 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(); + }); + + it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => { + const tool: DynamicTool = OPEN_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + esClient, + 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: [ + { + match_phrase: { + 'kibana.alert.workflow_status': 'open', + }, }, - }, - { - range: { - '@timestamp': { - format: 'strict_date_optional_time', - gte: 'now-1d/d', - lte: 'now/d', + { + range: { + '@timestamp': { + format: 'strict_date_optional_time', + gte: 'now-1d/d', + lte: 'now/d', + }, }, }, - }, - ], - must: [], - must_not: [ - { - exists: { - field: 'kibana.alert.building_block_type', + ], + must: [], + must_not: [ + { + exists: { + field: 'kibana.alert.building_block_type', + }, }, - }, - ], - should: [], + ], + should: [], + }, }, - }, - ], - }, - }, - runtime_mappings: {}, - size: 20, - sort: [ - { - 'kibana.alert.risk_score': { - order: 'desc', + ], }, }, - { - '@timestamp': { - order: 'desc', + runtime_mappings: {}, + size: 20, + sort: [ + { + 'kibana.alert.risk_score': { + order: 'desc', + }, }, - }, - ], - }, - ignore_unavailable: true, - index: ['alerts-index'], + { + '@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 = getOpenAlertsTool({ - alertsIndexPattern, - allow: requestWithMissingParams.body.allow, - allowReplacement: requestWithMissingParams.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request: requestWithMissingParams, - size: requestWithMissingParams.body.size, + 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_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: requestWithMissingParams.body.allow, + allowReplacement: requestWithMissingParams.body.allowReplacement, + esClient, + onNewReplacements: jest.fn(), + replacements, + request: requestWithMissingParams, + size: requestWithMissingParams.body.size, + ...rest, + }); + + expect(tool).toBeNull(); }); - expect(tool).toBeNull(); - }); - - it('returns null when alertsIndexPattern is undefined', () => { - const tool = getOpenAlertsTool({ - // alertsIndexPattern is undefined - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - size: request.body.size, + it('returns null when alertsIndexPattern is undefined', () => { + const tool = OPEN_ALERTS_TOOL.getTool({ + // alertsIndexPattern is undefined + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + esClient, + onNewReplacements: jest.fn(), + replacements, + request, + size: request.body.size, + ...rest, + }); + + expect(tool).toBeNull(); }); - expect(tool).toBeNull(); - }); - - it('returns null when size is undefined', () => { - const tool = getOpenAlertsTool({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - // size is undefined + it('returns null when size is undefined', () => { + const tool = OPEN_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + esClient, + onNewReplacements: jest.fn(), + replacements, + request, + ...rest, + // size is undefined + }); + + expect(tool).toBeNull(); }); - expect(tool).toBeNull(); - }); - - it('returns null when size out of range', () => { - const tool = getOpenAlertsTool({ - alertsIndexPattern, - allow: request.body.allow, - allowReplacement: request.body.allowReplacement, - esClient, - onNewReplacements: jest.fn(), - replacements, - request, - size: MAX_SIZE + 1, // <-- size is out of range + it('returns null when size out of range', () => { + const tool = OPEN_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + esClient, + onNewReplacements: jest.fn(), + replacements, + request, + size: MAX_SIZE + 1, // <-- size is out of range + ...rest, + }); + + expect(tool).toBeNull(); }); - expect(tool).toBeNull(); - }); - - it('returns a tool instance with the expected tags', () => { - const tool = getOpenAlertsTool({ - 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']); + it('returns a tool instance with the expected tags', () => { + const tool = OPEN_ALERTS_TOOL.getTool({ + alertsIndexPattern, + allow: request.body.allow, + allowReplacement: request.body.allowReplacement, + esClient, + onNewReplacements: jest.fn(), + replacements, + request, + size: request.body.size, + ...rest, + }) as DynamicTool; + + expect(tool.tags).toEqual(['alerts', 'open-alerts']); + }); }); }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts index 31837e8bc960ed..7a9ec4a11ad120 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts @@ -6,86 +6,87 @@ */ import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; -import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import type { KibanaRequest } from '@kbn/core-http-server'; import { getAnonymizedValue, transformRawData } from '@kbn/elastic-assistant-common'; -import type { Tool } from 'langchain/tools'; import { DynamicTool } from 'langchain/tools'; import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; -import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import { getOpenAlertsQuery } from './get_open_alerts_query'; import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers'; +import { APP_ID } from '../../../../common'; + +export interface OpenAlertsToolParams extends AssistantToolParams { + alertsIndexPattern: string; + size: number; +} 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 alerts, or null if the request - * doesn't have all the required parameters. + * A tool for querying open alerts, or null if the request doesn't have all the required parameters. */ -export const getOpenAlertsTool = ({ - 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; - } +export const OPEN_ALERTS_TOOL: AssistantTool = { + id: 'open-alerts-tool', + name: 'OpenAlertsTool', + description: OPEN_ALERTS_TOOL_DESCRIPTION, + sourceRegister: APP_ID, + isSupported: (params: AssistantToolParams): params is OpenAlertsToolParams => { + const { alertsIndexPattern, request, size } = params; + return ( + requestHasRequiredAnonymizationParams(request) && + alertsIndexPattern != null && + size != null && + !sizeIsOutOfRange(size) + ); + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; - return new DynamicTool({ - name: 'open-alerts', - description: OPEN_ALERTS_TOOL_DESCRIPTION, - func: async () => { - const query = getOpenAlertsQuery({ - alertsIndexPattern, - allow: allow ?? [], - size, - }); + const { + alertsIndexPattern, + allow, + allowReplacement, + esClient, + onNewReplacements, + replacements, + size, + } = params as OpenAlertsToolParams; + return new DynamicTool({ + name: 'open-alerts', + description: OPEN_ALERTS_TOOL_DESCRIPTION, + func: async () => { + const query = getOpenAlertsQuery({ + alertsIndexPattern, + allow: allow ?? [], + size, + }); - const result = await esClient.search(query); + 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 + // 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 - }; + 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'], - }); + 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/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index e12fba23589d8d..3fcb7f2eca6350 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -111,7 +111,7 @@ import { registerRiskScoringTask } from './lib/entity_analytics/risk_score/tasks import { registerProtectionUpdatesNoteRoutes } from './endpoint/routes/protection_updates_note'; import { latestRiskScoreIndexPattern, allRiskScoreIndexPattern } from '../common/risk_engine'; import { isEndpointPackageV2 } from '../common/endpoint/utils/package_v2'; -import { getApplicableTools } from './assistant/tools'; +import { getAssistantTools } from './assistant/tools'; export type { SetupPlugins, StartPlugins, PluginSetup, PluginStart } from './plugin_contract'; @@ -508,7 +508,7 @@ export class Plugin implements ISecuritySolutionPlugin { this.licensing$ = plugins.licensing.license$; // Assistant Tool and Feature Registration - plugins.elasticAssistant.registerTools(APP_ID, getApplicableTools); + plugins.elasticAssistant.registerTools(APP_ID, getAssistantTools()); if (this.lists && plugins.taskManager && plugins.fleet) { // Exceptions, Artifacts and Manifests start From d4e52f5caa40b10c5a4ba357afe375dca211b9e0 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 13 Dec 2023 06:17:19 +0000 Subject: [PATCH 06/13] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/elastic_assistant/tsconfig.json | 1 - x-pack/plugins/security_solution/tsconfig.json | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 5cad4d4b52141c..53616fc2dc2b07 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/tsconfig.json b/x-pack/plugins/security_solution/tsconfig.json index aa2417daeabb09..2158fe97996dc2 100644 --- a/x-pack/plugins/security_solution/tsconfig.json +++ b/x-pack/plugins/security_solution/tsconfig.json @@ -184,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" ] } From 85657840e47af87533cb90899956ab9875dadb84 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 13 Dec 2023 08:48:27 -0700 Subject: [PATCH 07/13] Adds jest config for new security solution assistant server tests --- .../server/assistant/jest.config.js | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 x-pack/plugins/security_solution/server/assistant/jest.config.js 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 00000000000000..787dae7ce8d687 --- /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'), +}; From 0f4789872f60f03e814ad8a8c4d5717367e3c3a0 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 13 Dec 2023 12:01:14 -0700 Subject: [PATCH 08/13] Renames to match new tool names from merge w/ main --- ..._alert_counts_tool.test.ts => alert_counts_tool.test.ts} | 2 +- .../{get_alert_counts_tool.ts => alert_counts_tool.ts} | 0 ...ol.test.ts => esql_language_knowledge_base_tool.test.ts} | 2 +- ...ge_base_tool.ts => esql_language_knowledge_base_tool.ts} | 0 .../security_solution/server/assistant/tools/index.ts | 6 +++--- .../get_open_and_acknowledge_alerts_query.test.ts} | 2 +- .../get_open_and_acknowledged_alerts_query.ts} | 0 .../helpers.test.ts | 0 .../{open_alerts => open_and_acknowledge_alerts}/helpers.ts | 0 .../open_and_acknowledge_alerts_tool.ts} | 2 +- .../open_and_acknowledged_alerts_tool.test.ts} | 2 +- 11 files changed, 8 insertions(+), 8 deletions(-) rename x-pack/plugins/security_solution/server/assistant/tools/alert_counts/{get_alert_counts_tool.test.ts => alert_counts_tool.test.ts} (98%) rename x-pack/plugins/security_solution/server/assistant/tools/alert_counts/{get_alert_counts_tool.ts => alert_counts_tool.ts} (100%) rename x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/{get_esql_language_knowledge_base_tool.test.ts => esql_language_knowledge_base_tool.test.ts} (96%) rename x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/{get_esql_language_knowledge_base_tool.ts => esql_language_knowledge_base_tool.ts} (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_alerts/get_open_alerts_query.test.ts => open_and_acknowledge_alerts/get_open_and_acknowledge_alerts_query.test.ts} (96%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_alerts/get_open_alerts_query.ts => open_and_acknowledge_alerts/get_open_and_acknowledged_alerts_query.ts} (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_alerts => open_and_acknowledge_alerts}/helpers.test.ts (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_alerts => open_and_acknowledge_alerts}/helpers.ts (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_alerts/get_open_alerts_tool.ts => open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts} (97%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_alerts/get_open_alerts_tool.test.ts => open_and_acknowledge_alerts/open_and_acknowledged_alerts_tool.test.ts} (99%) diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts similarity index 98% rename from x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts index 37239362480956..267f9d5a62dbba 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts @@ -11,7 +11,7 @@ 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 './get_alert_counts_tool'; +import { ALERT_COUNTS_TOOL } from './alert_counts_tool'; import type { RetrievalQAChain } from 'langchain/chains'; describe('AlertCountsTool', () => { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts similarity index 100% rename from x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_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 similarity index 96% rename from x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.test.ts index c45a57ddf75a8d..51a7e7a518599c 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_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 @@ -7,7 +7,7 @@ import type { RetrievalQAChain } from 'langchain/chains'; import type { DynamicTool } from 'langchain/tools'; -import { ESQL_KNOWLEDGE_BASE_TOOL } from './get_esql_language_knowledge_base_tool'; +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'; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_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 similarity index 100% rename from x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/get_esql_language_knowledge_base_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index 44b7a5642328c1..a07f7b9bafaf86 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -6,9 +6,9 @@ */ import type { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; -import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/get_esql_language_knowledge_base_tool'; -import { OPEN_ALERTS_TOOL } from './open_alerts/get_open_alerts_tool'; -import { ALERT_COUNTS_TOOL } from './alert_counts/get_alert_counts_tool'; +import { ESQL_KNOWLEDGE_BASE_TOOL } from './esql_language_knowledge_base/esql_language_knowledge_base_tool'; +import { OPEN_ALERTS_TOOL } from './open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool'; +import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; export const getAssistantTools = (): AssistantTool[] => [ ALERT_COUNTS_TOOL, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledge_alerts_query.test.ts similarity index 96% rename from x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledge_alerts_query.test.ts index f09ad2682d0645..a0cf067099e920 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledge_alerts_query.test.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { getOpenAndAcknowledgedAlertsQuery } from './get_open_alerts_query'; +import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query'; describe('getOpenAndAcknowledgedAlertsQuery', () => { it('returns the expected query', () => { diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledged_alerts_query.ts similarity index 100% rename from x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_query.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledged_alerts_query.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/helpers.test.ts similarity index 100% rename from x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/helpers.test.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/helpers.ts similarity index 100% rename from x-pack/plugins/security_solution/server/assistant/tools/open_alerts/helpers.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/helpers.ts diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts similarity index 97% rename from x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts index 474cdc34564701..6ac65c25d050f3 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts @@ -11,7 +11,7 @@ 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_alerts_query'; +import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_alerts_query'; import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers'; import { APP_ID } from '../../../../common'; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledged_alerts_tool.test.ts similarity index 99% rename from x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledged_alerts_tool.test.ts index fa6fbf84510aaf..69a111ba9e944f 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_alerts/get_open_alerts_tool.test.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledged_alerts_tool.test.ts @@ -10,7 +10,7 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; -import { OPEN_ALERTS_TOOL } from './get_open_alerts_tool'; +import { OPEN_ALERTS_TOOL } from './open_and_acknowledge_alerts_tool'; import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; import { MAX_SIZE } from './helpers'; From 7ab2e42c018396b6a9e5d79b9cba645b7218e6f2 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 13 Dec 2023 12:34:13 -0700 Subject: [PATCH 09/13] Additional missed renames --- .../server/assistant/tools/index.ts | 4 ++-- ...pen_and_acknowledged_alerts_query.test.ts} | 0 .../get_open_and_acknowledged_alerts_query.ts | 0 .../helpers.test.ts | 0 .../helpers.ts | 0 .../open_and_acknowledged_alerts_tool.test.ts | 18 ++++++++--------- .../open_and_acknowledged_alerts_tool.ts} | 20 +++++++++---------- 7 files changed, 21 insertions(+), 21 deletions(-) rename x-pack/plugins/security_solution/server/assistant/tools/{open_and_acknowledge_alerts/get_open_and_acknowledge_alerts_query.test.ts => open_and_acknowledged_alerts/get_open_and_acknowledged_alerts_query.test.ts} (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_and_acknowledge_alerts => open_and_acknowledged_alerts}/get_open_and_acknowledged_alerts_query.ts (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_and_acknowledge_alerts => open_and_acknowledged_alerts}/helpers.test.ts (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_and_acknowledge_alerts => open_and_acknowledged_alerts}/helpers.ts (100%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_and_acknowledge_alerts => open_and_acknowledged_alerts}/open_and_acknowledged_alerts_tool.test.ts (91%) rename x-pack/plugins/security_solution/server/assistant/tools/{open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts => open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts} (82%) diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index a07f7b9bafaf86..790e674a4b390c 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -7,11 +7,11 @@ 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_ALERTS_TOOL } from './open_and_acknowledge_alerts/open_and_acknowledge_alerts_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_ALERTS_TOOL, + OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL, ]; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_alerts/get_open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_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/security_solution/server/assistant/tools/open_and_acknowledge_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_acknowledge_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 similarity index 91% rename from x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledged_alerts_tool.test.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts index 69a111ba9e944f..c7d7e601bd4d19 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_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 @@ -10,13 +10,13 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; -import { OPEN_ALERTS_TOOL } from './open_and_acknowledge_alerts_tool'; +import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts_tool'; import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; import type { RequestBody } from '@kbn/elastic-assistant-plugin/server/lib/langchain/types'; import { MAX_SIZE } from './helpers'; import type { RetrievalQAChain } from 'langchain/chains'; -describe('getOpenAlertsTool', () => { +describe('OpenAndAcknowledgedAlertsTool', () => { describe('getTool', () => { const alertsIndexPattern = 'alerts-index'; const esClient = { @@ -47,7 +47,7 @@ describe('getOpenAlertsTool', () => { }); it('returns a `DynamicTool` with a `func` that calls `esClient.search()` with the expected query', async () => { - const tool: DynamicTool = OPEN_ALERTS_TOOL.getTool({ + const tool: DynamicTool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, @@ -153,7 +153,7 @@ describe('getOpenAlertsTool', () => { RequestBody >; - const tool = OPEN_ALERTS_TOOL.getTool({ + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ alertsIndexPattern, allow: requestWithMissingParams.body.allow, allowReplacement: requestWithMissingParams.body.allowReplacement, @@ -169,7 +169,7 @@ describe('getOpenAlertsTool', () => { }); it('returns null when alertsIndexPattern is undefined', () => { - const tool = OPEN_ALERTS_TOOL.getTool({ + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ // alertsIndexPattern is undefined allow: request.body.allow, allowReplacement: request.body.allowReplacement, @@ -185,7 +185,7 @@ describe('getOpenAlertsTool', () => { }); it('returns null when size is undefined', () => { - const tool = OPEN_ALERTS_TOOL.getTool({ + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, @@ -201,7 +201,7 @@ describe('getOpenAlertsTool', () => { }); it('returns null when size out of range', () => { - const tool = OPEN_ALERTS_TOOL.getTool({ + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, @@ -217,7 +217,7 @@ describe('getOpenAlertsTool', () => { }); it('returns a tool instance with the expected tags', () => { - const tool = OPEN_ALERTS_TOOL.getTool({ + const tool = OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL.getTool({ alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, @@ -229,7 +229,7 @@ describe('getOpenAlertsTool', () => { ...rest, }) as DynamicTool; - expect(tool.tags).toEqual(['alerts', 'open-alerts']); + expect(tool.tags).toEqual(['alerts', 'open-and-acknowledged-alerts']); }); }); }); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts similarity index 82% rename from x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts rename to x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts index 6ac65c25d050f3..b4f6532ad6755d 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledge_alerts/open_and_acknowledge_alerts_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts @@ -15,24 +15,24 @@ import { getOpenAndAcknowledgedAlertsQuery } from './get_open_and_acknowledged_a import { getRawDataOrDefault, sizeIsOutOfRange } from './helpers'; import { APP_ID } from '../../../../common'; -export interface OpenAlertsToolParams extends AssistantToolParams { +export interface OpenAndAcknowledgedAlertsToolParams extends AssistantToolParams { alertsIndexPattern: string; size: number; } -export const OPEN_ALERTS_TOOL_DESCRIPTION = +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_ALERTS_TOOL: AssistantTool = { - id: 'open-alerts-tool', - name: 'OpenAlertsTool', - description: OPEN_ALERTS_TOOL_DESCRIPTION, +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_ID, - isSupported: (params: AssistantToolParams): params is OpenAlertsToolParams => { + isSupported: (params: AssistantToolParams): params is OpenAndAcknowledgedAlertsToolParams => { const { alertsIndexPattern, request, size } = params; return ( requestHasRequiredAnonymizationParams(request) && @@ -52,10 +52,10 @@ export const OPEN_ALERTS_TOOL: AssistantTool = { onNewReplacements, replacements, size, - } = params as OpenAlertsToolParams; + } = params as OpenAndAcknowledgedAlertsToolParams; return new DynamicTool({ name: 'open-alerts', - description: OPEN_ALERTS_TOOL_DESCRIPTION, + description: OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION, func: async () => { const query = getOpenAndAcknowledgedAlertsQuery({ alertsIndexPattern, @@ -87,7 +87,7 @@ export const OPEN_ALERTS_TOOL: AssistantTool = { ) ); }, - tags: ['alerts', 'open-alerts'], + tags: ['alerts', 'open-and-acknowledged-alerts'], }); }, }; From 2f65edb7a4599ce0abaa350484464d10185794dc Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 14 Dec 2023 00:07:54 -0700 Subject: [PATCH 10/13] Extracts plugin name from request and adds tests --- .../elastic_assistant/server/plugin.ts | 17 +++ .../server/routes/helpers.test.ts | 58 ++++++++ .../server/routes/helpers.ts | 48 +++++++ .../routes/post_actions_connector_execute.ts | 6 +- .../server/services/app_context.test.ts | 87 ++++++++++++ .../alert_counts/alert_counts_tool.test.ts | 40 ++++++ .../tools/alert_counts/alert_counts_tool.ts | 4 +- .../esql_language_knowledge_base_tool.test.ts | 74 ++++++---- .../esql_language_knowledge_base_tool.ts | 4 +- .../open_and_acknowledged_alerts_tool.test.ts | 129 +++++++++++++----- .../open_and_acknowledged_alerts_tool.ts | 4 +- .../security_solution/server/plugin.ts | 4 +- 12 files changed, 408 insertions(+), 67 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/helpers.ts create mode 100644 x-pack/plugins/elastic_assistant/server/services/app_context.test.ts diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 2919e7b7f80128..bd06165e57284c 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -118,10 +118,27 @@ export class ElasticAssistantPlugin 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); }, 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 00000000000000..c35a045c9666d8 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts @@ -0,0 +1,58 @@ +/* + * 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 { 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 }); + expect(pluginName).toEqual('superSolution'); + }); + + it('fails to extracts plugin name from undefined "x-kbn-context" request header', async () => { + const invalidRequest = { + headers: { + 'x-kbn-context': undefined, + }, + } as unknown as KibanaRequest; + const pluginName = getPluginNameFromRequest({ request: invalidRequest }); + expect(pluginName).toEqual('securitySolutionUI'); + }); + + it('fails to extracts plugin name from malformed "x-kbn-context" invalidRequest header', async () => { + const invalidRequest = { + headers: { + 'x-kbn-context': 'asdfku', + }, + } as unknown as KibanaRequest; + const pluginName = getPluginNameFromRequest({ request: invalidRequest }); + expect(pluginName).toEqual('securitySolutionUI'); + }); +}); 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 00000000000000..b4291cea8382e1 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -0,0 +1,48 @@ +/* + * 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; + 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 request Kibana Request + * + * @returns plugin name + */ +export const getPluginNameFromRequest = ({ + logger, + 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, using default assistant tools.'); + } + return DEFAULT_PLUGIN_NAME; +}; 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 aede10c8aa263d..3ba021d33a02e5 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 { getPluginNameFromRequest } from './helpers'; export const postActionsConnectorExecuteRoute = ( router: IRouter, @@ -59,9 +60,8 @@ export const postActionsConnectorExecuteRoute = ( logger.debug('Executing via langchain, assistantLangChain: true'); // Fetch any tools registered by the request's originating plugin - const assistantTools = (await context.elasticAssistant).getRegisteredTools( - 'securitySolution' - ); + const pluginName = getPluginNameFromRequest({ request, logger }); + const assistantTools = (await context.elasticAssistant).getRegisteredTools(pluginName); // get a scoped esClient for assistant memory const esClient = (await context.core).elasticsearch.client.asCurrentUser; 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 00000000000000..621995d3452be5 --- /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/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 index 267f9d5a62dbba..715d656944c86a 100644 --- 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 @@ -43,6 +43,46 @@ describe('AlertCountsTool', () => { 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({ 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 index 441795dc472ec7..7cdd2342cd1f3a 100644 --- 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 @@ -11,7 +11,7 @@ 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_ID } from '../../../../common'; +import { APP_UI_ID } from '../../../../common'; export interface AlertCountsToolParams extends AssistantToolParams { alertsIndexPattern: string; @@ -23,7 +23,7 @@ export const ALERT_COUNTS_TOOL: AssistantTool = { id: 'alert-counts-tool', name: 'AlertCountsTool', description: ALERT_COUNTS_TOOL_DESCRIPTION, - sourceRegister: APP_ID, + sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is AlertCountsToolParams => { const { request, alertsIndexPattern } = params; return requestHasRequiredAnonymizationParams(request) && alertsIndexPattern != null; 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 index 51a7e7a518599c..19e3bf7e7b40ac 100644 --- 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 @@ -12,32 +12,63 @@ 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'; -const chain = {} as RetrievalQAChain; - describe('EsqlLanguageKnowledgeBaseTool', () => { - describe('getTool', () => { - const esClient = { - search: jest.fn().mockResolvedValue({}), - } as unknown as ElasticsearchClient; - const request = { - body: { + 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, - 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 = { - esClient, - request, - }; + 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, - chain, modelExists: true, ...rest, }); @@ -48,7 +79,6 @@ describe('EsqlLanguageKnowledgeBaseTool', () => { it('returns null if modelExists is false (the ELSER model is not installed)', () => { const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ assistantLangChain: true, - chain, modelExists: false, // <-- ELSER model is not installed ...rest, }); @@ -59,7 +89,6 @@ describe('EsqlLanguageKnowledgeBaseTool', () => { it('should return a Tool instance if assistantLangChain and modelExists are true', () => { const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ assistantLangChain: true, - chain, modelExists: true, ...rest, }); @@ -70,7 +99,6 @@ describe('EsqlLanguageKnowledgeBaseTool', () => { it('should return a tool with the expected tags', () => { const tool = ESQL_KNOWLEDGE_BASE_TOOL.getTool({ assistantLangChain: true, - chain, modelExists: true, ...rest, }) as DynamicTool; 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 index 70891b661d2aff..3dc7dfd8d976c6 100644 --- 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 @@ -7,7 +7,7 @@ import { ChainTool } from 'langchain/tools'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; -import { APP_ID } from '../../../../common'; +import { APP_UI_ID } from '../../../../common'; export type EsqlKnowledgeBaseToolParams = AssistantToolParams; @@ -16,7 +16,7 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = { 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_ID, + sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is EsqlKnowledgeBaseToolParams => { const { assistantLangChain, modelExists } = params; return assistantLangChain && modelExists; 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 index c7d7e601bd4d19..c485c001704eba 100644 --- 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 @@ -11,47 +11,115 @@ import type { DynamicTool } from 'langchain/tools'; import { omit } from 'lodash/fp'; import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts_tool'; -import { mockAlertsFieldsApi } from '@kbn/elastic-assistant-plugin/server/__mocks__/alerts'; 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', () => { - describe('getTool', () => { - 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, + 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, - }, - } as unknown as KibanaRequest; - const assistantLangChain = true; - const chain = {} as unknown as RetrievalQAChain; - const modelExists = true; - const rest = { - assistantLangChain, - chain, - modelExists, - }; - - beforeEach(() => { - jest.clearAllMocks(); + ...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, - esClient, onNewReplacements: jest.fn(), replacements, request, @@ -157,7 +225,6 @@ describe('OpenAndAcknowledgedAlertsTool', () => { alertsIndexPattern, allow: requestWithMissingParams.body.allow, allowReplacement: requestWithMissingParams.body.allowReplacement, - esClient, onNewReplacements: jest.fn(), replacements, request: requestWithMissingParams, @@ -173,7 +240,6 @@ describe('OpenAndAcknowledgedAlertsTool', () => { // alertsIndexPattern is undefined allow: request.body.allow, allowReplacement: request.body.allowReplacement, - esClient, onNewReplacements: jest.fn(), replacements, request, @@ -189,7 +255,6 @@ describe('OpenAndAcknowledgedAlertsTool', () => { alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, - esClient, onNewReplacements: jest.fn(), replacements, request, @@ -205,7 +270,6 @@ describe('OpenAndAcknowledgedAlertsTool', () => { alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, - esClient, onNewReplacements: jest.fn(), replacements, request, @@ -221,7 +285,6 @@ describe('OpenAndAcknowledgedAlertsTool', () => { alertsIndexPattern, allow: request.body.allow, allowReplacement: request.body.allowReplacement, - esClient, onNewReplacements: jest.fn(), replacements, request, 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 index b4f6532ad6755d..252e83e851589d 100644 --- 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 @@ -13,7 +13,7 @@ import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-pl 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_ID } from '../../../../common'; +import { APP_UI_ID } from '../../../../common'; export interface OpenAndAcknowledgedAlertsToolParams extends AssistantToolParams { alertsIndexPattern: string; @@ -31,7 +31,7 @@ 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_ID, + sourceRegister: APP_UI_ID, isSupported: (params: AssistantToolParams): params is OpenAndAcknowledgedAlertsToolParams => { const { alertsIndexPattern, request, size } = params; return ( diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 159386673798ad..227a7ea7e1439d 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'; @@ -511,7 +511,7 @@ export class Plugin implements ISecuritySolutionPlugin { this.licensing$ = plugins.licensing.license$; // Assistant Tool and Feature Registration - plugins.elasticAssistant.registerTools(APP_ID, getAssistantTools()); + plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools()); if (this.lists && plugins.taskManager && plugins.fleet) { // Exceptions, Artifacts and Manifests start From dcf4c82e0ae03181786070439dbf41fa5b2b4eab Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 14 Dec 2023 00:22:45 -0700 Subject: [PATCH 11/13] Updating tool names --- .../server/assistant/tools/alert_counts/alert_counts_tool.ts | 2 +- .../open_and_acknowledged_alerts_tool.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 index 7cdd2342cd1f3a..e1c265c3dc239b 100644 --- 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 @@ -32,7 +32,7 @@ export const ALERT_COUNTS_TOOL: AssistantTool = { if (!this.isSupported(params)) return null; const { alertsIndexPattern, esClient } = params as AlertCountsToolParams; return new DynamicTool({ - name: 'alert-counts', + name: 'AlertCountsTool', description: ALERT_COUNTS_TOOL_DESCRIPTION, func: async () => { const query = getAlertsCountQuery(alertsIndexPattern); 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 index 252e83e851589d..210fdd42c1555f 100644 --- 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 @@ -54,7 +54,7 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = { size, } = params as OpenAndAcknowledgedAlertsToolParams; return new DynamicTool({ - name: 'open-alerts', + name: 'OpenAndAcknowledgedAlertsTool', description: OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION, func: async () => { const query = getOpenAndAcknowledgedAlertsQuery({ From 70186ab4c837c88bb6e3d99210ebf91c9a402369 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 14 Dec 2023 12:14:19 -0700 Subject: [PATCH 12/13] Last comments from review --- .../server/routes/helpers.test.ts | 25 ++++++++----- .../server/routes/helpers.ts | 9 +++-- .../routes/post_actions_connector_execute.ts | 8 +++-- .../alert_counts/alert_counts_tool.test.ts | 24 +++++++++++-- .../tools/alert_counts/alert_counts_tool.ts | 2 +- .../get_alert_counts_query.test.ts | 35 +++++++++++++------ .../alert_counts/get_alert_counts_query.ts | 35 +++++++++++++------ 7 files changed, 103 insertions(+), 35 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts index c35a045c9666d8..384e1f88657364 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.test.ts @@ -8,7 +8,7 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { RequestBody } from '../lib/langchain/types'; -import { getPluginNameFromRequest } from './helpers'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from './helpers'; describe('getPluginNameFromRequest', () => { const contextRequestHeaderEncoded = encodeURIComponent( @@ -32,27 +32,36 @@ describe('getPluginNameFromRequest', () => { }); it('extracts plugin name from "x-kbn-context" request header', async () => { - const pluginName = getPluginNameFromRequest({ request }); + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); expect(pluginName).toEqual('superSolution'); }); - it('fails to extracts plugin name from undefined "x-kbn-context" request header', async () => { + 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 }); - expect(pluginName).toEqual('securitySolutionUI'); + const pluginName = getPluginNameFromRequest({ + request: invalidRequest, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); + expect(pluginName).toEqual(DEFAULT_PLUGIN_NAME); }); - it('fails to extracts plugin name from malformed "x-kbn-context" invalidRequest header', async () => { + 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 }); - expect(pluginName).toEqual('securitySolutionUI'); + 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 index b4291cea8382e1..99d4493c16cca0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -11,6 +11,7 @@ import { RequestBody } from '../lib/langchain/types'; interface GetPluginNameFromRequestParams { request: KibanaRequest; + defaultPluginName: string; logger?: Logger; } @@ -26,12 +27,14 @@ export const DEFAULT_PLUGIN_NAME = 'securitySolutionUI'; * 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 { @@ -42,7 +45,9 @@ export const getPluginNameFromRequest = ({ )?.name; } } catch (err) { - logger?.error('Error determining source plugin, using default assistant tools.'); + logger?.error( + `Error determining source plugin for selecting tools, using ${defaultPluginName}.` + ); } - return DEFAULT_PLUGIN_NAME; + return defaultPluginName; }; 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 3ba021d33a02e5..9c1d8601da5320 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,7 +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 { getPluginNameFromRequest } from './helpers'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from './helpers'; export const postActionsConnectorExecuteRoute = ( router: IRouter, @@ -60,7 +60,11 @@ export const postActionsConnectorExecuteRoute = ( logger.debug('Executing via langchain, assistantLangChain: true'); // Fetch any tools registered by the request's originating plugin - const pluginName = getPluginNameFromRequest({ request, logger }); + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); const assistantTools = (await context.elasticAssistant).getRegisteredTools(pluginName); // get a scoped esClient for assistant memory 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 index 715d656944c86a..b128189f3c4f3f 100644 --- 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 @@ -103,11 +103,31 @@ describe('AlertCountsTool', () => { filter: [ { bool: { - filter: [{ match_phrase: { 'kibana.alert.workflow_status': 'open' } }], + must: [], + filter: [ + { + bool: { + should: [ + { match_phrase: { 'kibana.alert.workflow_status': 'open' } }, + { match_phrase: { 'kibana.alert.workflow_status': 'acknowledged' } }, + ], + minimum_should_match: 1, + }, + }, + { + range: { + '@timestamp': { + gte: 'now-1d/d', + lte: 'now/d', + format: 'strict_date_optional_time', + }, + }, + }, + ], + should: [], must_not: [{ exists: { field: 'kibana.alert.building_block_type' } }], }, }, - { range: { '@timestamp': { gte: 'now/d', lte: 'now/d' } } }, ], }, }, 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 index e1c265c3dc239b..837ceb3e753417 100644 --- 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 @@ -17,7 +17,7 @@ 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'; + 'Call this for the counts of last 24 hours of open and acknowledged alerts in the environment, grouped by their severity'; export const ALERT_COUNTS_TOOL: AssistantTool = { id: 'alert-counts-tool', diff --git a/x-pack/plugins/security_solution/server/assistant/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 index 27e210d53d51d6..46aa6c75e30e1e 100644 --- a/x-pack/plugins/security_solution/server/assistant/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 @@ -26,13 +26,36 @@ describe('getAlertsCountQuery', () => { filter: [ { bool: { + must: [], filter: [ { - match_phrase: { - 'kibana.alert.workflow_status': 'open', + bool: { + should: [ + { + match_phrase: { + 'kibana.alert.workflow_status': 'open', + }, + }, + { + match_phrase: { + 'kibana.alert.workflow_status': 'acknowledged', + }, + }, + ], + minimum_should_match: 1, + }, + }, + { + range: { + '@timestamp': { + gte: 'now-1d/d', + lte: 'now/d', + format: 'strict_date_optional_time', + }, }, }, ], + should: [], must_not: [ { exists: { @@ -42,14 +65,6 @@ describe('getAlertsCountQuery', () => { ], }, }, - { - range: { - '@timestamp': { - gte: 'now/d', - lte: 'now/d', - }, - }, - }, ], }, }, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts index 10ca556ad59e11..89d613bb0d5542 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts @@ -19,13 +19,36 @@ export const getAlertsCountQuery = (alertsIndexPattern: string) => ({ filter: [ { bool: { + must: [], filter: [ { - match_phrase: { - 'kibana.alert.workflow_status': 'open', + bool: { + should: [ + { + match_phrase: { + 'kibana.alert.workflow_status': 'open', + }, + }, + { + match_phrase: { + 'kibana.alert.workflow_status': 'acknowledged', + }, + }, + ], + minimum_should_match: 1, + }, + }, + { + range: { + '@timestamp': { + gte: 'now-1d/d', + lte: 'now/d', + format: 'strict_date_optional_time', + }, }, }, ], + should: [], must_not: [ { exists: { @@ -35,14 +58,6 @@ export const getAlertsCountQuery = (alertsIndexPattern: string) => ({ ], }, }, - { - range: { - '@timestamp': { - gte: 'now/d', - lte: 'now/d', - }, - }, - }, ], }, }, From 91d868b1491e61d6b594e1b4465b49a957c7729f Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 14 Dec 2023 12:22:57 -0700 Subject: [PATCH 13/13] Reverts alert counts query updates for acknowledged rules --- .../alert_counts/alert_counts_tool.test.ts | 24 ++----------- .../tools/alert_counts/alert_counts_tool.ts | 2 +- .../get_alert_counts_query.test.ts | 35 ++++++------------- .../alert_counts/get_alert_counts_query.ts | 35 ++++++------------- 4 files changed, 23 insertions(+), 73 deletions(-) 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 index b128189f3c4f3f..715d656944c86a 100644 --- 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 @@ -103,31 +103,11 @@ describe('AlertCountsTool', () => { filter: [ { bool: { - must: [], - filter: [ - { - bool: { - should: [ - { match_phrase: { 'kibana.alert.workflow_status': 'open' } }, - { match_phrase: { 'kibana.alert.workflow_status': 'acknowledged' } }, - ], - minimum_should_match: 1, - }, - }, - { - range: { - '@timestamp': { - gte: 'now-1d/d', - lte: 'now/d', - format: 'strict_date_optional_time', - }, - }, - }, - ], - should: [], + 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' } } }, ], }, }, 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 index 837ceb3e753417..e1c265c3dc239b 100644 --- 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 @@ -17,7 +17,7 @@ export interface AlertCountsToolParams extends AssistantToolParams { alertsIndexPattern: string; } export const ALERT_COUNTS_TOOL_DESCRIPTION = - 'Call this for the counts of last 24 hours of open and acknowledged alerts in the environment, grouped by their severity'; + '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', diff --git a/x-pack/plugins/security_solution/server/assistant/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 index 46aa6c75e30e1e..27e210d53d51d6 100644 --- a/x-pack/plugins/security_solution/server/assistant/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 @@ -26,36 +26,13 @@ describe('getAlertsCountQuery', () => { filter: [ { bool: { - must: [], filter: [ { - bool: { - should: [ - { - match_phrase: { - 'kibana.alert.workflow_status': 'open', - }, - }, - { - match_phrase: { - 'kibana.alert.workflow_status': 'acknowledged', - }, - }, - ], - minimum_should_match: 1, - }, - }, - { - range: { - '@timestamp': { - gte: 'now-1d/d', - lte: 'now/d', - format: 'strict_date_optional_time', - }, + match_phrase: { + 'kibana.alert.workflow_status': 'open', }, }, ], - should: [], must_not: [ { exists: { @@ -65,6 +42,14 @@ describe('getAlertsCountQuery', () => { ], }, }, + { + range: { + '@timestamp': { + gte: 'now/d', + lte: 'now/d', + }, + }, + }, ], }, }, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts index 89d613bb0d5542..10ca556ad59e11 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/alert_counts/get_alert_counts_query.ts @@ -19,36 +19,13 @@ export const getAlertsCountQuery = (alertsIndexPattern: string) => ({ filter: [ { bool: { - must: [], filter: [ { - bool: { - should: [ - { - match_phrase: { - 'kibana.alert.workflow_status': 'open', - }, - }, - { - match_phrase: { - 'kibana.alert.workflow_status': 'acknowledged', - }, - }, - ], - minimum_should_match: 1, - }, - }, - { - range: { - '@timestamp': { - gte: 'now-1d/d', - lte: 'now/d', - format: 'strict_date_optional_time', - }, + match_phrase: { + 'kibana.alert.workflow_status': 'open', }, }, ], - should: [], must_not: [ { exists: { @@ -58,6 +35,14 @@ export const getAlertsCountQuery = (alertsIndexPattern: string) => ({ ], }, }, + { + range: { + '@timestamp': { + gte: 'now/d', + lte: 'now/d', + }, + }, + }, ], }, },