From a207f5761be2b1632710212a880756fbb2aee85c Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 11 Jan 2024 08:22:02 -0700 Subject: [PATCH] [Security Solution] [Elastic AI Assistant] Adds support for plugin feature registration (#174317) ## Summary Resolves https://github.com/elastic/kibana/issues/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: https://github.com/elastic/kibana/pull/166269 * Includes OpenAPI docs bundling script as detailed in: https://github.com/elastic/kibana/pull/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](https://github.com/elastic/kibana/pull/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> --- .buildkite/scripts/steps/checks.sh | 1 + .../elastic_assistant_codegen.sh | 10 ++ .../impl/capabilities/README.md | 51 +++++++++++ .../impl/capabilities/index.ts | 19 ++++ .../kbn-elastic-assistant-common/index.ts | 3 + .../api/capabilities/capabilities.test.tsx | 43 +++++++++ .../api/capabilities/capabilities.tsx | 44 +++++++++ .../capabilities/use_capabilities.test.tsx | 49 ++++++++++ .../api/capabilities/use_capabilities.tsx | 52 +++++++++++ .../impl/assistant_context/index.test.tsx | 56 ++---------- .../impl/assistant_context/index.tsx | 11 ++- .../mock/test_providers/test_providers.tsx | 69 ++++++++------ .../elastic_assistant/common/constants.ts | 3 + x-pack/plugins/elastic_assistant/package.json | 5 +- .../scripts/openapi/bundle.js | 18 ++++ .../scripts/openapi/generate.js | 18 ++++ .../server/__mocks__/request.ts | 26 +++++- .../server/__mocks__/request_context.ts | 2 + .../server/__mocks__/server.ts | 63 +++++++++---- .../server/__mocks__/test_adapters.ts | 2 + .../elastic_assistant/server/plugin.ts | 47 ++++++---- .../get_capabilities_route.test.ts | 47 ++++++++++ .../capabilities/get_capabilities_route.ts | 60 ++++++++++++ .../routes/evaluate/post_evaluate.test.ts | 59 ++++++++++++ .../server/routes/evaluate/post_evaluate.ts | 14 ++- .../server/routes/helpers.ts | 3 +- .../get_knowledge_base_status.ts | 1 + .../get_capabilities_route.gen.ts | 19 ++++ .../get_capabilities_route.schema.yaml | 41 +++++++++ .../server/schemas/evaluate/post_evaluate.ts | 2 + .../server/services/app_context.test.ts | 91 +++++++++++++++++++ .../server/services/app_context.ts | 54 ++++++++++- .../plugins/elastic_assistant/server/types.ts | 30 +++++- .../plugins/elastic_assistant/tsconfig.json | 3 +- .../public/assistant/provider.tsx | 5 - .../rule_status_failed_callout.test.tsx | 60 +++++++----- .../security_solution/server/plugin.ts | 4 + 37 files changed, 930 insertions(+), 155 deletions(-) create mode 100755 .buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.test.tsx create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx create mode 100644 x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js create mode 100644 x-pack/plugins/elastic_assistant/scripts/openapi/generate.js create mode 100644 x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.gen.ts create mode 100644 x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.schema.yaml diff --git a/.buildkite/scripts/steps/checks.sh b/.buildkite/scripts/steps/checks.sh index 4d9461a6a2bff4c..42a0eb9ee61c6f4 100755 --- a/.buildkite/scripts/steps/checks.sh +++ b/.buildkite/scripts/steps/checks.sh @@ -23,6 +23,7 @@ export DISABLE_BOOTSTRAP_VALIDATION=false .buildkite/scripts/steps/checks/ftr_configs.sh .buildkite/scripts/steps/checks/saved_objects_compat_changes.sh .buildkite/scripts/steps/checks/saved_objects_definition_change.sh +.buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh .buildkite/scripts/steps/code_generation/security_solution_codegen.sh .buildkite/scripts/steps/code_generation/osquery_codegen.sh .buildkite/scripts/steps/checks/yarn_deduplicate.sh diff --git a/.buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh b/.buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh new file mode 100755 index 000000000000000..ba742625b421f06 --- /dev/null +++ b/.buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash + +set -euo pipefail + +source .buildkite/scripts/common/util.sh + +echo --- Elastic Assistant OpenAPI Code Generation + +(cd x-pack/plugins/elastic_assistant && yarn openapi:generate) +check_for_changed_files "yarn openapi:generate" true diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md new file mode 100644 index 000000000000000..5a471245e04493f --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md @@ -0,0 +1,51 @@ +### Feature Capabilities + +Feature capabilities are an object describing specific capabilities of the assistant, like whether a feature like streaming is enabled, and are defined in the sibling `./index.ts` file within this `kbn-elastic-assistant-common` package. These capabilities can be registered for a given plugin through the assistant server, and so do not need to be plumbed through the `ElasticAssistantProvider`. + +Storage and accessor functions are made available via the `AppContextService`, and exposed to clients via the`/internal/elastic_assistant/capabilities` route, which can be fetched by clients using the `useCapabilities()` UI hook. + +### Registering Capabilities + +To register a capability on plugin start, add the following in the consuming plugin's `start()`, specifying any number of capabilities you would like to explicitly declare: + +```ts +plugins.elasticAssistant.registerFeatures(APP_UI_ID, { + assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, + assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled, +}); +``` + +### Declaring Feature Capabilities +Default feature capabilities are declared in `x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`: + +```ts +export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: boolean }; + +export const defaultAssistantFeatures = Object.freeze({ + assistantModelEvaluation: false, + assistantStreamingEnabled: false, +}); +``` + +### Using Capabilities Client Side +Capabilities can be fetched client side using the `useCapabilities()` hook ala: + +```ts +const { data: capabilities } = useCapabilities({ http, toasts }); +const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? defaultAssistantFeatures; +``` + +### 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](https://github.com/elastic/kibana/pull/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`. diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts new file mode 100644 index 000000000000000..1d404309f73e391 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -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. + */ + +/** + * Interface for features available to the elastic assistant + */ +export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: boolean }; + +/** + * Default features available to the elastic assistant + */ +export const defaultAssistantFeatures = Object.freeze({ + assistantModelEvaluation: false, + assistantStreamingEnabled: false, +}); diff --git a/x-pack/packages/kbn-elastic-assistant-common/index.ts b/x-pack/packages/kbn-elastic-assistant-common/index.ts index f17e13a33af3d4a..c64b02160d6e433 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/index.ts @@ -5,6 +5,9 @@ * 2.0. */ +export { defaultAssistantFeatures } from './impl/capabilities'; +export type { AssistantFeatures } from './impl/capabilities'; + export { getAnonymizedValue } from './impl/data_anonymization/get_anonymized_value'; export { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.test.tsx new file mode 100644 index 000000000000000..b41d7ac1445549b --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.test.tsx @@ -0,0 +1,43 @@ +/* + * 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 { HttpSetup } from '@kbn/core-http-browser'; + +import { getCapabilities } from './capabilities'; +import { API_ERROR } from '../../translations'; + +jest.mock('@kbn/core-http-browser'); + +const mockHttp = { + fetch: jest.fn(), +} as unknown as HttpSetup; + +describe('Capabilities API tests', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('getCapabilities', () => { + it('calls the internal assistant API for fetching assistant capabilities', async () => { + await getCapabilities({ http: mockHttp }); + + expect(mockHttp.fetch).toHaveBeenCalledWith('/internal/elastic_assistant/capabilities', { + method: 'GET', + signal: undefined, + version: '1', + }); + }); + + it('returns API_ERROR when the response status is error', async () => { + (mockHttp.fetch as jest.Mock).mockResolvedValue({ status: API_ERROR }); + + const result = await getCapabilities({ http: mockHttp }); + + expect(result).toEqual({ status: API_ERROR }); + }); + }); +}); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx new file mode 100644 index 000000000000000..794b89e1775f8b7 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx @@ -0,0 +1,44 @@ +/* + * 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 { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser'; +import { AssistantFeatures } from '@kbn/elastic-assistant-common'; + +export interface GetCapabilitiesParams { + http: HttpSetup; + signal?: AbortSignal | undefined; +} + +export type GetCapabilitiesResponse = AssistantFeatures; + +/** + * API call for fetching assistant capabilities + * + * @param {Object} options - The options object. + * @param {HttpSetup} options.http - HttpSetup + * @param {AbortSignal} [options.signal] - AbortSignal + * + * @returns {Promise} + */ +export const getCapabilities = async ({ + http, + signal, +}: GetCapabilitiesParams): Promise => { + try { + const path = `/internal/elastic_assistant/capabilities`; + + const response = await http.fetch(path, { + method: 'GET', + signal, + version: '1', + }); + + return response as GetCapabilitiesResponse; + } catch (error) { + return error as IHttpFetchError; + } +}; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx new file mode 100644 index 000000000000000..c9e60b806d1bf37 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx @@ -0,0 +1,49 @@ +/* + * 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 { renderHook } from '@testing-library/react-hooks'; + +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import type { ReactNode } from 'react'; +import React from 'react'; +import { useCapabilities, UseCapabilitiesParams } from './use_capabilities'; + +const statusResponse = { assistantModelEvaluation: true, assistantStreamingEnabled: false }; + +const http = { + fetch: jest.fn().mockResolvedValue(statusResponse), +}; +const toasts = { + addError: jest.fn(), +}; +const defaultProps = { http, toasts } as unknown as UseCapabilitiesParams; + +const createWrapper = () => { + const queryClient = new QueryClient(); + // eslint-disable-next-line react/display-name + return ({ children }: { children: ReactNode }) => ( + {children} + ); +}; + +describe('useFetchRelatedCases', () => { + it(`should make http request to fetch capabilities`, () => { + renderHook(() => useCapabilities(defaultProps), { + wrapper: createWrapper(), + }); + + expect(defaultProps.http.fetch).toHaveBeenCalledWith( + '/internal/elastic_assistant/capabilities', + { + method: 'GET', + version: '1', + signal: new AbortController().signal, + } + ); + expect(toasts.addError).not.toHaveBeenCalled(); + }); +}); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx new file mode 100644 index 000000000000000..5d52a2801fb9ea7 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx @@ -0,0 +1,52 @@ +/* + * 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 { UseQueryResult } from '@tanstack/react-query'; +import { useQuery } from '@tanstack/react-query'; +import type { HttpSetup, IHttpFetchError, ResponseErrorBody } from '@kbn/core-http-browser'; +import type { IToasts } from '@kbn/core-notifications-browser'; +import { i18n } from '@kbn/i18n'; +import { getCapabilities, GetCapabilitiesResponse } from './capabilities'; + +const CAPABILITIES_QUERY_KEY = ['elastic-assistant', 'capabilities']; + +export interface UseCapabilitiesParams { + http: HttpSetup; + toasts?: IToasts; +} +/** + * Hook for getting the feature capabilities of the assistant + * + * @param {Object} options - The options object. + * @param {HttpSetup} options.http - HttpSetup + * @param {IToasts} options.toasts - IToasts + * + * @returns {useQuery} hook for getting the status of the Knowledge Base + */ +export const useCapabilities = ({ + http, + toasts, +}: UseCapabilitiesParams): UseQueryResult => { + return useQuery({ + queryKey: CAPABILITIES_QUERY_KEY, + queryFn: async ({ signal }) => { + return getCapabilities({ http, signal }); + }, + retry: false, + keepPreviousData: true, + // Deprecated, hoist to `queryCache` w/in `QueryClient. See: https://stackoverflow.com/a/76961109 + onError: (error: IHttpFetchError) => { + if (error.name !== 'AbortError') { + toasts?.addError(error.body && error.body.message ? new Error(error.body.message) : error, { + title: i18n.translate('xpack.elasticAssistant.capabilities.statusError', { + defaultMessage: 'Error fetching capabilities', + }), + }); + } + }, + }); +}; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx index 84a2ac40a6f2486..a8dc5b1aa1db764 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx @@ -6,53 +6,14 @@ */ import { renderHook } from '@testing-library/react-hooks'; -import React from 'react'; -import { AssistantProvider, useAssistantContext } from '.'; -import { httpServiceMock } from '@kbn/core-http-browser-mocks'; -import { actionTypeRegistryMock } from '@kbn/triggers-actions-ui-plugin/public/application/action_type_registry.mock'; -import { AssistantAvailability } from '../..'; +import { useAssistantContext } from '.'; import { useLocalStorage } from 'react-use'; +import { TestProviders } from '../mock/test_providers/test_providers'; jest.mock('react-use', () => ({ useLocalStorage: jest.fn().mockReturnValue(['456', jest.fn()]), })); -const actionTypeRegistry = actionTypeRegistryMock.create(); -const mockGetInitialConversations = jest.fn(() => ({})); -const mockGetComments = jest.fn(() => []); -const mockHttp = httpServiceMock.createStartContract({ basePath: '/test' }); -const mockAssistantAvailability: AssistantAvailability = { - hasAssistantPrivilege: false, - hasConnectorsAllPrivilege: true, - hasConnectorsReadPrivilege: true, - isAssistantEnabled: true, -}; - -const ContextWrapper: React.FC = ({ children }) => ( - - {children} - -); describe('AssistantContext', () => { beforeEach(() => jest.clearAllMocks()); @@ -66,30 +27,29 @@ describe('AssistantContext', () => { }); test('it should return the httpFetch function', async () => { - const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper }); - const http = await result.current.http; + const { result } = renderHook(useAssistantContext, { wrapper: TestProviders }); const path = '/path/to/resource'; - await http.fetch(path); + await result.current.http.fetch(path); - expect(mockHttp.fetch).toBeCalledWith(path); + expect(result.current.http.fetch).toBeCalledWith(path); }); test('getConversationId defaults to provided id', async () => { - const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper }); + const { result } = renderHook(useAssistantContext, { wrapper: TestProviders }); const id = result.current.getConversationId('123'); expect(id).toEqual('123'); }); test('getConversationId uses local storage id when no id is provided ', async () => { - const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper }); + const { result } = renderHook(useAssistantContext, { wrapper: TestProviders }); const id = result.current.getConversationId(); expect(id).toEqual('456'); }); test('getConversationId defaults to Welcome when no local storage id and no id is provided ', async () => { (useLocalStorage as jest.Mock).mockReturnValue([undefined, jest.fn()]); - const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper }); + const { result } = renderHook(useAssistantContext, { wrapper: TestProviders }); const id = result.current.getConversationId(); expect(id).toEqual('Welcome'); }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 50a3211f74f3cda..3f3102a4ea6bff4 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -13,6 +13,7 @@ import type { IToasts } from '@kbn/core-notifications-browser'; import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public'; import { useLocalStorage } from 'react-use'; import type { DocLinksStart } from '@kbn/core-doc-links-browser'; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { WELCOME_CONVERSATION_TITLE } from '../assistant/use_conversation/translations'; import { updatePromptContexts } from './helpers'; import type { @@ -37,6 +38,7 @@ import { } from './constants'; import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings'; import { AssistantAvailability, AssistantTelemetry } from './types'; +import { useCapabilities } from '../assistant/api/capabilities/use_capabilities'; export interface ShowAssistantOverlayProps { showOverlay: boolean; @@ -53,7 +55,6 @@ export interface AssistantProviderProps { actionTypeRegistry: ActionTypeRegistryContract; alertsIndexPattern?: string; assistantAvailability: AssistantAvailability; - assistantStreamingEnabled?: boolean; assistantTelemetry?: AssistantTelemetry; augmentMessageCodeBlocks: (currentConversation: Conversation) => CodeBlockDetails[][]; baseAllow: string[]; @@ -87,7 +88,6 @@ export interface AssistantProviderProps { }) => EuiCommentProps[]; http: HttpSetup; getInitialConversations: () => Record; - modelEvaluatorEnabled?: boolean; nameSpace?: string; setConversations: React.Dispatch>>; setDefaultAllow: React.Dispatch>; @@ -163,7 +163,6 @@ export const AssistantProvider: React.FC = ({ actionTypeRegistry, alertsIndexPattern, assistantAvailability, - assistantStreamingEnabled = false, assistantTelemetry, augmentMessageCodeBlocks, baseAllow, @@ -179,7 +178,6 @@ export const AssistantProvider: React.FC = ({ getComments, http, getInitialConversations, - modelEvaluatorEnabled = false, nameSpace = DEFAULT_ASSISTANT_NAMESPACE, setConversations, setDefaultAllow, @@ -298,6 +296,11 @@ export const AssistantProvider: React.FC = ({ [localStorageLastConversationId] ); + // Fetch assistant capabilities + const { data: capabilities } = useCapabilities({ http, toasts }); + const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = + capabilities ?? defaultAssistantFeatures; + const value = useMemo( () => ({ actionTypeRegistry, diff --git a/x-pack/packages/security-solution/ecs_data_quality_dashboard/impl/data_quality/mock/test_providers/test_providers.tsx b/x-pack/packages/security-solution/ecs_data_quality_dashboard/impl/data_quality/mock/test_providers/test_providers.tsx index b2bd63f8101aa1c..175380cc5169ad9 100644 --- a/x-pack/packages/security-solution/ecs_data_quality_dashboard/impl/data_quality/mock/test_providers/test_providers.tsx +++ b/x-pack/packages/security-solution/ecs_data_quality_dashboard/impl/data_quality/mock/test_providers/test_providers.tsx @@ -13,6 +13,7 @@ import { euiDarkVars } from '@kbn/ui-theme'; import React from 'react'; import { ThemeProvider } from 'styled-components'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import { DataQualityProvider } from '../../data_quality_panel/data_quality_context'; interface Props { @@ -39,38 +40,52 @@ export const TestProvidersComponent: React.FC = ({ children, isILMAvailab hasConnectorsReadPrivilege: true, isAssistantEnabled: true, }; + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + logger: { + log: jest.fn(), + warn: jest.fn(), + error: () => {}, + }, + }); return ( ({ eui: euiDarkVars, darkMode: true })}> - - + - {children} - - + + {children} + + + ); diff --git a/x-pack/plugins/elastic_assistant/common/constants.ts b/x-pack/plugins/elastic_assistant/common/constants.ts index 100aebf3952875f..5634b0b1881bdb6 100755 --- a/x-pack/plugins/elastic_assistant/common/constants.ts +++ b/x-pack/plugins/elastic_assistant/common/constants.ts @@ -17,3 +17,6 @@ export const KNOWLEDGE_BASE = `${BASE_PATH}/knowledge_base/{resource?}`; // Model Evaluation export const EVALUATE = `${BASE_PATH}/evaluate`; + +// Capabilities +export const CAPABILITIES = `${BASE_PATH}/capabilities`; diff --git a/x-pack/plugins/elastic_assistant/package.json b/x-pack/plugins/elastic_assistant/package.json index b6f19d2ec7a3c23..37a0b9abb02b179 100644 --- a/x-pack/plugins/elastic_assistant/package.json +++ b/x-pack/plugins/elastic_assistant/package.json @@ -5,6 +5,9 @@ "private": true, "license": "Elastic License 2.0", "scripts": { - "evaluate-model": "node ./scripts/model_evaluator" + "evaluate-model": "node ./scripts/model_evaluator", + "openapi:generate": "node scripts/openapi/generate", + "openapi:generate:debug": "node --inspect-brk scripts/openapi/generate", + "openapi:bundle": "node scripts/openapi/bundle" } } \ No newline at end of file diff --git a/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js b/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js new file mode 100644 index 000000000000000..15d431a947582f7 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js @@ -0,0 +1,18 @@ +/* + * 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. + */ + +require('../../../../../src/setup_node_env'); +const { bundle } = require('@kbn/openapi-bundler'); +const { resolve } = require('path'); + +const ELASTIC_ASSISTANT_ROOT = resolve(__dirname, '../..'); + +bundle({ + rootDir: ELASTIC_ASSISTANT_ROOT, + sourceGlob: './server/schemas/**/*.schema.yaml', + outputFilePath: './target/openapi/elastic_assistant.bundled.schema.yaml', +}); diff --git a/x-pack/plugins/elastic_assistant/scripts/openapi/generate.js b/x-pack/plugins/elastic_assistant/scripts/openapi/generate.js new file mode 100644 index 000000000000000..2863fb25db58080 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/scripts/openapi/generate.js @@ -0,0 +1,18 @@ +/* + * 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. + */ + +require('../../../../../src/setup_node_env'); +const { generate } = require('@kbn/openapi-generator'); +const { resolve } = require('path'); + +const ELASTIC_ASSISTANT_ROOT = resolve(__dirname, '../..'); + +generate({ + rootDir: ELASTIC_ASSISTANT_ROOT, + sourceGlob: './server/schemas/**/*.schema.yaml', + templateName: 'zod_operation_schema', +}); diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts index deb85a88215cfe2..930374567533bfa 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts @@ -5,7 +5,11 @@ * 2.0. */ import { httpServerMock } from '@kbn/core/server/mocks'; -import { KNOWLEDGE_BASE } from '../../common/constants'; +import { CAPABILITIES, EVALUATE, KNOWLEDGE_BASE } from '../../common/constants'; +import { + PostEvaluateBodyInputs, + PostEvaluatePathQueryInputs, +} from '../schemas/evaluate/post_evaluate'; export const requestMock = { create: httpServerMock.createKibanaRequest, @@ -31,3 +35,23 @@ export const getDeleteKnowledgeBaseRequest = (resource?: string) => path: KNOWLEDGE_BASE, query: { resource }, }); + +export const getGetCapabilitiesRequest = () => + requestMock.create({ + method: 'get', + path: CAPABILITIES, + }); + +export const getPostEvaluateRequest = ({ + body, + query, +}: { + body: PostEvaluateBodyInputs; + query: PostEvaluatePathQueryInputs; +}) => + requestMock.create({ + body, + method: 'post', + path: EVALUATE, + query, + }); 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 750c13debb3fd52..3273cbf50b83d28 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(), + getRegisteredFeatures: jest.fn(), getRegisteredTools: jest.fn(), logger: loggingSystemMock.createLogger(), telemetry: coreMock.createSetup().analytics, @@ -74,6 +75,7 @@ const createElasticAssistantRequestContextMock = ( ): jest.Mocked => { return { actions: clients.elasticAssistant.actions as unknown as ActionsPluginStart, + getRegisteredFeatures: jest.fn(), getRegisteredTools: jest.fn(), logger: clients.elasticAssistant.logger, telemetry: clients.elasticAssistant.telemetry, diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/server.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/server.ts index 7ac44e1beedf1df..f08e66d1b5e80f0 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/server.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/server.ts @@ -5,35 +5,59 @@ * 2.0. */ import { httpServiceMock } from '@kbn/core/server/mocks'; -import type { RequestHandler, RouteConfig, KibanaRequest } from '@kbn/core/server'; -import type { RequestHandlerContext } from '@kbn/core-http-request-handler-context-server'; +import type { + RequestHandler, + RouteConfig, + KibanaRequest, + RequestHandlerContext, +} from '@kbn/core/server'; import { requestMock } from './request'; import { responseMock as responseFactoryMock } from './response'; import { requestContextMock } from './request_context'; import { responseAdapter } from './test_adapters'; +import type { RegisteredVersionedRoute } from '@kbn/core-http-router-server-mocks'; interface Route { - config: RouteConfig; + validate: RouteConfig< + unknown, + unknown, + unknown, + 'get' | 'post' | 'delete' | 'patch' | 'put' + >['validate']; handler: RequestHandler; } -const getRoute = (routerMock: MockServer['router']): Route => { - const routeCalls = [ - ...routerMock.get.mock.calls, - ...routerMock.post.mock.calls, - ...routerMock.put.mock.calls, - ...routerMock.patch.mock.calls, - ...routerMock.delete.mock.calls, - ]; - - const [route] = routeCalls; - if (!route) { - throw new Error('No route registered!'); +const HTTP_METHODS = ['get', 'post', 'put', 'patch', 'delete'] as const; + +const getClassicRoute = (routerMock: MockServer['router']): Route | undefined => { + const method = HTTP_METHODS.find((m) => routerMock[m].mock.calls.length > 0); + if (!method) { + return undefined; } - const [config, handler] = route; - return { config, handler }; + const [config, handler] = routerMock[method].mock.calls[0]; + return { validate: config.validate, handler }; +}; + +const getVersionedRoute = (router: MockServer['router']): Route => { + const method = HTTP_METHODS.find((m) => router.versioned[m].mock.calls.length > 0); + if (!method) { + throw new Error('No route registered!'); + } + const config = router.versioned[method].mock.calls[0][0]; + const routePath = config.path; + + const route: RegisteredVersionedRoute = router.versioned.getRoute(method, routePath); + const firstVersion = Object.values(route.versions)[0]; + + return { + validate: + firstVersion.config.validate === false + ? false + : firstVersion.config.validate.request || false, + handler: firstVersion.handler, + }; }; const buildResultMock = () => ({ ok: jest.fn((x) => x), badRequest: jest.fn((x) => x) }); @@ -63,7 +87,7 @@ class MockServer { } private getRoute(): Route { - return getRoute(this.router); + return getClassicRoute(this.router) ?? getVersionedRoute(this.router); } // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -72,7 +96,7 @@ class MockServer { } private validateRequest(request: KibanaRequest): KibanaRequest { - const validations = this.getRoute().config.validate; + const validations = this.getRoute().validate; if (!validations) { return request; } @@ -88,6 +112,7 @@ class MockServer { return validatedRequest; } } + const createMockServer = () => new MockServer(); export const serverMock = { diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/test_adapters.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/test_adapters.ts index a3fde9d64212f53..7b163138c3c2c8b 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/test_adapters.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/test_adapters.ts @@ -43,6 +43,8 @@ const buildResponses = (method: Method, calls: MockCall[]): ResponseCall[] => { status: call.statusCode, body: call.body, })); + case 'notFound': + return calls.map(() => ({ status: 404, body: undefined })); default: throw new Error(`Encountered unexpected call to response.${method}`); } diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index f142df46beb8b19..bbc2c63381fc927 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -18,6 +18,7 @@ import { } from '@kbn/core/server'; import { once } from 'lodash'; +import { AssistantFeatures } from '@kbn/elastic-assistant-common'; import { events } from './lib/telemetry/event_based_telemetry'; import { AssistantTool, @@ -36,11 +37,17 @@ import { postEvaluateRoute, postKnowledgeBaseRoute, } from './routes'; -import { appContextService, GetRegisteredTools } from './services/app_context'; +import { + appContextService, + GetRegisteredFeatures, + GetRegisteredTools, +} from './services/app_context'; +import { getCapabilitiesRoute } from './routes/capabilities/get_capabilities_route'; interface CreateRouteHandlerContextParams { core: CoreSetup; logger: Logger; + getRegisteredFeatures: GetRegisteredFeatures; getRegisteredTools: GetRegisteredTools; telemetry: AnalyticsServiceSetup; } @@ -63,6 +70,7 @@ export class ElasticAssistantPlugin private createRouteHandlerContext = ({ core, logger, + getRegisteredFeatures, getRegisteredTools, telemetry, }: CreateRouteHandlerContextParams): IContextProvider< @@ -74,6 +82,7 @@ export class ElasticAssistantPlugin return { actions: pluginsStart.actions, + getRegisteredFeatures, getRegisteredTools, logger, telemetry, @@ -89,6 +98,9 @@ export class ElasticAssistantPlugin this.createRouteHandlerContext({ core: core as CoreSetup, logger: this.logger, + getRegisteredFeatures: (pluginName: string) => { + return appContextService.getRegisteredFeatures(pluginName); + }, getRegisteredTools: (pluginName: string) => { return appContextService.getRegisteredTools(pluginName); }, @@ -112,40 +124,37 @@ export class ElasticAssistantPlugin postActionsConnectorExecuteRoute(router, getElserId); // Evaluate postEvaluateRoute(router, getElserId); + // Capabilities + getCapabilitiesRoute(router); return { actions: plugins.actions, + getRegisteredFeatures: (pluginName: string) => { + return appContextService.getRegisteredFeatures(pluginName); + }, getRegisteredTools: (pluginName: string) => { return appContextService.getRegisteredTools(pluginName); }, }; } - public start(core: CoreStart, plugins: ElasticAssistantPluginStartDependencies) { + public start( + core: CoreStart, + plugins: ElasticAssistantPluginStartDependencies + ): ElasticAssistantPluginStart { this.logger.debug('elasticAssistant: Started'); appContextService.start({ logger: this.logger }); return { - /** - * Actions plugin start contract - */ actions: plugins.actions, - - /** - * Get the registered tools for a given plugin name. - * @param pluginName - */ + getRegisteredFeatures: (pluginName: string) => { + return appContextService.getRegisteredFeatures(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 - */ + registerFeatures: (pluginName: string, features: Partial) => { + return appContextService.registerFeatures(pluginName, features); + }, registerTools: (pluginName: string, tools: AssistantTool[]) => { return appContextService.registerTools(pluginName, tools); }, diff --git a/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.test.ts new file mode 100644 index 000000000000000..b0437bbcb7209fb --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.test.ts @@ -0,0 +1,47 @@ +/* + * 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 { getCapabilitiesRoute } from './get_capabilities_route'; +import { serverMock } from '../../__mocks__/server'; +import { requestContextMock } from '../../__mocks__/request_context'; +import { getGetCapabilitiesRequest } from '../../__mocks__/request'; +import { getPluginNameFromRequest } from '../helpers'; + +jest.mock('../helpers'); + +describe('Get Capabilities Route', () => { + let server: ReturnType; + let { context } = requestContextMock.createTools(); + + beforeEach(() => { + server = serverMock.create(); + ({ context } = requestContextMock.createTools()); + + getCapabilitiesRoute(server.router); + }); + + describe('Status codes', () => { + it('returns 200 with capabilities', async () => { + const response = await server.inject( + getGetCapabilitiesRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + }); + + it('returns 500 if an error is thrown in fetching capabilities', async () => { + (getPluginNameFromRequest as jest.Mock).mockImplementation(() => { + throw new Error('Mocked error'); + }); + const response = await server.inject( + getGetCapabilitiesRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.ts b/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.ts new file mode 100644 index 000000000000000..46fc486b82a4854 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.ts @@ -0,0 +1,60 @@ +/* + * 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 { IKibanaResponse, IRouter } from '@kbn/core/server'; +import { transformError } from '@kbn/securitysolution-es-utils'; + +import { CAPABILITIES } from '../../../common/constants'; +import { ElasticAssistantRequestHandlerContext } from '../../types'; + +import { GetCapabilitiesResponse } from '../../schemas/capabilities/get_capabilities_route.gen'; +import { buildResponse } from '../../lib/build_response'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; + +/** + * Get the assistant capabilities for the requesting plugin + * + * @param router IRouter for registering routes + */ +export const getCapabilitiesRoute = (router: IRouter) => { + router.versioned + .get({ + access: 'internal', + path: CAPABILITIES, + options: { + tags: ['access:elasticAssistant'], + }, + }) + .addVersion( + { + version: '1', + validate: {}, + }, + async (context, request, response): Promise> => { + const resp = buildResponse(response); + const assistantContext = await context.elasticAssistant; + const logger = assistantContext.logger; + + try { + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName); + + return response.ok({ body: registeredFeatures }); + } catch (err) { + const error = transformError(err); + return resp.error({ + body: error.message, + statusCode: error.statusCode, + }); + } + } + ); +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts new file mode 100644 index 000000000000000..3ae64f1d89f3b39 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts @@ -0,0 +1,59 @@ +/* + * 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 { postEvaluateRoute } from './post_evaluate'; +import { serverMock } from '../../__mocks__/server'; +import { requestContextMock } from '../../__mocks__/request_context'; +import { getPostEvaluateRequest } from '../../__mocks__/request'; +import { + PostEvaluateBodyInputs, + PostEvaluatePathQueryInputs, +} from '../../schemas/evaluate/post_evaluate'; + +const defaultBody: PostEvaluateBodyInputs = { + dataset: undefined, + evalPrompt: undefined, +}; + +const defaultQueryParams: PostEvaluatePathQueryInputs = { + agents: 'agents', + datasetName: undefined, + evaluationType: undefined, + evalModel: undefined, + models: 'models', + outputIndex: '.kibana-elastic-ai-assistant-', + projectName: undefined, + runName: undefined, +}; + +describe('Post Evaluate Route', () => { + let server: ReturnType; + let { context } = requestContextMock.createTools(); + const mockGetElser = jest.fn().mockResolvedValue('.elser_model_2'); + + beforeEach(() => { + server = serverMock.create(); + ({ context } = requestContextMock.createTools()); + + postEvaluateRoute(server.router, mockGetElser); + }); + + describe('Capabilities', () => { + it('returns a 404 if evaluate feature is not registered', async () => { + context.elasticAssistant.getRegisteredFeatures.mockReturnValueOnce({ + assistantModelEvaluation: false, + assistantStreamingEnabled: false, + }); + + const response = await server.inject( + getPostEvaluateRequest({ body: defaultBody, query: defaultQueryParams }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(404); + }); + }); +}); 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 ff3291f6b703f9c..aa041175b75ee07 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 @@ -29,6 +29,7 @@ import { } from '../../lib/model_evaluator/output_index/utils'; import { fetchLangSmithDataset, getConnectorName, getLangSmithTracer, getLlmType } from './utils'; import { RequestBody } from '../../lib/langchain/types'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; /** * To support additional Agent Executors from the UI, add them to this map @@ -53,11 +54,22 @@ export const postEvaluateRoute = ( query: buildRouteValidation(PostEvaluatePathQuery), }, }, - // TODO: Limit route based on experimental feature async (context, request, response) => { const assistantContext = await context.elasticAssistant; const logger = assistantContext.logger; const telemetry = assistantContext.telemetry; + + // Validate evaluation feature is enabled + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName); + if (!registeredFeatures.assistantModelEvaluation) { + return response.notFound(); + } + try { const evaluationId = uuidv4(); const { diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts index 99d4493c16cca00..a418827c4829d24 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -7,10 +7,9 @@ import { KibanaRequest } from '@kbn/core-http-server'; import { Logger } from '@kbn/core/server'; -import { RequestBody } from '../lib/langchain/types'; interface GetPluginNameFromRequestParams { - request: KibanaRequest; + request: KibanaRequest; defaultPluginName: string; logger?: Logger; } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 6cc683fd4d8b83d..c61887a436267d8 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -22,6 +22,7 @@ import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } f * Get the status of the Knowledge Base index, pipeline, and resources (collection of documents) * * @param router IRouter for registering routes + * @param getElser Function to get the default Elser ID */ export const getKnowledgeBaseStatusRoute = ( router: IRouter, diff --git a/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.gen.ts b/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.gen.ts new file mode 100644 index 000000000000000..609d83fb0b93134 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.gen.ts @@ -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. + */ + +import { z } from 'zod'; + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + */ + +export type GetCapabilitiesResponse = z.infer; +export const GetCapabilitiesResponse = z.object({ + assistantModelEvaluation: z.boolean(), + assistantStreamingEnabled: z.boolean(), +}); diff --git a/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.schema.yaml b/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.schema.yaml new file mode 100644 index 000000000000000..6278d83411d1067 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.schema.yaml @@ -0,0 +1,41 @@ +openapi: 3.0.0 +info: + title: Get Capabilities API endpoint + version: '1' +paths: + /internal/elastic_assistant/capabilities: + get: + operationId: GetCapabilities + x-codegen-enabled: true + description: Get Elastic Assistant capabilities for the requesting plugin + summary: Get Elastic Assistant capabilities + tags: + - Capabilities API + responses: + '200': + description: Successful response + content: + application/json: + schema: + type: object + properties: + assistantModelEvaluation: + type: boolean + assistantStreamingEnabled: + type: boolean + required: + - assistantModelEvaluation + - assistantStreamingEnabled + '400': + description: Generic Error + content: + application/json: + schema: + type: object + properties: + statusCode: + type: number + error: + type: string + message: + type: string diff --git a/x-pack/plugins/elastic_assistant/server/schemas/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/schemas/evaluate/post_evaluate.ts index c9c0ee1f00e5194..f520bf9bf93b6b4 100644 --- a/x-pack/plugins/elastic_assistant/server/schemas/evaluate/post_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/schemas/evaluate/post_evaluate.ts @@ -35,6 +35,8 @@ export const PostEvaluatePathQuery = t.type({ runName: t.union([t.string, t.undefined]), }); +export type PostEvaluatePathQueryInputs = t.TypeOf; + export type DatasetItem = t.TypeOf; export const DatasetItem = t.type({ id: t.union([t.string, t.undefined]), 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 index 621995d3452be55..9c9c7ea0dd2fb1a 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts @@ -8,6 +8,7 @@ import { appContextService, ElasticAssistantAppContext } from './app_context'; import { loggerMock } from '@kbn/logging-mocks'; import { AssistantTool } from '../types'; +import { AssistantFeatures, defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; // Mock Logger const mockLogger = loggerMock.create(); @@ -48,6 +49,19 @@ describe('AppContextService', () => { expect(appContextService.getRegisteredTools('super').length).toBe(0); }); + + it('should return default registered features when stopped ', () => { + appContextService.start(mockAppContext); + appContextService.registerFeatures('super', { + assistantModelEvaluation: true, + assistantStreamingEnabled: true, + }); + appContextService.stop(); + + expect(appContextService.getRegisteredFeatures('super')).toEqual( + expect.objectContaining(defaultAssistantFeatures) + ); + }); }); describe('registering tools', () => { @@ -84,4 +98,81 @@ describe('AppContextService', () => { expect(appContextService.getRegisteredTools(pluginName).length).toEqual(1); }); }); + + describe('registering features', () => { + it('should register and get features for a single plugin', () => { + const pluginName = 'pluginName'; + const features: AssistantFeatures = { + assistantModelEvaluation: true, + assistantStreamingEnabled: true, + }; + + appContextService.start(mockAppContext); + appContextService.registerFeatures(pluginName, features); + + // Check if getRegisteredFeatures returns the correct tools + const retrievedFeatures = appContextService.getRegisteredFeatures(pluginName); + expect(retrievedFeatures).toEqual(features); + }); + + it('should register and get features for multiple plugins', () => { + const pluginOne = 'plugin1'; + const featuresOne: AssistantFeatures = { + assistantModelEvaluation: true, + assistantStreamingEnabled: false, + }; + const pluginTwo = 'plugin2'; + const featuresTwo: AssistantFeatures = { + assistantModelEvaluation: false, + assistantStreamingEnabled: true, + }; + + appContextService.start(mockAppContext); + appContextService.registerFeatures(pluginOne, featuresOne); + appContextService.registerFeatures(pluginTwo, featuresTwo); + + expect(appContextService.getRegisteredFeatures(pluginOne)).toEqual(featuresOne); + expect(appContextService.getRegisteredFeatures(pluginTwo)).toEqual(featuresTwo); + }); + + it('should update features if registered again', () => { + const pluginName = 'pluginName'; + const featuresOne: AssistantFeatures = { + assistantModelEvaluation: true, + assistantStreamingEnabled: false, + }; + const featuresTwo: AssistantFeatures = { + assistantModelEvaluation: false, + assistantStreamingEnabled: true, + }; + + appContextService.start(mockAppContext); + appContextService.registerFeatures(pluginName, featuresOne); + appContextService.registerFeatures(pluginName, featuresTwo); + + expect(appContextService.getRegisteredFeatures(pluginName)).toEqual(featuresTwo); + }); + + it('should return default features if pluginName not present', () => { + appContextService.start(mockAppContext); + + expect(appContextService.getRegisteredFeatures('super')).toEqual( + expect.objectContaining(defaultAssistantFeatures) + ); + }); + + it('allows registering a subset of all available features', () => { + const pluginName = 'pluginName'; + const featuresSubset: Partial = { + assistantModelEvaluation: true, + }; + + appContextService.start(mockAppContext); + appContextService.registerFeatures(pluginName, featuresSubset); + + expect(appContextService.getRegisteredFeatures(pluginName)).toEqual( + expect.objectContaining({ ...defaultAssistantFeatures, ...featuresSubset }) + ); + }); + }); }); 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 bd7a7c0cc3203b2..cb425540635d985 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -6,11 +6,14 @@ */ import type { Logger } from '@kbn/core/server'; -import type { AssistantTool } from '../types'; +import { defaultAssistantFeatures, AssistantFeatures } from '@kbn/elastic-assistant-common'; +import { AssistantTool } from '../types'; export type PluginName = string; export type RegisteredToolsStorage = Map>; +export type RegisteredFeaturesStorage = Map; export type GetRegisteredTools = (pluginName: string) => AssistantTool[]; +export type GetRegisteredFeatures = (pluginName: string) => AssistantFeatures; export interface ElasticAssistantAppContext { logger: Logger; } @@ -23,6 +26,7 @@ export interface ElasticAssistantAppContext { class AppContextService { private logger: Logger | undefined; private registeredTools: RegisteredToolsStorage = new Map>(); + private registeredFeatures: RegisteredFeaturesStorage = new Map(); public start(appContext: ElasticAssistantAppContext) { this.logger = appContext.logger; @@ -30,6 +34,7 @@ class AppContextService { public stop() { this.registeredTools.clear(); + this.registeredFeatures.clear(); } /** @@ -44,7 +49,7 @@ class AppContextService { this.logger?.debug(`tools: ${tools.map((tool) => tool.name).join(', ')}`); if (!this.registeredTools.has(pluginName)) { - this.logger?.debug('plugin has no tools, making new set'); + this.logger?.debug('plugin has no tools, initializing...'); this.registeredTools.set(pluginName, new Set()); } tools.forEach((tool) => this.registeredTools.get(pluginName)?.add(tool)); @@ -64,6 +69,51 @@ class AppContextService { return tools; } + + /** + * Register features to be used by the Elastic Assistant + * + * @param pluginName + * @param features + */ + public registerFeatures(pluginName: string, features: Partial) { + this.logger?.debug('AppContextService:registerFeatures'); + this.logger?.debug(`pluginName: ${pluginName}`); + this.logger?.debug( + `features: ${Object.entries(features) + .map(([feature, enabled]) => `${feature}:${enabled}`) + .join(', ')}` + ); + + if (!this.registeredFeatures.has(pluginName)) { + this.logger?.debug('plugin has no features, initializing...'); + this.registeredFeatures.set(pluginName, defaultAssistantFeatures); + } + + const registeredFeatures = this.registeredFeatures.get(pluginName); + if (registeredFeatures != null) { + this.registeredFeatures.set(pluginName, { ...registeredFeatures, ...features }); + } + } + + /** + * Get the registered features + * + * @param pluginName + */ + public getRegisteredFeatures(pluginName: string): AssistantFeatures { + const features = this.registeredFeatures?.get(pluginName) ?? defaultAssistantFeatures; + + this.logger?.debug('AppContextService:getRegisteredFeatures'); + this.logger?.debug(`pluginName: ${pluginName}`); + this.logger?.debug( + `features: ${Object.entries(features) + .map(([feature, enabled]) => `${feature}:${enabled}`) + .join(', ')}` + ); + + return features; + } } 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 c45966b9b80a24e..dafb6ad6b9bb3a9 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -20,8 +20,9 @@ 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 { AssistantFeatures } from '@kbn/elastic-assistant-common'; import { RequestBody } from './lib/langchain/types'; -import type { GetRegisteredTools } from './services/app_context'; +import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; export const PLUGIN_ID = 'elasticAssistant' as const; @@ -32,15 +33,37 @@ export interface ElasticAssistantPluginSetup { /** The plugin start interface */ export interface ElasticAssistantPluginStart { + /** + * Actions plugin start contract. + */ actions: ActionsPluginStart; /** - * Register tools to be used by the elastic assistant + * Register features to be used by the elastic assistant. + * + * Note: Be sure to use the pluginName that is sent in the request headers by your plugin to ensure it is extracted + * and the correct features are available. See {@link getPluginNameFromRequest} for more details. + * + * @param pluginName Name of the plugin the features should be registered to + * @param features Partial to be registered with for the given plugin + */ + registerFeatures: (pluginName: string, features: Partial) => void; + /** + * Get the registered features for a given plugin name. + * @param pluginName Name of the plugin to get the features for + */ + getRegisteredFeatures: GetRegisteredFeatures; + /** + * Register tools to be used by the elastic assistant. + * + * Note: Be sure to use the pluginName that is sent in the request headers by your plugin to ensure it is extracted + * and the correct tools are selected. See {@link getPluginNameFromRequest} for more details. + * * @param pluginName Name of the plugin the tool should be registered to * @param tools AssistantTools to be registered with for the given plugin */ registerTools: (pluginName: string, tools: AssistantTool[]) => void; /** - * Get the registered tools + * Get the registered tools for a given plugin name. * @param pluginName Name of the plugin to get the tools for */ getRegisteredTools: GetRegisteredTools; @@ -56,6 +79,7 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { actions: ActionsPluginStart; + getRegisteredFeatures: GetRegisteredFeatures; getRegisteredTools: GetRegisteredTools; logger: Logger; telemetry: AnalyticsServiceSetup; diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 2fd22f015ad8da9..dfca7893b203659 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -17,7 +17,6 @@ "@kbn/core", "@kbn/core-http-server", "@kbn/licensing-plugin", - "@kbn/core-http-request-handler-context-server", "@kbn/securitysolution-es-utils", "@kbn/securitysolution-io-ts-utils", "@kbn/actions-plugin", @@ -34,6 +33,8 @@ "@kbn/ml-plugin", "@kbn/apm-utils", "@kbn/core-analytics-server", + "@kbn/elastic-assistant-common", + "@kbn/core-http-router-server-mocks", ], "exclude": [ "target/**/*", diff --git a/x-pack/plugins/security_solution/public/assistant/provider.tsx b/x-pack/plugins/security_solution/public/assistant/provider.tsx index a9f9e14a8d3e0ba..43ae8eb69612886 100644 --- a/x-pack/plugins/security_solution/public/assistant/provider.tsx +++ b/x-pack/plugins/security_solution/public/assistant/provider.tsx @@ -22,7 +22,6 @@ import { useAnonymizationStore } from './use_anonymization_store'; import { useAssistantAvailability } from './use_assistant_availability'; import { APP_ID } from '../../common/constants'; import { useAppToasts } from '../common/hooks/use_app_toasts'; -import { useIsExperimentalFeatureEnabled } from '../common/hooks/use_experimental_features'; import { useSignalIndex } from '../detections/containers/detection_engine/alerts/use_signal_index'; const ASSISTANT_TITLE = i18n.translate('xpack.securitySolution.assistant.title', { @@ -39,8 +38,6 @@ export const AssistantProvider: React.FC = ({ children }) => { docLinks: { ELASTIC_WEBSITE_URL, DOC_LINK_VERSION }, } = useKibana().services; const basePath = useBasePath(); - const isModelEvaluationEnabled = useIsExperimentalFeatureEnabled('assistantModelEvaluation'); - const assistantStreamingEnabled = useIsExperimentalFeatureEnabled('assistantStreamingEnabled'); const { conversations, setConversations } = useConversationStore(); const getInitialConversation = useCallback(() => { @@ -78,8 +75,6 @@ export const AssistantProvider: React.FC = ({ children }) => { getInitialConversations={getInitialConversation} getComments={getComments} http={http} - assistantStreamingEnabled={assistantStreamingEnabled} - modelEvaluatorEnabled={isModelEvaluationEnabled} nameSpace={nameSpace} setConversations={setConversations} setDefaultAllow={setDefaultAllow} diff --git a/x-pack/plugins/security_solution/public/detections/components/rules/rule_execution_status/rule_status_failed_callout.test.tsx b/x-pack/plugins/security_solution/public/detections/components/rules/rule_execution_status/rule_status_failed_callout.test.tsx index 5998073d65cdb66..a1f7b488b1d0a6a 100644 --- a/x-pack/plugins/security_solution/public/detections/components/rules/rule_execution_status/rule_status_failed_callout.test.tsx +++ b/x-pack/plugins/security_solution/public/detections/components/rules/rule_execution_status/rule_status_failed_callout.test.tsx @@ -15,6 +15,7 @@ import { AssistantProvider } from '@kbn/elastic-assistant'; import type { AssistantAvailability } from '@kbn/elastic-assistant'; import { httpServiceMock } from '@kbn/core-http-browser-mocks'; import { actionTypeRegistryMock } from '@kbn/triggers-actions-ui-plugin/public/application/action_type_registry.mock'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; jest.mock('../../../../common/lib/kibana'); @@ -32,29 +33,44 @@ const mockAssistantAvailability: AssistantAvailability = { hasConnectorsReadPrivilege: true, isAssistantEnabled: true, }; +const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + logger: { + log: jest.fn(), + warn: jest.fn(), + error: () => {}, + }, +}); + const ContextWrapper: React.FC = ({ children }) => ( - - {children} - + + + {children} + + ); describe('RuleStatusFailedCallOut', () => { diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 01154dc06f5f6e0..27bdcfe796e7666 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -514,6 +514,10 @@ export class Plugin implements ISecuritySolutionPlugin { // Assistant Tool and Feature Registration plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools()); + plugins.elasticAssistant.registerFeatures(APP_UI_ID, { + assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, + assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled, + }); if (this.lists && plugins.taskManager && plugins.fleet) { // Exceptions, Artifacts and Manifests start