File tree Expand file tree Collapse file tree 5 files changed +36
-7
lines changed
x-pack/solutions/security/plugins/elastic_assistant/server Expand file tree Collapse file tree 5 files changed +36
-7
lines changed Original file line number Diff line number Diff line change @@ -194,7 +194,7 @@ export class ElasticAssistantPlugin
194194 getRegisteredFeatures : ( pluginName : string ) => {
195195 return appContextService . getRegisteredFeatures ( pluginName ) ;
196196 } ,
197- getRegisteredTools : ( pluginName : string ) => {
197+ getRegisteredTools : ( pluginName : string | string [ ] ) => {
198198 return appContextService . getRegisteredTools ( pluginName ) ;
199199 } ,
200200 } ;
@@ -220,7 +220,7 @@ export class ElasticAssistantPlugin
220220 getRegisteredFeatures : ( pluginName : string ) => {
221221 return appContextService . getRegisteredFeatures ( pluginName ) ;
222222 } ,
223- getRegisteredTools : ( pluginName : string ) => {
223+ getRegisteredTools : ( pluginName : string | string [ ] ) => {
224224 return appContextService . getRegisteredTools ( pluginName ) ;
225225 } ,
226226 registerFeatures : ( pluginName : string , features : Partial < AssistantFeatures > ) => {
Original file line number Diff line number Diff line change @@ -293,8 +293,10 @@ export const langChainExecute = async ({
293293 const assistantContext = context . elasticAssistant ;
294294 // We don't (yet) support invoking these tools interactively
295295 const unsupportedTools = new Set ( [ 'attack-discovery' , DEFEND_INSIGHTS_ID ] ) ;
296+ const pluginNames = Array . from ( new Set ( [ pluginName , DEFAULT_PLUGIN_NAME ] ) ) ;
297+
296298 const assistantTools = assistantContext
297- . getRegisteredTools ( pluginName )
299+ . getRegisteredTools ( pluginNames )
298300 . filter ( ( tool ) => ! unsupportedTools . has ( tool . id ) ) ;
299301
300302 // get a scoped esClient for assistant memory
Original file line number Diff line number Diff line change @@ -107,7 +107,7 @@ export class RequestContextFactory implements IRequestContextFactory {
107107
108108 getCurrentUser,
109109
110- getRegisteredTools : ( pluginName : string ) => {
110+ getRegisteredTools : ( pluginName : string | string [ ] ) => {
111111 return appContextService . getRegisteredTools ( pluginName ) ;
112112 } ,
113113
Original file line number Diff line number Diff line change @@ -35,6 +35,14 @@ describe('AppContextService', () => {
3535 isSupported : jest . fn ( ) ,
3636 getTool : jest . fn ( ) ,
3737 } ;
38+ const toolThree : AssistantTool = {
39+ id : 'tool-three' ,
40+ name : 'ToolThree' ,
41+ description : 'Description 3' ,
42+ sourceRegister : 'Source3' ,
43+ isSupported : jest . fn ( ) ,
44+ getTool : jest . fn ( ) ,
45+ } ;
3846
3947 beforeEach ( ( ) => {
4048 appContextService . stop ( ) ;
@@ -99,6 +107,17 @@ describe('AppContextService', () => {
99107 } ) ;
100108 } ) ;
101109
110+ it ( 'get tools for multiple plugins' , ( ) => {
111+ const pluginName1 = 'pluginName1' ;
112+ const pluginName2 = 'pluginName2' ;
113+
114+ appContextService . start ( mockAppContext ) ;
115+ appContextService . registerTools ( pluginName2 , [ toolOne , toolThree ] ) ;
116+ appContextService . registerTools ( pluginName1 , [ toolOne ] ) ;
117+
118+ expect ( appContextService . getRegisteredTools ( [ pluginName1 , pluginName2 ] ) . length ) . toEqual ( 2 ) ;
119+ } ) ;
120+
102121 describe ( 'registering features' , ( ) => {
103122 it ( 'should register and get features for a single plugin' , ( ) => {
104123 const pluginName = 'pluginName' ;
Original file line number Diff line number Diff line change @@ -12,7 +12,7 @@ import { AssistantTool } from '../types';
1212export type PluginName = string ;
1313export type RegisteredToolsStorage = Map < PluginName , Set < AssistantTool > > ;
1414export type RegisteredFeaturesStorage = Map < PluginName , AssistantFeatures > ;
15- export type GetRegisteredTools = ( pluginName : string ) => AssistantTool [ ] ;
15+ export type GetRegisteredTools = ( pluginName : string | string [ ] ) => AssistantTool [ ] ;
1616export type GetRegisteredFeatures = ( pluginName : string ) => AssistantFeatures ;
1717export interface ElasticAssistantAppContext {
1818 logger : Logger ;
@@ -67,8 +67,16 @@ class AppContextService {
6767 *
6868 * @param pluginName
6969 */
70- public getRegisteredTools ( pluginName : string ) : AssistantTool [ ] {
71- const tools = Array . from ( this . registeredTools ?. get ( pluginName ) ?? new Set < AssistantTool > ( ) ) ;
70+ public getRegisteredTools ( pluginName : string | string [ ] ) : AssistantTool [ ] {
71+ const pluginNames = Array . isArray ( pluginName ) ? pluginName : [ pluginName ] ;
72+
73+ const tools = [
74+ ...new Set (
75+ pluginNames
76+ . map ( ( name ) => this . registeredTools ?. get ( name ) ?? new Set < AssistantTool > ( ) )
77+ . flatMap ( ( set ) => [ ...set ] )
78+ ) ,
79+ ] ;
7280
7381 this . logger ?. debug ( 'AppContextService:getRegisteredTools' ) ;
7482 this . logger ?. debug ( `pluginName: ${ pluginName } ` ) ;
You can’t perform that action at this time.
0 commit comments