From cd307887e1c2f0a5fd2bdcc6a3477e07017b45ff Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 4 Jan 2024 20:37:35 -0700 Subject: [PATCH 01/12] Adds ability to register features with plugin server --- .buildkite/scripts/steps/checks.sh | 1 + .../elastic_assistant_codegen.sh | 10 ++++ .../impl/assistant/api/capabilities.tsx | 44 ++++++++++++++ .../impl/assistant/api/use_capabilities.tsx | 57 +++++++++++++++++++ .../impl/assistant_context/index.test.tsx | 56 +++--------------- .../impl/assistant_context/index.tsx | 14 +++-- .../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_context.ts | 2 + .../elastic_assistant/server/plugin.ts | 39 ++++++++++++- .../capabilities/get_capabilities_route.ts | 55 ++++++++++++++++++ .../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 | 27 +++++++++ .../server/services/app_context.ts | 53 ++++++++++++++++- .../plugins/elastic_assistant/server/types.ts | 34 ++++++++++- .../public/assistant/provider.tsx | 5 -- .../security_solution/server/plugin.ts | 4 ++ 22 files changed, 416 insertions(+), 66 deletions(-) create mode 100755 .buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities.tsx create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/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.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 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/impl/assistant/api/capabilities.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities.tsx new file mode 100644 index 00000000000000..6dd514ffa6e8cd --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/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 type { AssistantFeatures } from '@kbn/elastic-assistant-plugin/server/types'; + +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/use_capabilities.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/use_capabilities.tsx new file mode 100644 index 00000000000000..24847808a06245 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/use_capabilities.tsx @@ -0,0 +1,57 @@ +/* + * 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( + CAPABILITIES_QUERY_KEY, + 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..035ca242cd84c9 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 @@ -37,6 +37,7 @@ import { } from './constants'; import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings'; import { AssistantAvailability, AssistantTelemetry } from './types'; +import { useCapabilities } from '../assistant/api/use_capabilities'; export interface ShowAssistantOverlayProps { showOverlay: boolean; @@ -53,7 +54,6 @@ export interface AssistantProviderProps { actionTypeRegistry: ActionTypeRegistryContract; alertsIndexPattern?: string; assistantAvailability: AssistantAvailability; - assistantStreamingEnabled?: boolean; assistantTelemetry?: AssistantTelemetry; augmentMessageCodeBlocks: (currentConversation: Conversation) => CodeBlockDetails[][]; baseAllow: string[]; @@ -87,7 +87,6 @@ export interface AssistantProviderProps { }) => EuiCommentProps[]; http: HttpSetup; getInitialConversations: () => Record; - modelEvaluatorEnabled?: boolean; nameSpace?: string; setConversations: React.Dispatch>>; setDefaultAllow: React.Dispatch>; @@ -163,7 +162,6 @@ export const AssistantProvider: React.FC = ({ actionTypeRegistry, alertsIndexPattern, assistantAvailability, - assistantStreamingEnabled = false, assistantTelemetry, augmentMessageCodeBlocks, baseAllow, @@ -179,7 +177,6 @@ export const AssistantProvider: React.FC = ({ getComments, http, getInitialConversations, - modelEvaluatorEnabled = false, nameSpace = DEFAULT_ASSISTANT_NAMESPACE, setConversations, setDefaultAllow, @@ -298,6 +295,15 @@ export const AssistantProvider: React.FC = ({ [localStorageLastConversationId] ); + // Fetch assistant capabilities + // TODO: Where are constants going these days, server/common, or kbn-elastic-assistant-common pkg? + const { data: capabilities } = useCapabilities({ http, toasts }); + const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = + capabilities ?? { + assistantModelEvaluation: false, + assistantStreamingEnabled: false, + }; + const value = useMemo( () => ({ actionTypeRegistry, 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..89a2307f00fdc7 --- /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: './common/api/**/*.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..b5d813648d7d94 --- /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: './**/*.schema.yaml', + templateName: 'zod_operation_schema', +}); 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/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index f142df46beb8b1..00c72bb3aaec63 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -20,6 +20,7 @@ import { once } from 'lodash'; import { events } from './lib/telemetry/event_based_telemetry'; import { + AssistantFeatures, AssistantTool, ElasticAssistantPluginSetup, ElasticAssistantPluginSetupDependencies, @@ -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,8 +124,13 @@ 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); }, @@ -130,6 +147,26 @@ export class ElasticAssistantPlugin */ actions: plugins.actions, + /** + * Get the registered features for a given plugin name. + * @param pluginName + */ + getRegisteredFeatures: (pluginName: string) => { + return appContextService.getRegisteredFeatures(pluginName); + }, + + /** + * Register features 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 features + */ + registerFeatures: (pluginName: string, features: AssistantFeatures) => { + return appContextService.registerFeatures(pluginName, features); + }, + /** * Get the registered tools for a given plugin name. * @param pluginName 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..f726cb10bd42ff --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.ts @@ -0,0 +1,55 @@ +/* + * 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'; + +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.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..4f6e5d2fd89d4b --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.schema.yaml @@ -0,0 +1,27 @@ +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. + tags: + - Capabilities API + responses: + '200': + description: Successful response + content: + application/json: + schema: + type: object + properties: + assistantModelEvaluation: + type: boolean + assistantStreamingEnabled: + type: boolean + required: + - assistantModelEvaluation + - assistantStreamingEnabled 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..3214bf678fa572 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,13 @@ */ import type { Logger } from '@kbn/core/server'; -import type { AssistantTool } from '../types'; +import { assistantFeatures, AssistantFeatures, 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 +25,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 +33,7 @@ class AppContextService { public stop() { this.registeredTools.clear(); + this.registeredFeatures.clear(); } /** @@ -44,7 +48,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 +68,51 @@ class AppContextService { return tools; } + + /** + * Register features to be used by the Elastic Assistant + * + * @param pluginName + * @param features + */ + public registerFeatures(pluginName: string, features: AssistantFeatures) { + 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, assistantFeatures); + } + + 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) ?? assistantFeatures; + + 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..cb1eed5f941bd7 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -21,7 +21,7 @@ import { Tool } from 'langchain/dist/tools/base'; import { RetrievalQAChain } from 'langchain/chains'; import { ElasticsearchClient } from '@kbn/core/server'; 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; @@ -34,7 +34,26 @@ export interface ElasticAssistantPluginSetup { export interface ElasticAssistantPluginStart { 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 AssistantFeatures to be registered with for the given plugin + */ + registerFeatures: (pluginName: string, features: AssistantFeatures) => void; + /** + * Get the registered features + * @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 */ @@ -56,6 +75,7 @@ export interface ElasticAssistantPluginStartDependencies { export interface ElasticAssistantApiRequestHandlerContext { actions: ActionsPluginStart; + getRegisteredFeatures: GetRegisteredFeatures; getRegisteredTools: GetRegisteredTools; logger: Logger; telemetry: AnalyticsServiceSetup; @@ -99,3 +119,13 @@ export interface AssistantToolParams { request: KibanaRequest; size?: number; } + +/** + * 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, +}); 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/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 53141fc1751ceb..3afddd82bddcc8 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -513,6 +513,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 From 4728d2dc9cd642e05f33ecc507b15f5917e6f0af Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Fri, 5 Jan 2024 15:17:42 +0000 Subject: [PATCH 02/12] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/packages/kbn-elastic-assistant/tsconfig.json | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/packages/kbn-elastic-assistant/tsconfig.json b/x-pack/packages/kbn-elastic-assistant/tsconfig.json index eea97cfc917dc1..130e86915e144e 100644 --- a/x-pack/packages/kbn-elastic-assistant/tsconfig.json +++ b/x-pack/packages/kbn-elastic-assistant/tsconfig.json @@ -30,5 +30,6 @@ "@kbn/ui-theme", "@kbn/core-doc-links-browser", "@kbn/core", + "@kbn/elastic-assistant-plugin", ] } From f78a0bf5a3a7464f69b3dde59186e30a923c96ec Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 5 Jan 2024 09:46:19 -0700 Subject: [PATCH 03/12] Fixes tests and moves capabilities to common package --- .../impl/capabilities/index.ts | 19 +++++ .../kbn-elastic-assistant-common/index.ts | 3 + .../impl/assistant/api/capabilities.tsx | 2 +- .../kbn-elastic-assistant/tsconfig.json | 1 - .../mock/test_providers/test_providers.tsx | 69 +++++++++++-------- .../elastic_assistant/server/plugin.ts | 2 +- .../server/services/app_context.ts | 3 +- .../plugins/elastic_assistant/server/types.ts | 11 +-- .../rule_status_failed_callout.test.tsx | 60 ++++++++++------ 9 files changed, 107 insertions(+), 63 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts 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..90251c3d307658 --- /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 assistantFeatures]: boolean }; + +/** + * Default features available to the elastic assistant + */ +export const assistantFeatures = 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..2bc8987c1f7bf5 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 { assistantFeatures } 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.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities.tsx index 6dd514ffa6e8cd..794b89e1775f8b 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities.tsx @@ -6,7 +6,7 @@ */ import { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser'; -import type { AssistantFeatures } from '@kbn/elastic-assistant-plugin/server/types'; +import { AssistantFeatures } from '@kbn/elastic-assistant-common'; export interface GetCapabilitiesParams { http: HttpSetup; diff --git a/x-pack/packages/kbn-elastic-assistant/tsconfig.json b/x-pack/packages/kbn-elastic-assistant/tsconfig.json index 130e86915e144e..eea97cfc917dc1 100644 --- a/x-pack/packages/kbn-elastic-assistant/tsconfig.json +++ b/x-pack/packages/kbn-elastic-assistant/tsconfig.json @@ -30,6 +30,5 @@ "@kbn/ui-theme", "@kbn/core-doc-links-browser", "@kbn/core", - "@kbn/elastic-assistant-plugin", ] } 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/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 00c72bb3aaec63..f1c0defe987a78 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -18,9 +18,9 @@ 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 { - AssistantFeatures, AssistantTool, ElasticAssistantPluginSetup, ElasticAssistantPluginSetupDependencies, 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 3214bf678fa572..13d9cad945ab68 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -6,7 +6,8 @@ */ import type { Logger } from '@kbn/core/server'; -import { assistantFeatures, AssistantFeatures, AssistantTool } from '../types'; +import { assistantFeatures, AssistantFeatures } from '@kbn/elastic-assistant-common'; +import { AssistantTool } from '../types'; export type PluginName = string; export type RegisteredToolsStorage = Map>; diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index cb1eed5f941bd7..f856afcd5e5066 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -20,6 +20,7 @@ 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 { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; @@ -119,13 +120,3 @@ export interface AssistantToolParams { request: KibanaRequest; size?: number; } - -/** - * 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, -}); 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', () => { From 42ab68d7b49e1686c89c5dbecd8495085d7b9a0a Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 5 Jan 2024 11:38:57 -0700 Subject: [PATCH 04/12] Use default features in assistant context --- .../kbn-elastic-assistant/impl/assistant_context/index.tsx | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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 035ca242cd84c9..08bd5a919eea84 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 { assistantFeatures } from '@kbn/elastic-assistant-common'; import { WELCOME_CONVERSATION_TITLE } from '../assistant/use_conversation/translations'; import { updatePromptContexts } from './helpers'; import type { @@ -296,13 +297,9 @@ export const AssistantProvider: React.FC = ({ ); // Fetch assistant capabilities - // TODO: Where are constants going these days, server/common, or kbn-elastic-assistant-common pkg? const { data: capabilities } = useCapabilities({ http, toasts }); const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = - capabilities ?? { - assistantModelEvaluation: false, - assistantStreamingEnabled: false, - }; + capabilities ?? assistantFeatures; const value = useMemo( () => ({ From e55e6768068fb2296634eefe16cec11e116583f4 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Fri, 5 Jan 2024 18:46:33 +0000 Subject: [PATCH 05/12] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/elastic_assistant/tsconfig.json | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 2fd22f015ad8da..bcec58e760b2b1 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -34,6 +34,7 @@ "@kbn/ml-plugin", "@kbn/apm-utils", "@kbn/core-analytics-server", + "@kbn/elastic-assistant-common", ], "exclude": [ "target/**/*", From 2eba0bccb8aa32489a0a0f168f6c71218ee19b34 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 5 Jan 2024 16:15:12 -0700 Subject: [PATCH 06/12] Fixes OpenAPI generation and bundling --- .../elastic_assistant/scripts/openapi/bundle.js | 2 +- .../scripts/openapi/generate.js | 2 +- .../get_capabilities_route.schema.yaml | 16 +++++++++++++++- 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js b/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js index 89a2307f00fdc7..15d431a947582f 100644 --- a/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js +++ b/x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js @@ -13,6 +13,6 @@ const ELASTIC_ASSISTANT_ROOT = resolve(__dirname, '../..'); bundle({ rootDir: ELASTIC_ASSISTANT_ROOT, - sourceGlob: './common/api/**/*.schema.yaml', + 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 index b5d813648d7d94..2863fb25db5808 100644 --- a/x-pack/plugins/elastic_assistant/scripts/openapi/generate.js +++ b/x-pack/plugins/elastic_assistant/scripts/openapi/generate.js @@ -13,6 +13,6 @@ const ELASTIC_ASSISTANT_ROOT = resolve(__dirname, '../..'); generate({ rootDir: ELASTIC_ASSISTANT_ROOT, - sourceGlob: './**/*.schema.yaml', + sourceGlob: './server/schemas/**/*.schema.yaml', templateName: 'zod_operation_schema', }); 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 index 4f6e5d2fd89d4b..6278d83411d106 100644 --- 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 @@ -7,7 +7,8 @@ paths: get: operationId: GetCapabilities x-codegen-enabled: true - description: Get Elastic Assistant capabilities for the requesting plugin. + description: Get Elastic Assistant capabilities for the requesting plugin + summary: Get Elastic Assistant capabilities tags: - Capabilities API responses: @@ -25,3 +26,16 @@ paths: required: - assistantModelEvaluation - assistantStreamingEnabled + '400': + description: Generic Error + content: + application/json: + schema: + type: object + properties: + statusCode: + type: number + error: + type: string + message: + type: string From 064ccfe3d042c760d3c74089a459fd7df4056122 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 5 Jan 2024 16:55:49 -0700 Subject: [PATCH 07/12] Adds tests for feature registration --- .../server/services/app_context.test.ts | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) 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..ae90c44d0250f7 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, assistantFeatures } 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(assistantFeatures) + ); + }); }); describe('registering tools', () => { @@ -84,4 +98,67 @@ 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(assistantFeatures) + ); + }); + }); }); From e7ab3e9a925fbf92ffa16852014f5c3c2128d802 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 9 Jan 2024 17:23:14 -0700 Subject: [PATCH 08/12] Increases test coverage, and updates server mocks to work with versioned router --- .../api/capabilities/capabilities.test.tsx | 43 +++++++++++++ .../api/{ => capabilities}/capabilities.tsx | 0 .../capabilities/use_capabilities.test.tsx | 49 +++++++++++++++ .../{ => capabilities}/use_capabilities.tsx | 37 +++++------ .../impl/assistant_context/index.tsx | 2 +- .../server/__mocks__/request.ts | 26 +++++++- .../server/__mocks__/server.ts | 63 +++++++++++++------ .../server/__mocks__/test_adapters.ts | 2 + .../get_capabilities_route.test.ts | 47 ++++++++++++++ .../capabilities/get_capabilities_route.ts | 5 ++ .../routes/evaluate/post_evaluate.test.ts | 59 +++++++++++++++++ .../server/schemas/evaluate/post_evaluate.ts | 2 + 12 files changed, 293 insertions(+), 42 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.test.tsx rename x-pack/packages/kbn-elastic-assistant/impl/assistant/api/{ => capabilities}/capabilities.tsx (100%) create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx rename x-pack/packages/kbn-elastic-assistant/impl/assistant/api/{ => capabilities}/use_capabilities.tsx (64%) 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/evaluate/post_evaluate.test.ts 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.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx similarity index 100% rename from x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities.tsx rename to x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx 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/use_capabilities.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx similarity index 64% rename from x-pack/packages/kbn-elastic-assistant/impl/assistant/api/use_capabilities.tsx rename to x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx index 24847808a06245..5d52a2801fb9ea 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/use_capabilities.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx @@ -31,27 +31,22 @@ export const useCapabilities = ({ http, toasts, }: UseCapabilitiesParams): UseQueryResult => { - return useQuery( - CAPABILITIES_QUERY_KEY, - async ({ signal }) => { + 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', - }), - } - ); - } - }, - } - ); + 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.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 08bd5a919eea84..cb85f2b2cb6c07 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 @@ -38,7 +38,7 @@ import { } from './constants'; import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings'; import { AssistantAvailability, AssistantTelemetry } from './types'; -import { useCapabilities } from '../assistant/api/use_capabilities'; +import { useCapabilities } from '../assistant/api/capabilities/use_capabilities'; export interface ShowAssistantOverlayProps { showOverlay: boolean; 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__/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/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 index f726cb10bd42ff..46fc486b82a485 100644 --- 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 @@ -15,6 +15,11 @@ import { GetCapabilitiesResponse } from '../../schemas/capabilities/get_capabili 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({ 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/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]), From 99f535bb5573190b2ca9b8f72606806d77b1a419 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 10 Jan 2024 00:30:02 +0000 Subject: [PATCH 09/12] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/elastic_assistant/tsconfig.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index bcec58e760b2b1..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", @@ -35,6 +34,7 @@ "@kbn/apm-utils", "@kbn/core-analytics-server", "@kbn/elastic-assistant-common", + "@kbn/core-http-router-server-mocks", ], "exclude": [ "target/**/*", From 6f17d5e4de092dfc2f4322249d76bf21d4a70e2b Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 10 Jan 2024 14:58:28 -0700 Subject: [PATCH 10/12] Allows partial feature registration, and adds docs --- .../impl/capabilities/README.md | 53 +++++++++++++++++++ .../elastic_assistant/server/plugin.ts | 42 +++------------ .../server/services/app_context.test.ts | 14 +++++ .../server/services/app_context.ts | 2 +- .../plugins/elastic_assistant/server/types.ts | 11 ++-- 5 files changed, 82 insertions(+), 40 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md 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..85dc21691de976 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md @@ -0,0 +1,53 @@ +### 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 + +```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 assistantFeatures]: boolean }; + +export const assistantFeatures = 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 ?? 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`. diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index f1c0defe987a78..bbc2c63381fc92 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -137,52 +137,24 @@ export class ElasticAssistantPlugin }; } - 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 features for a given plugin name. - * @param pluginName - */ getRegisteredFeatures: (pluginName: string) => { return appContextService.getRegisteredFeatures(pluginName); }, - - /** - * Register features 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 features - */ - registerFeatures: (pluginName: string, features: AssistantFeatures) => { - return appContextService.registerFeatures(pluginName, features); - }, - - /** - * Get the registered tools for a given plugin name. - * @param 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/services/app_context.test.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts index ae90c44d0250f7..bd8ff9c56e9ea6 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 @@ -160,5 +160,19 @@ describe('AppContextService', () => { expect.objectContaining(assistantFeatures) ); }); + + 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({ ...assistantFeatures, ...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 13d9cad945ab68..e7754f30a4674f 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.ts @@ -76,7 +76,7 @@ class AppContextService { * @param pluginName * @param features */ - public registerFeatures(pluginName: string, features: AssistantFeatures) { + public registerFeatures(pluginName: string, features: Partial) { this.logger?.debug('AppContextService:registerFeatures'); this.logger?.debug(`pluginName: ${pluginName}`); this.logger?.debug( diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index f856afcd5e5066..dafb6ad6b9bb3a 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -33,6 +33,9 @@ export interface ElasticAssistantPluginSetup { /** The plugin start interface */ export interface ElasticAssistantPluginStart { + /** + * Actions plugin start contract. + */ actions: ActionsPluginStart; /** * Register features to be used by the elastic assistant. @@ -41,11 +44,11 @@ export interface ElasticAssistantPluginStart { * 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 AssistantFeatures to be registered with for the given plugin + * @param features Partial to be registered with for the given plugin */ - registerFeatures: (pluginName: string, features: AssistantFeatures) => void; + registerFeatures: (pluginName: string, features: Partial) => void; /** - * Get the registered features + * Get the registered features for a given plugin name. * @param pluginName Name of the plugin to get the features for */ getRegisteredFeatures: GetRegisteredFeatures; @@ -60,7 +63,7 @@ export interface ElasticAssistantPluginStart { */ 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; From 45cd452cbc4abd07121110215bd9a382328ab082 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 10 Jan 2024 15:00:07 -0700 Subject: [PATCH 11/12] Fix codeblock in docs --- .../kbn-elastic-assistant-common/impl/capabilities/README.md | 2 -- 1 file changed, 2 deletions(-) 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 index 85dc21691de976..f40b952fd5af22 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md @@ -8,8 +8,6 @@ Storage and accessor functions are made available via the `AppContextService`, a 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 - ```ts plugins.elasticAssistant.registerFeatures(APP_UI_ID, { assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, From 125109daf4dc8c7556729bbb10055b6add21028d Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 10 Jan 2024 16:14:50 -0700 Subject: [PATCH 12/12] Rename assistantFeatures to defaultAssistantFeatures --- .../impl/capabilities/README.md | 6 +++--- .../impl/capabilities/index.ts | 4 ++-- x-pack/packages/kbn-elastic-assistant-common/index.ts | 2 +- .../impl/assistant_context/index.tsx | 4 ++-- .../elastic_assistant/server/services/app_context.test.ts | 8 ++++---- .../elastic_assistant/server/services/app_context.ts | 6 +++--- 6 files changed, 15 insertions(+), 15 deletions(-) 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 index f40b952fd5af22..5a471245e04493 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md @@ -19,9 +19,9 @@ plugins.elasticAssistant.registerFeatures(APP_UI_ID, { 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 assistantFeatures]: boolean }; +export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: boolean }; -export const assistantFeatures = Object.freeze({ +export const defaultAssistantFeatures = Object.freeze({ assistantModelEvaluation: false, assistantStreamingEnabled: false, }); @@ -32,7 +32,7 @@ Capabilities can be fetched client side using the `useCapabilities()` hook ala: ```ts const { data: capabilities } = useCapabilities({ http, toasts }); -const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures; +const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? defaultAssistantFeatures; ``` ### Using Capabilities Server Side 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 index 90251c3d307658..1d404309f73e39 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -8,12 +8,12 @@ /** * Interface for features available to the elastic assistant */ -export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean }; +export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: boolean }; /** * Default features available to the elastic assistant */ -export const assistantFeatures = Object.freeze({ +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 2bc8987c1f7bf5..c64b02160d6e43 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/index.ts @@ -5,7 +5,7 @@ * 2.0. */ -export { assistantFeatures } from './impl/capabilities'; +export { defaultAssistantFeatures } from './impl/capabilities'; export type { AssistantFeatures } from './impl/capabilities'; export { getAnonymizedValue } from './impl/data_anonymization/get_anonymized_value'; 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 cb85f2b2cb6c07..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,7 +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 { assistantFeatures } from '@kbn/elastic-assistant-common'; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { WELCOME_CONVERSATION_TITLE } from '../assistant/use_conversation/translations'; import { updatePromptContexts } from './helpers'; import type { @@ -299,7 +299,7 @@ export const AssistantProvider: React.FC = ({ // Fetch assistant capabilities const { data: capabilities } = useCapabilities({ http, toasts }); const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = - capabilities ?? assistantFeatures; + capabilities ?? defaultAssistantFeatures; const value = useMemo( () => ({ 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 bd8ff9c56e9ea6..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,7 +8,7 @@ import { appContextService, ElasticAssistantAppContext } from './app_context'; import { loggerMock } from '@kbn/logging-mocks'; import { AssistantTool } from '../types'; -import { AssistantFeatures, assistantFeatures } from '@kbn/elastic-assistant-common'; +import { AssistantFeatures, defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; // Mock Logger const mockLogger = loggerMock.create(); @@ -59,7 +59,7 @@ describe('AppContextService', () => { appContextService.stop(); expect(appContextService.getRegisteredFeatures('super')).toEqual( - expect.objectContaining(assistantFeatures) + expect.objectContaining(defaultAssistantFeatures) ); }); }); @@ -157,7 +157,7 @@ describe('AppContextService', () => { appContextService.start(mockAppContext); expect(appContextService.getRegisteredFeatures('super')).toEqual( - expect.objectContaining(assistantFeatures) + expect.objectContaining(defaultAssistantFeatures) ); }); @@ -171,7 +171,7 @@ describe('AppContextService', () => { appContextService.registerFeatures(pluginName, featuresSubset); expect(appContextService.getRegisteredFeatures(pluginName)).toEqual( - expect.objectContaining({ ...assistantFeatures, ...featuresSubset }) + 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 e7754f30a4674f..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,7 +6,7 @@ */ import type { Logger } from '@kbn/core/server'; -import { assistantFeatures, AssistantFeatures } from '@kbn/elastic-assistant-common'; +import { defaultAssistantFeatures, AssistantFeatures } from '@kbn/elastic-assistant-common'; import { AssistantTool } from '../types'; export type PluginName = string; @@ -87,7 +87,7 @@ class AppContextService { if (!this.registeredFeatures.has(pluginName)) { this.logger?.debug('plugin has no features, initializing...'); - this.registeredFeatures.set(pluginName, assistantFeatures); + this.registeredFeatures.set(pluginName, defaultAssistantFeatures); } const registeredFeatures = this.registeredFeatures.get(pluginName); @@ -102,7 +102,7 @@ class AppContextService { * @param pluginName */ public getRegisteredFeatures(pluginName: string): AssistantFeatures { - const features = this.registeredFeatures?.get(pluginName) ?? assistantFeatures; + const features = this.registeredFeatures?.get(pluginName) ?? defaultAssistantFeatures; this.logger?.debug('AppContextService:getRegisteredFeatures'); this.logger?.debug(`pluginName: ${pluginName}`);