diff --git a/.buildkite/scripts/steps/checks.sh b/.buildkite/scripts/steps/checks.sh index 4d9461a6a2bff4..42a0eb9ee61c6f 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 00000000000000..ba742625b421f0 --- /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 00000000000000..5a471245e04493 --- /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 00000000000000..1d404309f73e39 --- /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 f17e13a33af3d4..c64b02160d6e43 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 00000000000000..b41d7ac1445549 --- /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 00000000000000..794b89e1775f8b --- /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 00000000000000..c9e60b806d1bf3 --- /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 00000000000000..5d52a2801fb9ea --- /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 84a2ac40a6f248..a8dc5b1aa1db76 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 50a3211f74f3cd..3f3102a4ea6bff 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 b2bd63f8101aa1..175380cc5169ad 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 100aebf3952875..5634b0b1881bdb 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 b6f19d2ec7a3c2..37a0b9abb02b17 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 00000000000000..15d431a947582f --- /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 00000000000000..2863fb25db5808 --- /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 deb85a88215cfe..930374567533bf 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 750c13debb3fd5..3273cbf50b83d2 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 7ac44e1beedf1d..f08e66d1b5e80f 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 a3fde9d64212f5..7b163138c3c2c8 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 f142df46beb8b1..bbc2c63381fc92 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 00000000000000..b0437bbcb7209f --- /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 00000000000000..46fc486b82a485 --- /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 00000000000000..3ae64f1d89f3b3 --- /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 ff3291f6b703f9..aa041175b75ee0 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 99d4493c16cca0..a418827c4829d2 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 6cc683fd4d8b83..c61887a436267d 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 00000000000000..609d83fb0b9313 --- /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 00000000000000..6278d83411d106 --- /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 c9c0ee1f00e519..f520bf9bf93b6b 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 621995d3452be5..9c9c7ea0dd2fb1 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 bd7a7c0cc3203b..cb425540635d98 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 c45966b9b80a24..dafb6ad6b9bb3a 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 2fd22f015ad8da..dfca7893b20365 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 a9f9e14a8d3e0b..43ae8eb6961288 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 5998073d65cdb6..a1f7b488b1d0a6 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 01154dc06f5f6e..27bdcfe796e766 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