Skip to content

Commit

Permalink
[Security Solution] [Elastic AI Assistant] Adds support for plugin fe…
Browse files Browse the repository at this point in the history
…ature registration (elastic#174317)

## Summary

Resolves elastic#172509

Adds ability to register feature capabilities through the assistant
server so they no longer need to be plumbed through the
`ElasticAssistantProvider`, which now also makes them available server
side.

Adds new `/internal/elastic_assistant/capabilities` route and
`useCapabilities()` UI hook for fetching capabilities.

### OpenAPI Codegen

Implemented using the new OpenAPI codegen and bundle packages:
* Includes OpenAPI codegen script and CI action as detailed in:
elastic#166269
* Includes OpenAPI docs bundling script as detailed in:
elastic#171526

To run codegen/bundling locally, cd to
`x-pack/plugins/elastic_assistant/` and run any of the following
commands:

```bash
yarn openapi:generate
yarn openapi:generate:debug
yarn openapi:bundle
```

> [!NOTE]
> At the moment `yarn openapi:bundle` will output an empty bundled
schema since `get_capabilities_route` is an internal route, this is to
be expected. Also, if you don't see the file in your IDE, it's probably
because `target` directories are ignored, so you may need to manually
find/open the bundled schema at it's target location:
`/x-pack/plugins/elastic_assistant/target/openapi/elastic_assistant.bundled.schema.yaml`

### Registering Capabilities 

To register a capability on plugin start, add the following in the
consuming plugin's `start()`:

```ts
plugins.elasticAssistant.registerFeatures(APP_UI_ID, {
  assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation,
  assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled,
});
```

### Declaring Feature Capabilities
Feature capabilities are declared in
`x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`:

```ts
/**
 * Interfaces for features available to the elastic assistant
 */
export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean };

export const assistantFeatures = Object.freeze({
  assistantModelEvaluation: false,
  assistantStreamingEnabled: false,
});
```
### Using Capabilities Client Side
And can be fetched client side using the `useCapabilities()` hook ala:

```ts
// Fetch assistant capabilities
const { data: capabilities } = useCapabilities({ http, toasts });
const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures;
```

### Using Capabilities Server Side
Or server side within a route (or elsewhere) via the `assistantContext`:

```ts
const assistantContext = await context.elasticAssistant;
const pluginName = getPluginNameFromRequest({ request, logger });
const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName);
if (!registeredFeatures.assistantModelEvaluation) {
  return response.notFound();
}
```

> [!NOTE]
> Note, just as with [registering arbitrary
tools](elastic#172234), features are
registered for a specific plugin, where the plugin name that corresponds
to your application is defined in the `x-kbn-context` header of requests
made from your application, which may be different than your plugin's
registered `APP_ID`.

Perhaps this separation of concerns from one plugin to another isn't
necessary, but it was easy to add matching the behavior of registering
arbitrary tools. We can remove this granularity in favor of global
features if desired.


### Test Steps

* Verify `/internal/elastic_assistant/capabilities` route is called on
security solution page load in dev tools, and that by default the
`Evaluation` UI in setting does is not displayed and `404`'s if manually
called.
* Set the below experimental feature flag in your `kibana.dev.yml` and
observe the feature being enabled by inspecting the capabilities api
response, and that the evaluation feature becomes available:
```
xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation']
```
* Run the `yarn openapi:*` codegen scripts above and ensure they execute
as expected (code is generated/bundled)

### Checklist

- [x]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [X] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
  • Loading branch information
2 people authored and eokoneyo committed Jan 12, 2024
1 parent bdc4d24 commit a207f57
Show file tree
Hide file tree
Showing 37 changed files with 930 additions and 155 deletions.
1 change: 1 addition & 0 deletions .buildkite/scripts/steps/checks.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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`.
Original file line number Diff line number Diff line change
@@ -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,
});
3 changes: 3 additions & 0 deletions x-pack/packages/kbn-elastic-assistant-common/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
@@ -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 });
});
});
});
Original file line number Diff line number Diff line change
@@ -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<GetCapabilitiesResponse | IHttpFetchError>}
*/
export const getCapabilities = async ({
http,
signal,
}: GetCapabilitiesParams): Promise<GetCapabilitiesResponse | IHttpFetchError> => {
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;
}
};
Original file line number Diff line number Diff line change
@@ -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 }) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
};

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();
});
});
Original file line number Diff line number Diff line change
@@ -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<GetCapabilitiesResponse, IHttpFetchError> => {
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<ResponseErrorBody>) => {
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',
}),
});
}
},
});
};
Original file line number Diff line number Diff line change
Expand Up @@ -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 }) => (
<AssistantProvider
actionTypeRegistry={actionTypeRegistry}
assistantAvailability={mockAssistantAvailability}
assistantStreamingEnabled
augmentMessageCodeBlocks={jest.fn()}
baseAllow={[]}
baseAllowReplacement={[]}
basePath={'https://localhost:5601/kbn'}
defaultAllow={[]}
defaultAllowReplacement={[]}
docLinks={{
ELASTIC_WEBSITE_URL: 'https://www.elastic.co/',
DOC_LINK_VERSION: 'current',
}}
getInitialConversations={mockGetInitialConversations}
getComments={mockGetComments}
http={mockHttp}
setConversations={jest.fn()}
setDefaultAllow={jest.fn()}
setDefaultAllowReplacement={jest.fn()}
>
{children}
</AssistantProvider>
);

describe('AssistantContext', () => {
beforeEach(() => jest.clearAllMocks());
Expand All @@ -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');
});
Expand Down
Loading

0 comments on commit a207f57

Please sign in to comment.