[Security Solution] [Elastic AI Assistant] Adds support for arbitrary tool registration

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<ElasticAssistantApiRequestHandlerContext> => {
   return {
     actions: clients.elasticAssistant.actions as unknown as ActionsPluginStart,
+    getRegisteredTools: jest.fn(),
     logger: clients.elasticAssistant.logger,
   };
 };

x-pack/plugins/elastic_assistant/server/index.ts

@@ -13,6 +13,10 @@ export async function plugin(initializerContext: PluginInitializerContext) {
 }
 
 export type {
-  ElasticAssistantPluginSetup as EcsDataQualityDashboardPluginSetup,
-  ElasticAssistantPluginStart as EcsDataQualityDashboardPluginStart,
+  ElasticAssistantPluginSetup,
+  ElasticAssistantPluginStart,
+  ElasticAssistantPluginSetupDependencies,
+  ElasticAssistantPluginStartDependencies,
+  AssistantTool,
+  AssistantToolParams,
 } from './types';

x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts

@@ -16,7 +16,7 @@ import { KNOWLEDGE_BASE_INDEX_PATTERN } from '../../../routes/knowledge_base/con
 import type { AgentExecutorParams, AgentExecutorResponse } from '../executors/types';
 import { withAssistantSpan } from '../tracers/with_assistant_span';
 import { APMTracer } from '../tracers/apm_tracer';
-import { getApplicableTools } from '../tools';
+import { AssistantToolParams } from '../../../types';
 
 export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor';
 
@@ -31,6 +31,7 @@ export const callAgentExecutor = async ({
   allow,
   allowReplacement,
   assistantLangChain,
+  assistantTools = [],
   connectorId,
   elserId,
   esClient,
@@ -71,7 +72,8 @@ export const callAgentExecutor = async ({
   // Create a chain that uses the ELSER backed ElasticsearchStore, override k=10 for esql query generation for now
   const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever(10));
 
-  const tools: Tool[] = getApplicableTools({
+  // Fetch any applicable tools that the source plugin may have registered
+  const assistantToolParams: AssistantToolParams = {
     allow,
     allowReplacement,
     alertsIndexPattern,
@@ -83,7 +85,8 @@ export const callAgentExecutor = async ({
     replacements,
     request,
     size,
-  });
+  };
+  const tools: Tool[] = assistantTools.flatMap((tool) => tool.getTool(assistantToolParams) ?? []);
 
   logger.debug(`applicable tools: ${JSON.stringify(tools.map((t) => t.name).join(', '), null, 2)}`);
 

x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts

@@ -12,13 +12,15 @@ import { Logger } from '@kbn/logging';
 import { KibanaRequest } from '@kbn/core-http-server';
 import type { LangChainTracer } from 'langchain/callbacks';
 import { RequestBody, ResponseBody } from '../types';
+import type { AssistantTool } from '../../../types';
 
 export interface AgentExecutorParams {
   alertsIndexPattern?: string;
   actions: ActionsPluginStart;
   allow?: string[];
   allowReplacement?: string[];
   assistantLangChain: boolean;
+  assistantTools?: AssistantTool[];
   connectorId: string;
   esClient: ElasticsearchClient;
   kbResource: string | undefined;