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

spong · 2023-11-30T00:42:13Z

Summary

Adds support for arbitrary tool registration on plugin start, allowing any plugin consuming the elastic_assistant plugin to register tools specific to their use case.

This was achieved by introducing an AppContextService, with functions registerTools and getRegisteredTools.

When registering within another plugin's start contract, call registerTools, e.g.

// Assistant Tool and Feature Registration
plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools());

and to get any tools for a given plugin, call getRegisteredTools with the plugin's name. E.g here within the assistant request handler:

// Fetch any tools registered by the request's originating plugin
const pluginName = getPluginNameFromRequest({ request, logger });
const assistantTools = (await context.elasticAssistant).getRegisteredTools(pluginName);

Note

Use the plugin name that corresponds to your application as defined in the x-kbn-context header of requests made from your application. This is the value used to fetch relevant tools, which may be different than your plugin's registered APP_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.

[DEBUG][plugins.elasticAssistant] AppContextService:getRegisteredTools
[DEBUG][plugins.elasticAssistant] pluginName: securitySolutionUI
[DEBUG][plugins.elasticAssistant] tools: AlertCountsTool, ESQLKnowledgeBaseTool, OpenAndAcknowledgedAlertsTool
[DEBUG][plugins.elasticAssistant] applicable tools: "alert-counts, ESQLKnowledgeBaseTool, open-alerts"

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 support~~
~~Documentation was added for features that require explanation or tutorials~~
Unit or functional tests were updated or added to match the most common scenarios

…fix'

x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts

andrew-goldstein · 2023-12-14T00:05:33Z

x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts

+  },
+  getTool(params: AssistantToolParams) {
+    if (!this.isSupported(params)) return null;
+    const { alertsIndexPattern, esClient } = params as AlertCountsToolParams;


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.)

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.

andrew-goldstein · 2023-12-14T00:13:36Z

...ion/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts

+}
+
+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';


thanks for adding acknowledged to the tool description per this comment in #173121 🙏

elasticmachine · 2023-12-14T07:22:34Z

Pinging @elastic/security-solution (Team: SecuritySolution)

x-pack/plugins/elastic_assistant/server/routes/helpers.ts

x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts

andrew-goldstein · 2023-12-14T19:00:37Z

x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts

@@ -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 });


Nice solution to ensure getRegisteredTools() returns tools scoped to the plugin! 🎉

andrew-goldstein

Thanks @spong for making it possible for any plugin to register LangChain tools 🙏
✅ Desk tested locally
LGTM 🚀

kibana-ci · 2023-12-14T20:57:54Z

💚 Build Succeeded

Buildkite Build
Commit: 91d868b

Metrics [docs]

Public APIs missing comments

Total count of every public API that lacks a comment. Target amount is 0. Run node scripts/build_api_docs --plugin [yourplugin] --stats comments for more detailed information.

id	before	after	diff
`elasticAssistant`	2	28	+26

Unknown metric groups

API count

id	before	after	diff
`elasticAssistant`	4	36	+32

History

💚 Build #183555 succeeded dcf4c82
💔 Build #183513 failed 7ab2e42
💚 Build #183453 succeeded 92211a2
💔 Build #183337 failed d4e52f5
💔 Build #183332 failed 9c79925
💔 Build #182786 failed bef4775

To update your PR or re-run it, just comment with:
@elasticmachine merge upstream

cc @spong

…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 <[email protected]>

…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 <[email protected]>

spong added 2 commits November 29, 2023 17:35

Support registration of arbitrary tools on plugin start

9aeed7c

Merge branch 'main' of github.com:elastic/kibana into tool-registry

f7b50dc

spong added release_note:skip Skip the PR/issue when compiling release notes backport:skip This commit does not require backporting Team: SecuritySolution Security Solutions Team working on SIEM, Endpoint, Timeline, Resolver, etc. Feature:Security Assistant Security Assistant labels Nov 30, 2023

spong requested review from andrew-goldstein and YulNaumenko November 30, 2023 00:42

spong self-assigned this Nov 30, 2023

kibanamachine and others added 5 commits November 30, 2023 00:49

[CI] Auto-commit changed files from 'node scripts/lint_ts_projects --…

5df33d5

…fix'

Merge branch 'main' of github.com:elastic/kibana into tool-registry

f2f2e51

Update mocks

38a3941

Merge branch 'main' into tool-registry

5096191

