[Security Solution] [Elastic AI Assistant] Adds support for arbitrary tool registration #172234
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
spong
added
release_note:skip
x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts
Outdated
Show resolved
Hide resolved
| }, | ||
| getTool(params: AssistantToolParams) { | ||
| if (!this.isSupported(params)) return null; | ||
| const { alertsIndexPattern, esClient } = params as AlertCountsToolParams; |
There was a problem hiding this comment.
Not required for this PR, but per our offline discussion the cast feels strange here because isSupported is a type guard that validates params. (This comment is applicable to all implementations of getTool in the PR.)
There was a problem hiding this comment.
I gave another pass at this and was not successful in getting type narrowing working with the isSupported function. Tried hoisting the function, re-declaring, and still needing the as cast. Will try again with fresh 👀's tomorrow.
| } | ||
|
|
||
| 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'; |
There was a problem hiding this comment.
thanks for adding acknowledged to the tool description per this comment in #173121 🙏
|
Pinging @elastic/security-solution (Team: SecuritySolution) |
x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts
Show resolved
Hide resolved
| @@ -58,6 +59,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 pluginName = getPluginNameFromRequest({ request, logger }); | |||
There was a problem hiding this comment.
Nice solution to ensure getRegisteredTools() returns tools scoped to the plugin! 🎉
andrew-goldstein
approved these changes
Dec 14, 2023
There was a problem hiding this comment.
Thanks @spong for making it possible for any plugin to register LangChain tools 🙏
✅ Desk tested locally
LGTM 🚀
💚 Build Succeeded
Metrics [docs]Public APIs missing comments
History
To update your PR or re-run it, just comment with: cc @spong |
Merged
2 tasks
spong
added a commit
that referenced
this pull request
Jan 11, 2024
…ature registration (#174317) ## Summary Resolves #172509 Adds ability to register feature capabilities through the assistant server so they no longer need to be plumbed through the `ElasticAssistantProvider`, which now also makes them available server side. Adds new `/internal/elastic_assistant/capabilities` route and `useCapabilities()` UI hook for fetching capabilities. ### OpenAPI Codegen Implemented using the new OpenAPI codegen and bundle packages: * Includes OpenAPI codegen script and CI action as detailed in: #166269 * Includes OpenAPI docs bundling script as detailed in: #171526 To run codegen/bundling locally, cd to `x-pack/plugins/elastic_assistant/` and run any of the following commands: ```bash yarn openapi:generate yarn openapi:generate:debug yarn openapi:bundle ``` > [!NOTE] > At the moment `yarn openapi:bundle` will output an empty bundled schema since `get_capabilities_route` is an internal route, this is to be expected. Also, if you don't see the file in your IDE, it's probably because `target` directories are ignored, so you may need to manually find/open the bundled schema at it's target location: `/x-pack/plugins/elastic_assistant/target/openapi/elastic_assistant.bundled.schema.yaml` ### Registering Capabilities To register a capability on plugin start, add the following in the consuming plugin's `start()`: ```ts plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled, }); ``` ### Declaring Feature Capabilities Feature capabilities are declared in `x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`: ```ts /** * Interfaces for features available to the elastic assistant */ export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean }; export const assistantFeatures = Object.freeze({ assistantModelEvaluation: false, assistantStreamingEnabled: false, }); ``` ### Using Capabilities Client Side And can be fetched client side using the `useCapabilities()` hook ala: ```ts // Fetch assistant capabilities const { data: capabilities } = useCapabilities({ http, toasts }); const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures; ``` ### Using Capabilities Server Side Or server side within a route (or elsewhere) via the `assistantContext`: ```ts const assistantContext = await context.elasticAssistant; const pluginName = getPluginNameFromRequest({ request, logger }); const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName); if (!registeredFeatures.assistantModelEvaluation) { return response.notFound(); } ``` > [!NOTE] > Note, just as with [registering arbitrary tools](#172234), features are registered for a specific plugin, where the plugin name that corresponds to your application is defined in the `x-kbn-context` header of requests made from your application, which may be different than your plugin's registered `APP_ID`. Perhaps this separation of concerns from one plugin to another isn't necessary, but it was easy to add matching the behavior of registering arbitrary tools. We can remove this granularity in favor of global features if desired. ### Test Steps * Verify `/internal/elastic_assistant/capabilities` route is called on security solution page load in dev tools, and that by default the `Evaluation` UI in setting does is not displayed and `404`'s if manually called. * Set the below experimental feature flag in your `kibana.dev.yml` and observe the feature being enabled by inspecting the capabilities api response, and that the evaluation feature becomes available: ``` xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation'] ``` * Run the `yarn openapi:*` codegen scripts above and ensure they execute as expected (code is generated/bundled) ### Checklist - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
eokoneyo
pushed a commit
to eokoneyo/kibana
that referenced
this pull request
Jan 12, 2024
…ature registration (elastic#174317) ## Summary Resolves elastic#172509 Adds ability to register feature capabilities through the assistant server so they no longer need to be plumbed through the `ElasticAssistantProvider`, which now also makes them available server side. Adds new `/internal/elastic_assistant/capabilities` route and `useCapabilities()` UI hook for fetching capabilities. ### OpenAPI Codegen Implemented using the new OpenAPI codegen and bundle packages: * Includes OpenAPI codegen script and CI action as detailed in: elastic#166269 * Includes OpenAPI docs bundling script as detailed in: elastic#171526 To run codegen/bundling locally, cd to `x-pack/plugins/elastic_assistant/` and run any of the following commands: ```bash yarn openapi:generate yarn openapi:generate:debug yarn openapi:bundle ``` > [!NOTE] > At the moment `yarn openapi:bundle` will output an empty bundled schema since `get_capabilities_route` is an internal route, this is to be expected. Also, if you don't see the file in your IDE, it's probably because `target` directories are ignored, so you may need to manually find/open the bundled schema at it's target location: `/x-pack/plugins/elastic_assistant/target/openapi/elastic_assistant.bundled.schema.yaml` ### Registering Capabilities To register a capability on plugin start, add the following in the consuming plugin's `start()`: ```ts plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled, }); ``` ### Declaring Feature Capabilities Feature capabilities are declared in `x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`: ```ts /** * Interfaces for features available to the elastic assistant */ export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean }; export const assistantFeatures = Object.freeze({ assistantModelEvaluation: false, assistantStreamingEnabled: false, }); ``` ### Using Capabilities Client Side And can be fetched client side using the `useCapabilities()` hook ala: ```ts // Fetch assistant capabilities const { data: capabilities } = useCapabilities({ http, toasts }); const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures; ``` ### Using Capabilities Server Side Or server side within a route (or elsewhere) via the `assistantContext`: ```ts const assistantContext = await context.elasticAssistant; const pluginName = getPluginNameFromRequest({ request, logger }); const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName); if (!registeredFeatures.assistantModelEvaluation) { return response.notFound(); } ``` > [!NOTE] > Note, just as with [registering arbitrary tools](elastic#172234), features are registered for a specific plugin, where the plugin name that corresponds to your application is defined in the `x-kbn-context` header of requests made from your application, which may be different than your plugin's registered `APP_ID`. Perhaps this separation of concerns from one plugin to another isn't necessary, but it was easy to add matching the behavior of registering arbitrary tools. We can remove this granularity in favor of global features if desired. ### Test Steps * Verify `/internal/elastic_assistant/capabilities` route is called on security solution page load in dev tools, and that by default the `Evaluation` UI in setting does is not displayed and `404`'s if manually called. * Set the below experimental feature flag in your `kibana.dev.yml` and observe the feature being enabled by inspecting the capabilities api response, and that the evaluation feature becomes available: ``` xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation'] ``` * Run the `yarn openapi:*` codegen scripts above and ensure they execute as expected (code is generated/bundled) ### Checklist - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
semd
pushed a commit
to semd/kibana
that referenced
this pull request
Jan 12, 2024
…ature registration (elastic#174317) ## Summary Resolves elastic#172509 Adds ability to register feature capabilities through the assistant server so they no longer need to be plumbed through the `ElasticAssistantProvider`, which now also makes them available server side. Adds new `/internal/elastic_assistant/capabilities` route and `useCapabilities()` UI hook for fetching capabilities. ### OpenAPI Codegen Implemented using the new OpenAPI codegen and bundle packages: * Includes OpenAPI codegen script and CI action as detailed in: elastic#166269 * Includes OpenAPI docs bundling script as detailed in: elastic#171526 To run codegen/bundling locally, cd to `x-pack/plugins/elastic_assistant/` and run any of the following commands: ```bash yarn openapi:generate yarn openapi:generate:debug yarn openapi:bundle ``` > [!NOTE] > At the moment `yarn openapi:bundle` will output an empty bundled schema since `get_capabilities_route` is an internal route, this is to be expected. Also, if you don't see the file in your IDE, it's probably because `target` directories are ignored, so you may need to manually find/open the bundled schema at it's target location: `/x-pack/plugins/elastic_assistant/target/openapi/elastic_assistant.bundled.schema.yaml` ### Registering Capabilities To register a capability on plugin start, add the following in the consuming plugin's `start()`: ```ts plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled, }); ``` ### Declaring Feature Capabilities Feature capabilities are declared in `x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`: ```ts /** * Interfaces for features available to the elastic assistant */ export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean }; export const assistantFeatures = Object.freeze({ assistantModelEvaluation: false, assistantStreamingEnabled: false, }); ``` ### Using Capabilities Client Side And can be fetched client side using the `useCapabilities()` hook ala: ```ts // Fetch assistant capabilities const { data: capabilities } = useCapabilities({ http, toasts }); const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures; ``` ### Using Capabilities Server Side Or server side within a route (or elsewhere) via the `assistantContext`: ```ts const assistantContext = await context.elasticAssistant; const pluginName = getPluginNameFromRequest({ request, logger }); const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName); if (!registeredFeatures.assistantModelEvaluation) { return response.notFound(); } ``` > [!NOTE] > Note, just as with [registering arbitrary tools](elastic#172234), features are registered for a specific plugin, where the plugin name that corresponds to your application is defined in the `x-kbn-context` header of requests made from your application, which may be different than your plugin's registered `APP_ID`. Perhaps this separation of concerns from one plugin to another isn't necessary, but it was easy to add matching the behavior of registering arbitrary tools. We can remove this granularity in favor of global features if desired. ### Test Steps * Verify `/internal/elastic_assistant/capabilities` route is called on security solution page load in dev tools, and that by default the `Evaluation` UI in setting does is not displayed and `404`'s if manually called. * Set the below experimental feature flag in your `kibana.dev.yml` and observe the feature being enabled by inspecting the capabilities api response, and that the evaluation feature becomes available: ``` xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation'] ``` * Run the `yarn openapi:*` codegen scripts above and ensure they execute as expected (code is generated/bundled) ### Checklist - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
CoenWarmer
pushed a commit
to CoenWarmer/kibana
that referenced
this pull request
Feb 15, 2024
…ature registration (elastic#174317) ## Summary Resolves elastic#172509 Adds ability to register feature capabilities through the assistant server so they no longer need to be plumbed through the `ElasticAssistantProvider`, which now also makes them available server side. Adds new `/internal/elastic_assistant/capabilities` route and `useCapabilities()` UI hook for fetching capabilities. ### OpenAPI Codegen Implemented using the new OpenAPI codegen and bundle packages: * Includes OpenAPI codegen script and CI action as detailed in: elastic#166269 * Includes OpenAPI docs bundling script as detailed in: elastic#171526 To run codegen/bundling locally, cd to `x-pack/plugins/elastic_assistant/` and run any of the following commands: ```bash yarn openapi:generate yarn openapi:generate:debug yarn openapi:bundle ``` > [!NOTE] > At the moment `yarn openapi:bundle` will output an empty bundled schema since `get_capabilities_route` is an internal route, this is to be expected. Also, if you don't see the file in your IDE, it's probably because `target` directories are ignored, so you may need to manually find/open the bundled schema at it's target location: `/x-pack/plugins/elastic_assistant/target/openapi/elastic_assistant.bundled.schema.yaml` ### Registering Capabilities To register a capability on plugin start, add the following in the consuming plugin's `start()`: ```ts plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled, }); ``` ### Declaring Feature Capabilities Feature capabilities are declared in `x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`: ```ts /** * Interfaces for features available to the elastic assistant */ export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean }; export const assistantFeatures = Object.freeze({ assistantModelEvaluation: false, assistantStreamingEnabled: false, }); ``` ### Using Capabilities Client Side And can be fetched client side using the `useCapabilities()` hook ala: ```ts // Fetch assistant capabilities const { data: capabilities } = useCapabilities({ http, toasts }); const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures; ``` ### Using Capabilities Server Side Or server side within a route (or elsewhere) via the `assistantContext`: ```ts const assistantContext = await context.elasticAssistant; const pluginName = getPluginNameFromRequest({ request, logger }); const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName); if (!registeredFeatures.assistantModelEvaluation) { return response.notFound(); } ``` > [!NOTE] > Note, just as with [registering arbitrary tools](elastic#172234), features are registered for a specific plugin, where the plugin name that corresponds to your application is defined in the `x-kbn-context` header of requests made from your application, which may be different than your plugin's registered `APP_ID`. Perhaps this separation of concerns from one plugin to another isn't necessary, but it was easy to add matching the behavior of registering arbitrary tools. We can remove this granularity in favor of global features if desired. ### Test Steps * Verify `/internal/elastic_assistant/capabilities` route is called on security solution page load in dev tools, and that by default the `Evaluation` UI in setting does is not displayed and `404`'s if manually called. * Set the below experimental feature flag in your `kibana.dev.yml` and observe the feature being enabled by inspecting the capabilities api response, and that the evaluation feature becomes available: ``` xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation'] ``` * Run the `yarn openapi:*` codegen scripts above and ensure they execute as expected (code is generated/bundled) ### Checklist - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Adds support for arbitrary tool registration on plugin start, allowing any plugin consuming the
elastic_assistantplugin to register tools specific to their use case.This was achieved by introducing an
AppContextService, with functionsregisterToolsandgetRegisteredTools.When registering within another plugin's start contract, call
registerTools, e.g.and to get any tools for a given plugin, call
getRegisteredToolswith the plugin's name. E.g here within the assistant request handler:Note
Use the plugin name that corresponds to your application as defined in the
x-kbn-contextheader of requests made from your application. This is the value used to fetch relevant tools, which may be different than your plugin's registeredAPP_ID.Test instructions
Since no new tool functionality has been enabled here, you can just follow the test instructions from #173121, and ensure that enabling/disabling the KB or RAG on Alerts still functions as expected. You should see the registered tools as a debug logger, e.g.
Checklist
Delete any items that are not applicable to this PR.
Any text added follows EUI's writing guidelines, uses sentence case text and includes i18n supportDocumentation was added for features that require explanation or tutorials