Merge branch 'main' into tool-registry

36be9e6

spong mentioned this pull request Dec 5, 2023

[Security Solution] [Elastic AI Assistant] Retrieval Augmented Generation (RAG) for Alerts #172542

Merged

spong added 6 commits December 6, 2023 13:29

Merge branch 'main' of github.com:elastic/kibana into tool-registry

c713281

Merge branch 'main' of github.com:elastic/kibana into tool-registry

c10e546

Move new RAG on alerts tools to Security Solution plugin

52d57cf

Merge branch 'main' of github.com:elastic/kibana into tool-registry

bef4775

Adds tool metadata and separates params validation from tool creation

68e1ceb

Merge branch 'main' of github.com:elastic/kibana into tool-registry

9c79925

spong force-pushed the tool-registry branch from 230835e to 9c79925 Compare December 13, 2023 06:10

kibanamachine and others added 3 commits December 13, 2023 06:17

[CI] Auto-commit changed files from 'node scripts/lint_ts_projects --…

d4e52f5

…fix'

Adds jest config for new security solution assistant server tests

8565784

Merge branch 'main' of github.com:elastic/kibana into tool-registry

92211a2

spong mentioned this pull request Dec 13, 2023

[Security Solution] [Elastic AI Assistant] Include acknowledged alerts in the context sent to the LLM (Retrieval Augmented Generation (RAG) for Alerts) #173121

Merged

spong added 3 commits December 13, 2023 11:53

Merge branch 'main' of github.com:elastic/kibana into tool-registry

cf5eea4

Renames to match new tool names from merge w/ main

0f47898

Additional missed renames

7ab2e42

andrew-goldstein reviewed Dec 14, 2023

View reviewed changes

x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts Outdated Show resolved Hide resolved

andrew-goldstein reviewed Dec 14, 2023

View reviewed changes

spong added 2 commits December 14, 2023 00:07

Extracts plugin name from request and adds tests

2f65edb

Merge branch 'main' of github.com:elastic/kibana into tool-registry

19b7532

spong marked this pull request as ready for review December 14, 2023 07:22

spong requested a review from a team as a code owner December 14, 2023 07:22

Updating tool names

dcf4c82

andrew-goldstein reviewed Dec 14, 2023

View reviewed changes

x-pack/plugins/elastic_assistant/server/routes/helpers.ts Outdated Show resolved Hide resolved

andrew-goldstein reviewed Dec 14, 2023

View reviewed changes

x-pack/plugins/elastic_assistant/server/routes/helpers.ts Outdated Show resolved Hide resolved

andrew-goldstein reviewed Dec 14, 2023

View reviewed changes

x-pack/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts Show resolved Hide resolved

andrew-goldstein reviewed Dec 14, 2023

View reviewed changes

spong added 2 commits December 14, 2023 12:14

Last comments from review

70186ab

Reverts alert counts query updates for acknowledged rules

91d868b

andrew-goldstein approved these changes Dec 14, 2023

View reviewed changes

spong merged commit a1d8272 into elastic:main Dec 14, 2023

spong deleted the tool-registry branch December 14, 2023 21:53

kibanamachine added the v8.13.0 label Dec 14, 2023

spong mentioned this pull request Jan 5, 2024

[Security Solution] [Elastic AI Assistant] Adds support for plugin feature registration #174317

Merged

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

spong commented Nov 30, 2023 •

edited

Loading

andrew-goldstein Dec 14, 2023

spong Dec 14, 2023

andrew-goldstein Dec 14, 2023

elasticmachine commented Dec 14, 2023

andrew-goldstein Dec 14, 2023

andrew-goldstein left a comment

kibana-ci commented Dec 14, 2023

API count

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

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

Conversation

spong commented Nov 30, 2023 • edited Loading

Summary

Test instructions

Checklist

andrew-goldstein Dec 14, 2023

Choose a reason for hiding this comment

spong Dec 14, 2023

Choose a reason for hiding this comment

andrew-goldstein Dec 14, 2023

Choose a reason for hiding this comment

elasticmachine commented Dec 14, 2023

andrew-goldstein Dec 14, 2023

Choose a reason for hiding this comment

andrew-goldstein left a comment

Choose a reason for hiding this comment

kibana-ci commented Dec 14, 2023

💚 Build Succeeded

Metrics [docs]

Public APIs missing comments

API count

History

spong commented Nov 30, 2023 •

edited

Loading