[Security Assistant] Adds client hooks and internal routes for managing Knowledge Base Entries (#184974)

## Summary

This PR adds client hooks and basic REST API's for accessing and
mutating Knowledge Base Entries. This is in support of @angorayc
building out the new Knowledge Base settings interface.

Change set includes:
- [X] Refactors existing KB client hooks from
`x-pack/packages/kbn-elastic-assistant/impl/knowledge_base` to be
co-located next to the API methods where we put all our other hooks:
`x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base`
- [X] Refactors existing KB API calls and associated tests out of
`kbn-elastic-assistant/impl/assistant/api/index.tsx` and into
`x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/api.tsx`
- [X] Adds new `find_knowledge_base_entries_route.schema.yaml` OAS for
the supporting
`/internal/elastic_assistant/knowledge_base/entries/_find` route
- [X] Refactors `SortOrder` out of existing OAS's into the shared
`schemas/common_attributes.schema.yaml`

### Client Hooks & Routes
Adds new `useKnowledgeBaseEntries()` hook and corresponding
`/knowledge_base/entries/_find` route for returning paginated KB Entries
to populate the KB table in settings. E.g.

``` ts
    const {
      assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault },
      http,
      toasts,
    } = useAssistantContext();
    const { data: kbEntries, isLoading: isLoadingEntries } = useKnowledgeBaseEntries({ http });
```


###### Sample Response
``` json
{
  "perPage": 20,
  "page": 1,
  "total": 145,
  "data": [
    {
      "timestamp": "2024-06-05T21:19:56.482Z",
      "id": "CtBF6o8BSQy1Bdxt2FHz",
      "createdAt": "2024-06-05T21:19:56.482Z",
      "createdBy": "u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0",
      "updatedAt": "2024-06-05T21:19:56.482Z",
      "updatedBy": "u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0",
      "users": [
        {
          "id": "u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0",
          "name": "elastic"
        }
      ],
      "metadata": {
        "kbResource": "security_labs",
        "source": "/Users/garrettspong/dev/kibana-main/x-pack/plugins/elastic_assistant/server/knowledge_base/security_labs/2022_elastic_global_threat_report_announcement.mdx",
        "required": false
      },
      "namespace": "default",
      "text": "[Source Content Here]",
      "vector": {
        "modelId": ".elser_model_2",
        "tokens": {
          "2": 0.06595266,
          ...
        }
      }
    },
    ...
  ]
}
```

Response is the full newly created `entry`. Same format for the entry as
above in the `_find` API, and the `KnowledgeBaseEntries` cache is
invalidated.


Adds new `useCreateKnowledgeBaseEntry()` hook and corresponding
`/knowledge_base/entries` route for creating new KB Entries

``` ts
    const entry: KnowledgeBaseEntryCreateProps = {
      metadata: {
        kbResource: 'user',
        required: true,
        source: 'user',
      },
      text: 'Useful information about the user',
    };
    const { mutate: createEntry, isLoading: isCreatingEntry } = useCreateKnowledgeBaseEntry({
      http,
    });
    await createEntry(entry);
```

Adds new `useDeleteKnowledgeBaseEntries()` hook and corresponding
`/knowledge_base/entries/_bulk_action` route for deleting existing KB
Entries. I left a TODO to plumb through `delete_by_query` so we can add
a filter bar to the table. Need to confirm if we can do pagination with
similarity search as well.

``` ts
    const { mutate: deleteEntries, isLoading: isDeletingEntries } = useDeleteKnowledgeBaseEntries({
      http,
    });
    await deleteEntries({ ids: ['YOE_CZABSQy1BdxtAGbs'] })
```

See `KnowledgeBaseEntryBulkCrudActionResponse` for response formats.
`KnowledgeBaseEntries` cache is invalidated upon delete.


### Checklist

Delete any items that are not applicable to this PR.

- [ ] ~Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md)~
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
* Feature currently behind feature flag. Documentation to be added
before flag is removed. Tracked in
https://github.com/elastic/security-docs/issues/5337
- [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
- [ ] API tests will need to be rounded out as we finalize functionality
behind the feature flag

---------

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>
This commit is contained in:
Garrett Spong 2024-06-18 19:48:07 -06:00 committed by GitHub
parent 1c7b5952b3
commit 1b872fbf9d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
40 changed files with 1161 additions and 288 deletions

View file

@ -27,5 +27,6 @@ export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_FIND = `${ELASTIC_AI_
// TODO: Update existing 'status' endpoint to take resource as query param as to not conflict with 'entries'
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/knowledge_base/{resource?}`;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL = `${ELASTIC_AI_ASSISTANT_URL}/knowledge_base/entries`;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_URL}/knowledge_base/_bulk_action`;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/knowledge_base/entries`;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND = `${ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL}/_find`;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL}/_bulk_action`;

View file

@ -10,6 +10,11 @@
*/
export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: boolean };
/**
* Type for keys of the assistant features
*/
export type AssistantFeatureKey = keyof AssistantFeatures;
/**
* Default features available to the elastic assistant
*/

View file

@ -17,6 +17,7 @@
import { z } from 'zod';
import { ArrayFromString } from '@kbn/zod-helpers';
import { SortOrder } from '../common_attributes.gen';
import { AnonymizationFieldResponse } from './bulk_crud_anonymization_fields_route.gen';
export type FindAnonymizationFieldsSortField = z.infer<typeof FindAnonymizationFieldsSortField>;
@ -30,11 +31,6 @@ export const FindAnonymizationFieldsSortField = z.enum([
export type FindAnonymizationFieldsSortFieldEnum = typeof FindAnonymizationFieldsSortField.enum;
export const FindAnonymizationFieldsSortFieldEnum = FindAnonymizationFieldsSortField.enum;
export type SortOrder = z.infer<typeof SortOrder>;
export const SortOrder = z.enum(['asc', 'desc']);
export type SortOrderEnum = typeof SortOrder.enum;
export const SortOrderEnum = SortOrder.enum;
export type FindAnonymizationFieldsRequestQuery = z.infer<
typeof FindAnonymizationFieldsRequestQuery
>;

View file

@ -36,7 +36,7 @@ paths:
description: Sort order
required: false
schema:
$ref: '#/components/schemas/SortOrder'
$ref: '../common_attributes.schema.yaml#/components/schemas/SortOrder'
- name: 'page'
in: query
description: Page number
@ -101,9 +101,3 @@ components:
- 'allowed'
- 'field'
- 'updated_at'
SortOrder:
type: string
enum:
- 'asc'
- 'desc'

View file

@ -45,3 +45,8 @@ export const User = z.object({
*/
name: z.string().optional(),
});
export type SortOrder = z.infer<typeof SortOrder>;
export const SortOrder = z.enum(['asc', 'desc']);
export type SortOrderEnum = typeof SortOrder.enum;
export const SortOrderEnum = SortOrder.enum;

View file

@ -28,3 +28,9 @@ components:
type: string
description: User name
SortOrder:
type: string
enum:
- 'asc'
- 'desc'

View file

@ -17,6 +17,7 @@
import { z } from 'zod';
import { ArrayFromString } from '@kbn/zod-helpers';
import { SortOrder } from '../common_attributes.gen';
import { ConversationResponse } from './common_attributes.gen';
export type FindConversationsSortField = z.infer<typeof FindConversationsSortField>;
@ -29,11 +30,6 @@ export const FindConversationsSortField = z.enum([
export type FindConversationsSortFieldEnum = typeof FindConversationsSortField.enum;
export const FindConversationsSortFieldEnum = FindConversationsSortField.enum;
export type SortOrder = z.infer<typeof SortOrder>;
export const SortOrder = z.enum(['asc', 'desc']);
export type SortOrderEnum = typeof SortOrder.enum;
export const SortOrderEnum = SortOrder.enum;
export type FindConversationsRequestQuery = z.infer<typeof FindConversationsRequestQuery>;
export const FindConversationsRequestQuery = z.object({
fields: ArrayFromString(z.string()).optional(),

View file

@ -36,7 +36,7 @@ paths:
description: Sort order
required: false
schema:
$ref: '#/components/schemas/SortOrder'
$ref: '../common_attributes.schema.yaml#/components/schemas/SortOrder'
- name: 'page'
in: query
description: Page number
@ -124,7 +124,7 @@ paths:
description: Sort order
required: false
schema:
$ref: '#/components/schemas/SortOrder'
$ref: '../common_attributes.schema.yaml#/components/schemas/SortOrder'
- name: 'page'
in: query
description: Page number
@ -188,9 +188,3 @@ components:
- 'is_default'
- 'title'
- 'updated_at'
SortOrder:
type: string
enum:
- 'asc'
- 'desc'

View file

@ -45,3 +45,4 @@ export * from './knowledge_base/crud_kb_route.gen';
export * from './knowledge_base/bulk_crud_knowledge_base_route.gen';
export * from './knowledge_base/common_attributes.gen';
export * from './knowledge_base/crud_knowledge_base_route.gen';
export * from './knowledge_base/find_knowledge_base_entries_route.gen';

View file

@ -121,4 +121,3 @@ components:
type: string
description: Knowledge Base Entry content

View file

@ -0,0 +1,69 @@
/*
* 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.
*/
/*
* NOTICE: Do not edit this file manually.
* This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator.
*
* info:
* title: Find Knowledge Base Entries API endpoint
* version: 1
*/
import { z } from 'zod';
import { ArrayFromString } from '@kbn/zod-helpers';
import { SortOrder } from '../common_attributes.gen';
import { KnowledgeBaseEntryResponse } from './common_attributes.gen';
export type FindKnowledgeBaseEntriesSortField = z.infer<typeof FindKnowledgeBaseEntriesSortField>;
export const FindKnowledgeBaseEntriesSortField = z.enum([
'created_at',
'is_default',
'title',
'updated_at',
]);
export type FindKnowledgeBaseEntriesSortFieldEnum = typeof FindKnowledgeBaseEntriesSortField.enum;
export const FindKnowledgeBaseEntriesSortFieldEnum = FindKnowledgeBaseEntriesSortField.enum;
export type FindKnowledgeBaseEntriesRequestQuery = z.infer<
typeof FindKnowledgeBaseEntriesRequestQuery
>;
export const FindKnowledgeBaseEntriesRequestQuery = z.object({
fields: ArrayFromString(z.string()).optional(),
/**
* Search query
*/
filter: z.string().optional(),
/**
* Field to sort by
*/
sort_field: FindKnowledgeBaseEntriesSortField.optional(),
/**
* Sort order
*/
sort_order: SortOrder.optional(),
/**
* Page number
*/
page: z.coerce.number().int().min(1).optional().default(1),
/**
* Knowledge Base Entries per page
*/
per_page: z.coerce.number().int().min(0).optional().default(20),
});
export type FindKnowledgeBaseEntriesRequestQueryInput = z.input<
typeof FindKnowledgeBaseEntriesRequestQuery
>;
export type FindKnowledgeBaseEntriesResponse = z.infer<typeof FindKnowledgeBaseEntriesResponse>;
export const FindKnowledgeBaseEntriesResponse = z.object({
page: z.number().int(),
perPage: z.number().int(),
total: z.number().int(),
data: z.array(KnowledgeBaseEntryResponse),
});

View file

@ -0,0 +1,102 @@
openapi: 3.0.0
info:
title: Find Knowledge Base Entries API endpoint
version: '1'
paths:
/internal/elastic_assistant/knowledge_base/entries/_find:
get:
operationId: FindKnowledgeBaseEntries
x-codegen-enabled: true
description: Finds Knowledge Base Entries that match the given query.
summary: Finds Knowledge Base Entries that match the given query.
tags:
- Knowledge Base Entries API
parameters:
- name: 'fields'
in: query
required: false
schema:
type: array
items:
type: string
- name: 'filter'
in: query
description: Search query
required: false
schema:
type: string
- name: 'sort_field'
in: query
description: Field to sort by
required: false
schema:
$ref: '#/components/schemas/FindKnowledgeBaseEntriesSortField'
- name: 'sort_order'
in: query
description: Sort order
required: false
schema:
$ref: '../common_attributes.schema.yaml#/components/schemas/SortOrder'
- name: 'page'
in: query
description: Page number
required: false
schema:
type: integer
minimum: 1
default: 1
- name: 'per_page'
in: query
description: Knowledge Base Entries per page
required: false
schema:
type: integer
minimum: 0
default: 20
responses:
200:
description: Successful response
content:
application/json:
schema:
type: object
properties:
page:
type: integer
perPage:
type: integer
total:
type: integer
data:
type: array
items:
$ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse'
required:
- page
- perPage
- total
- data
400:
description: Generic Error
content:
application/json:
schema:
type: object
properties:
statusCode:
type: number
error:
type: string
message:
type: string
components:
schemas:
FindKnowledgeBaseEntriesSortField:
type: string
enum:
- 'created_at'
- 'is_default'
- 'title'
- 'updated_at'

View file

@ -17,6 +17,7 @@
import { z } from 'zod';
import { ArrayFromString } from '@kbn/zod-helpers';
import { SortOrder } from '../common_attributes.gen';
import { PromptResponse } from './bulk_crud_prompts_route.gen';
export type FindPromptsSortField = z.infer<typeof FindPromptsSortField>;
@ -24,11 +25,6 @@ export const FindPromptsSortField = z.enum(['created_at', 'is_default', 'name',
export type FindPromptsSortFieldEnum = typeof FindPromptsSortField.enum;
export const FindPromptsSortFieldEnum = FindPromptsSortField.enum;
export type SortOrder = z.infer<typeof SortOrder>;
export const SortOrder = z.enum(['asc', 'desc']);
export type SortOrderEnum = typeof SortOrder.enum;
export const SortOrderEnum = SortOrder.enum;
export type FindPromptsRequestQuery = z.infer<typeof FindPromptsRequestQuery>;
export const FindPromptsRequestQuery = z.object({
fields: ArrayFromString(z.string()).optional(),

View file

@ -36,7 +36,7 @@ paths:
description: Sort order
required: false
schema:
$ref: '#/components/schemas/SortOrder'
$ref: '../common_attributes.schema.yaml#/components/schemas/SortOrder'
- name: 'page'
in: query
description: Page number
@ -100,9 +100,3 @@ components:
- 'is_default'
- 'name'
- 'updated_at'
SortOrder:
type: string
enum:
- 'asc'
- 'desc'

View file

@ -9,13 +9,7 @@ import { HttpSetup } from '@kbn/core-http-browser';
import { ApiConfig } from '@kbn/elastic-assistant-common';
import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/public/common';
import {
deleteKnowledgeBase,
fetchConnectorExecuteAction,
FetchConnectorExecuteAction,
getKnowledgeBaseStatus,
postKnowledgeBase,
} from '.';
import { fetchConnectorExecuteAction, FetchConnectorExecuteAction } from '.';
import { API_ERROR } from '../translations';
jest.mock('@kbn/core-http-browser');
@ -303,79 +297,4 @@ describe('API tests', () => {
expect(result).toEqual({ response, isStream: false, isError: false });
});
});
const knowledgeBaseArgs = {
resource: 'a-resource',
http: mockHttp,
};
describe('getKnowledgeBaseStatus', () => {
it('calls the knowledge base API when correct resource path', async () => {
await getKnowledgeBaseStatus(knowledgeBaseArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/a-resource',
{
method: 'GET',
signal: undefined,
version: '1',
}
);
});
it('returns error when error is an error', async () => {
const error = 'simulated error';
(mockHttp.fetch as jest.Mock).mockImplementation(() => {
throw new Error(error);
});
await expect(getKnowledgeBaseStatus(knowledgeBaseArgs)).resolves.toThrowError(
'simulated error'
);
});
});
describe('postKnowledgeBase', () => {
it('calls the knowledge base API when correct resource path', async () => {
await postKnowledgeBase(knowledgeBaseArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/a-resource',
{
method: 'POST',
signal: undefined,
version: '1',
}
);
});
it('returns error when error is an error', async () => {
const error = 'simulated error';
(mockHttp.fetch as jest.Mock).mockImplementation(() => {
throw new Error(error);
});
await expect(postKnowledgeBase(knowledgeBaseArgs)).resolves.toThrowError('simulated error');
});
});
describe('deleteKnowledgeBase', () => {
it('calls the knowledge base API when correct resource path', async () => {
await deleteKnowledgeBase(knowledgeBaseArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/a-resource',
{
method: 'DELETE',
signal: undefined,
version: '1',
}
);
});
it('returns error when error is an error', async () => {
const error = 'simulated error';
(mockHttp.fetch as jest.Mock).mockImplementation(() => {
throw new Error(error);
});
await expect(deleteKnowledgeBase(knowledgeBaseArgs)).resolves.toThrowError('simulated error');
});
});
});

View file

@ -6,19 +6,7 @@
*/
import { HttpSetup } from '@kbn/core/public';
import { IHttpFetchError } from '@kbn/core-http-browser';
import {
API_VERSIONS,
ApiConfig,
CreateKnowledgeBaseRequestParams,
CreateKnowledgeBaseResponse,
DeleteKnowledgeBaseRequestParams,
DeleteKnowledgeBaseResponse,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL,
ReadKnowledgeBaseRequestParams,
ReadKnowledgeBaseResponse,
Replacements,
} from '@kbn/elastic-assistant-common';
import { API_VERSIONS, ApiConfig, Replacements } from '@kbn/elastic-assistant-common';
import { API_ERROR } from '../translations';
import { getOptionalRequestParams } from '../helpers';
import { TraceOptions } from '../types';
@ -185,99 +173,3 @@ export const fetchConnectorExecuteAction = async ({
};
}
};
/**
* API call for getting the status of the Knowledge Base. Provide
* a resource to include the status of that specific resource.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {string} [options.resource] - Resource to get the status of, otherwise status of overall KB
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<ReadKnowledgeBaseResponse | IHttpFetchError>}
*/
export const getKnowledgeBaseStatus = async ({
http,
resource,
signal,
}: ReadKnowledgeBaseRequestParams & { http: HttpSetup; signal?: AbortSignal | undefined }): Promise<
ReadKnowledgeBaseResponse | IHttpFetchError
> => {
try {
const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
const response = await http.fetch(path, {
method: 'GET',
signal,
version: API_VERSIONS.internal.v1,
});
return response as ReadKnowledgeBaseResponse;
} catch (error) {
return error as IHttpFetchError;
}
};
/**
* API call for setting up the Knowledge Base. Provide a resource to set up a specific resource.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {string} [options.resource] - Resource to be added to the KB, otherwise sets up the base KB
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<CreateKnowledgeBaseResponse | IHttpFetchError>}
*/
export const postKnowledgeBase = async ({
http,
resource,
signal,
}: CreateKnowledgeBaseRequestParams & {
http: HttpSetup;
signal?: AbortSignal | undefined;
}): Promise<CreateKnowledgeBaseResponse | IHttpFetchError> => {
try {
const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
const response = await http.fetch(path, {
method: 'POST',
signal,
version: API_VERSIONS.internal.v1,
});
return response as CreateKnowledgeBaseResponse;
} catch (error) {
return error as IHttpFetchError;
}
};
/**
* API call for deleting the Knowledge Base. Provide a resource to delete that specific resource.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {string} [options.resource] - Resource to be deleted from the KB, otherwise delete the entire KB
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<DeleteKnowledgeBaseResponse | IHttpFetchError>}
*/
export const deleteKnowledgeBase = async ({
http,
resource,
signal,
}: DeleteKnowledgeBaseRequestParams & {
http: HttpSetup;
signal?: AbortSignal | undefined;
}): Promise<DeleteKnowledgeBaseResponse | IHttpFetchError> => {
try {
const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
const response = await http.fetch(path, {
method: 'DELETE',
signal,
version: API_VERSIONS.internal.v1,
});
return response as DeleteKnowledgeBaseResponse;
} catch (error) {
return error as IHttpFetchError;
}
};

View file

@ -0,0 +1,97 @@
/*
* 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 { deleteKnowledgeBase, getKnowledgeBaseStatus, postKnowledgeBase } from './api';
jest.mock('@kbn/core-http-browser');
const mockHttp = {
fetch: jest.fn(),
} as unknown as HttpSetup;
describe('API tests', () => {
beforeEach(() => {
jest.clearAllMocks();
});
const knowledgeBaseArgs = {
resource: 'a-resource',
http: mockHttp,
};
describe('getKnowledgeBaseStatus', () => {
it('calls the knowledge base API when correct resource path', async () => {
await getKnowledgeBaseStatus(knowledgeBaseArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/a-resource',
{
method: 'GET',
signal: undefined,
version: '1',
}
);
});
it('returns error when error is an error', async () => {
const error = 'simulated error';
(mockHttp.fetch as jest.Mock).mockImplementation(() => {
throw new Error(error);
});
await expect(getKnowledgeBaseStatus(knowledgeBaseArgs)).resolves.toThrowError(
'simulated error'
);
});
});
describe('postKnowledgeBase', () => {
it('calls the knowledge base API when correct resource path', async () => {
await postKnowledgeBase(knowledgeBaseArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/a-resource',
{
method: 'POST',
signal: undefined,
version: '1',
}
);
});
it('returns error when error is an error', async () => {
const error = 'simulated error';
(mockHttp.fetch as jest.Mock).mockImplementation(() => {
throw new Error(error);
});
await expect(postKnowledgeBase(knowledgeBaseArgs)).resolves.toThrowError('simulated error');
});
});
describe('deleteKnowledgeBase', () => {
it('calls the knowledge base API when correct resource path', async () => {
await deleteKnowledgeBase(knowledgeBaseArgs);
expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/a-resource',
{
method: 'DELETE',
signal: undefined,
version: '1',
}
);
});
it('returns error when error is an error', async () => {
const error = 'simulated error';
(mockHttp.fetch as jest.Mock).mockImplementation(() => {
throw new Error(error);
});
await expect(deleteKnowledgeBase(knowledgeBaseArgs)).resolves.toThrowError('simulated error');
});
});
});

View file

@ -0,0 +1,114 @@
/*
* 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 {
API_VERSIONS,
CreateKnowledgeBaseRequestParams,
CreateKnowledgeBaseResponse,
DeleteKnowledgeBaseRequestParams,
DeleteKnowledgeBaseResponse,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL,
ReadKnowledgeBaseRequestParams,
ReadKnowledgeBaseResponse,
} from '@kbn/elastic-assistant-common';
import { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser';
/**
* API call for getting the status of the Knowledge Base. Provide
* a resource to include the status of that specific resource.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {string} [options.resource] - Resource to get the status of, otherwise status of overall KB
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<ReadKnowledgeBaseResponse | IHttpFetchError>}
*/
export const getKnowledgeBaseStatus = async ({
http,
resource,
signal,
}: ReadKnowledgeBaseRequestParams & { http: HttpSetup; signal?: AbortSignal | undefined }): Promise<
ReadKnowledgeBaseResponse | IHttpFetchError
> => {
try {
const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
const response = await http.fetch(path, {
method: 'GET',
signal,
version: API_VERSIONS.internal.v1,
});
return response as ReadKnowledgeBaseResponse;
} catch (error) {
return error as IHttpFetchError;
}
};
/**
* API call for setting up the Knowledge Base. Provide a resource to set up a specific resource.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {string} [options.resource] - Resource to be added to the KB, otherwise sets up the base KB
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<CreateKnowledgeBaseResponse | IHttpFetchError>}
*/
export const postKnowledgeBase = async ({
http,
resource,
signal,
}: CreateKnowledgeBaseRequestParams & {
http: HttpSetup;
signal?: AbortSignal | undefined;
}): Promise<CreateKnowledgeBaseResponse | IHttpFetchError> => {
try {
const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
const response = await http.fetch(path, {
method: 'POST',
signal,
version: API_VERSIONS.internal.v1,
});
return response as CreateKnowledgeBaseResponse;
} catch (error) {
return error as IHttpFetchError;
}
};
/**
* API call for deleting the Knowledge Base. Provide a resource to delete that specific resource.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {string} [options.resource] - Resource to be deleted from the KB, otherwise delete the entire KB
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<DeleteKnowledgeBaseResponse | IHttpFetchError>}
*/
export const deleteKnowledgeBase = async ({
http,
resource,
signal,
}: DeleteKnowledgeBaseRequestParams & {
http: HttpSetup;
signal?: AbortSignal | undefined;
}): Promise<DeleteKnowledgeBaseResponse | IHttpFetchError> => {
try {
const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
const response = await http.fetch(path, {
method: 'DELETE',
signal,
version: API_VERSIONS.internal.v1,
});
return response as DeleteKnowledgeBaseResponse;
} catch (error) {
return error as IHttpFetchError;
}
};

View file

@ -0,0 +1,83 @@
/*
* 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 { useMutation } 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 {
API_VERSIONS,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL,
KnowledgeBaseEntryCreateProps,
KnowledgeBaseEntryResponse,
} from '@kbn/elastic-assistant-common';
import { useInvalidateKnowledgeBaseEntries } from './use_knowledge_base_entries';
const CREATE_KNOWLEDGE_BASE_ENTRY_MUTATION_KEY = [
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL,
API_VERSIONS.internal.v1,
];
export interface UseCreateKnowledgeBaseEntryParams {
http: HttpSetup;
signal?: AbortSignal;
toasts?: IToasts;
}
/**
* Hook for creating a Knowledge Base Entry
*
* @param {Object} options - The options object
* @param {HttpSetup} options.http - HttpSetup
* @param {AbortSignal} [options.signal] - AbortSignal
* @param {IToasts} [options.toasts] - IToasts
*
* @returns mutation hook for creating a Knowledge Base Entry
*
*/
export const useCreateKnowledgeBaseEntry = ({
http,
signal,
toasts,
}: UseCreateKnowledgeBaseEntryParams) => {
const invalidateKnowledgeBaseEntries = useInvalidateKnowledgeBaseEntries();
return useMutation(
CREATE_KNOWLEDGE_BASE_ENTRY_MUTATION_KEY,
(entry: KnowledgeBaseEntryCreateProps) => {
return http.post<KnowledgeBaseEntryResponse>(
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL,
{
body: JSON.stringify(entry),
version: API_VERSIONS.internal.v1,
signal,
}
);
},
{
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.knowledgeBase.entries.createErrorTitle',
{
defaultMessage: 'Error creating Knowledge Base Entry',
}
),
}
);
}
},
onSettled: () => {
invalidateKnowledgeBaseEntries();
},
}
);
};

View file

@ -0,0 +1,90 @@
/*
* 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 { useMutation } 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 {
API_VERSIONS,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION,
KnowledgeBaseEntryBulkActionBase,
KnowledgeBaseEntryBulkCrudActionResponse,
PerformKnowledgeBaseEntryBulkActionRequestBody,
} from '@kbn/elastic-assistant-common';
import { useInvalidateKnowledgeBaseEntries } from './use_knowledge_base_entries';
const DELETE_KNOWLEDGE_BASE_ENTRIES_MUTATION_KEY = [
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION,
API_VERSIONS.internal.v1,
];
export interface UseDeleteKnowledgeEntriesParams {
http: HttpSetup;
signal?: AbortSignal;
toasts?: IToasts;
}
/**
* Hook for deleting Knowledge Base Entries by id or query.
*
* @param {Object} options - The options object
* @param {HttpSetup} options.http - HttpSetup
* @param {AbortSignal} [options.signal] - AbortSignal
* @param {IToasts} [options.toasts] - IToasts
*
* @returns mutation hook for deleting Knowledge Base Entries
*
*/
export const useDeleteKnowledgeBaseEntries = ({
http,
signal,
toasts,
}: UseDeleteKnowledgeEntriesParams) => {
const invalidateKnowledgeBaseEntries = useInvalidateKnowledgeBaseEntries();
return useMutation(
DELETE_KNOWLEDGE_BASE_ENTRIES_MUTATION_KEY,
({ ids, query }: KnowledgeBaseEntryBulkActionBase) => {
const body: PerformKnowledgeBaseEntryBulkActionRequestBody = {
delete: {
query,
ids,
},
};
return http.post<KnowledgeBaseEntryBulkCrudActionResponse>(
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION,
{
body: JSON.stringify(body),
version: API_VERSIONS.internal.v1,
signal,
}
);
},
{
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.knowledgeBase.entries.deleteErrorTitle',
{
defaultMessage: 'Error deleting Knowledge Base Entries',
}
),
}
);
}
},
onSettled: () => {
invalidateKnowledgeBaseEntries();
},
}
);
};

View file

@ -0,0 +1,80 @@
/*
* 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/public';
import { useQuery, useQueryClient } from '@tanstack/react-query';
import {
API_VERSIONS,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND,
FindKnowledgeBaseEntriesResponse,
} from '@kbn/elastic-assistant-common';
import { useCallback } from 'react';
export interface UseKnowledgeBaseEntriesParams {
http: HttpSetup;
signal?: AbortSignal | undefined;
}
/**
* Hook for fetching Knowledge Base Entries.
*
* Note: RBAC is handled at kbDataClient layer, so unless user has KB feature privileges, this will only return system and their own user KB entries.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {Function} [options.onFetch] - transformation function for kb entries fetch result
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {useQuery} hook for fetching Knowledge Base Entries
*/
const query = {
page: 1,
perPage: 100,
};
export const KNOWLEDGE_BASE_ENTRY_QUERY_KEY = [
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND,
query.page,
query.perPage,
API_VERSIONS.internal.v1,
];
export const useKnowledgeBaseEntries = ({ http, signal }: UseKnowledgeBaseEntriesParams) =>
useQuery(
KNOWLEDGE_BASE_ENTRY_QUERY_KEY,
async () =>
http.fetch<FindKnowledgeBaseEntriesResponse>(
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND,
{
method: 'GET',
version: API_VERSIONS.internal.v1,
query,
signal,
}
),
{
keepPreviousData: true,
initialData: { page: 1, perPage: 100, total: 0, data: [] },
}
);
/**
* Use this hook to invalidate the Knowledge Base Entries cache. For example, adding,
* editing, or deleting any Knowledge Base entries should lead to cache invalidation.
*
* @returns {Function} - Function to invalidate the Knowledge Base Entries cache
*/
export const useInvalidateKnowledgeBaseEntries = () => {
const queryClient = useQueryClient();
return useCallback(() => {
queryClient.invalidateQueries(KNOWLEDGE_BASE_ENTRY_QUERY_KEY, {
refetchType: 'active',
});
}, [queryClient]);
};

View file

@ -7,14 +7,14 @@
import { act, renderHook } from '@testing-library/react-hooks';
import { useDeleteKnowledgeBase, UseDeleteKnowledgeBaseParams } from './use_delete_knowledge_base';
import { deleteKnowledgeBase as _deleteKnowledgeBase } from '../assistant/api';
import { deleteKnowledgeBase as _deleteKnowledgeBase } from './api';
import { useMutation as _useMutation } from '@tanstack/react-query';
const useMutationMock = _useMutation as jest.Mock;
const deleteKnowledgeBaseMock = _deleteKnowledgeBase as jest.Mock;
jest.mock('../assistant/api', () => {
const actual = jest.requireActual('../assistant/api');
jest.mock('./api', () => {
const actual = jest.requireActual('./api');
return {
...actual,
deleteKnowledgeBase: jest.fn((...args) => actual.deleteKnowledgeBase(...args)),

View file

@ -9,7 +9,7 @@ import { useMutation } from '@tanstack/react-query';
import type { IToasts } from '@kbn/core-notifications-browser';
import type { HttpSetup, IHttpFetchError, ResponseErrorBody } from '@kbn/core-http-browser';
import { i18n } from '@kbn/i18n';
import { deleteKnowledgeBase } from '../assistant/api';
import { deleteKnowledgeBase } from './api';
import { useInvalidateKnowledgeBaseStatus } from './use_knowledge_base_status';
const DELETE_KNOWLEDGE_BASE_MUTATION_KEY = ['elastic-assistant', 'delete-knowledge-base'];

View file

@ -7,12 +7,12 @@
import { act, renderHook } from '@testing-library/react-hooks';
import { useKnowledgeBaseStatus, UseKnowledgeBaseStatusParams } from './use_knowledge_base_status';
import { getKnowledgeBaseStatus as _getKnowledgeBaseStatus } from '../assistant/api';
import { getKnowledgeBaseStatus as _getKnowledgeBaseStatus } from './api';
const getKnowledgeBaseStatusMock = _getKnowledgeBaseStatus as jest.Mock;
jest.mock('../assistant/api', () => {
const actual = jest.requireActual('../assistant/api');
jest.mock('./api', () => {
const actual = jest.requireActual('./api');
return {
...actual,
getKnowledgeBaseStatus: jest.fn((...args) => actual.getKnowledgeBaseStatus(...args)),

View file

@ -12,7 +12,7 @@ import type { IToasts } from '@kbn/core-notifications-browser';
import { i18n } from '@kbn/i18n';
import { useCallback } from 'react';
import { ReadKnowledgeBaseResponse } from '@kbn/elastic-assistant-common';
import { getKnowledgeBaseStatus } from '../assistant/api';
import { getKnowledgeBaseStatus } from './api';
const KNOWLEDGE_BASE_STATUS_QUERY_KEY = ['elastic-assistant', 'knowledge-base-status'];

View file

@ -7,13 +7,13 @@
import { act, renderHook } from '@testing-library/react-hooks';
import { useSetupKnowledgeBase, UseSetupKnowledgeBaseParams } from './use_setup_knowledge_base';
import { postKnowledgeBase as _postKnowledgeBase } from '../assistant/api';
import { postKnowledgeBase as _postKnowledgeBase } from './api';
import { useMutation as _useMutation } from '@tanstack/react-query';
const postKnowledgeBaseMock = _postKnowledgeBase as jest.Mock;
const useMutationMock = _useMutation as jest.Mock;
jest.mock('../assistant/api', () => {
const actual = jest.requireActual('../assistant/api');
jest.mock('./api', () => {
const actual = jest.requireActual('./api');
return {
...actual,
postKnowledgeBase: jest.fn((...args) => actual.postKnowledgeBase(...args)),

View file

@ -9,7 +9,7 @@ import { useMutation } 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 { postKnowledgeBase } from '../assistant/api';
import { postKnowledgeBase } from './api';
import { useInvalidateKnowledgeBaseStatus } from './use_knowledge_base_status';
const SETUP_KNOWLEDGE_BASE_MUTATION_KEY = ['elastic-assistant', 'post-knowledge-base'];

View file

@ -10,8 +10,8 @@ import { EuiButton } from '@elastic/eui';
import { i18n } from '@kbn/i18n';
import { useAssistantContext } from '../..';
import { useSetupKnowledgeBase } from './use_setup_knowledge_base';
import { useKnowledgeBaseStatus } from './use_knowledge_base_status';
import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base';
import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
const ESQL_RESOURCE = 'esql';

View file

@ -11,7 +11,7 @@ import { fireEvent, render } from '@testing-library/react';
import { DEFAULT_LATEST_ALERTS } from '../assistant_context/constants';
import { KnowledgeBaseSettings } from './knowledge_base_settings';
import { TestProviders } from '../mock/test_providers/test_providers';
import { useKnowledgeBaseStatus } from './use_knowledge_base_status';
import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
import { mockSystemPrompts } from '../mock/system_prompt';
import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
@ -47,7 +47,7 @@ const defaultProps = {
setUpdatedKnowledgeBaseSettings,
};
const mockDelete = jest.fn();
jest.mock('./use_delete_knowledge_base', () => ({
jest.mock('../assistant/api/knowledge_base/use_delete_knowledge_base', () => ({
useDeleteKnowledgeBase: jest.fn(() => {
return {
mutate: mockDelete,
@ -57,7 +57,7 @@ jest.mock('./use_delete_knowledge_base', () => ({
}));
const mockSetup = jest.fn();
jest.mock('./use_setup_knowledge_base', () => ({
jest.mock('../assistant/api/knowledge_base/use_setup_knowledge_base', () => ({
useSetupKnowledgeBase: jest.fn(() => {
return {
mutate: mockSetup,
@ -66,7 +66,7 @@ jest.mock('./use_setup_knowledge_base', () => ({
}),
}));
jest.mock('./use_knowledge_base_status', () => ({
jest.mock('../assistant/api/knowledge_base/use_knowledge_base_status', () => ({
useKnowledgeBaseStatus: jest.fn(() => {
return {
data: {

View file

@ -30,9 +30,9 @@ import { AlertsSettings } from '../alerts/settings/alerts_settings';
import { useAssistantContext } from '../assistant_context';
import type { KnowledgeBaseConfig } from '../assistant/types';
import * as i18n from './translations';
import { useDeleteKnowledgeBase } from './use_delete_knowledge_base';
import { useKnowledgeBaseStatus } from './use_knowledge_base_status';
import { useSetupKnowledgeBase } from './use_setup_knowledge_base';
import { useDeleteKnowledgeBase } from '../assistant/api/knowledge_base/use_delete_knowledge_base';
import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base';
const ESQL_RESOURCE = 'esql';
const KNOWLEDGE_BASE_INDEX_PATTERN_OLD = '.kibana-elastic-ai-assistant-kb';

View file

@ -27,6 +27,9 @@ export const ANONYMIZATION_FIELDS_TABLE_MAX_PAGE_SIZE = 100;
export const MAX_PROMPTS_TO_UPDATE_IN_PARALLEL = 50;
export const PROMPTS_TABLE_MAX_PAGE_SIZE = 100;
// Knowledge Base
export const KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE = 100;
// Capabilities
export const CAPABILITIES = `${BASE_PATH}/capabilities`;

View file

@ -16,6 +16,7 @@ import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-ser
import {
KnowledgeBaseEntryCreateProps,
KnowledgeBaseEntryResponse,
Metadata,
} from '@kbn/elastic-assistant-common';
import pRetry from 'p-retry';
import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
@ -221,12 +222,12 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
/**
* Adds LangChain Documents to the knowledge base
*
* @param documents LangChain Documents to add to the knowledge base
* @param {Array<Document<Metadata>>} documents - LangChain Documents to add to the knowledge base
*/
public addKnowledgeBaseDocuments = async ({
documents,
}: {
documents: Document[];
documents: Array<Document<Metadata>>;
}): Promise<KnowledgeBaseEntryResponse[]> => {
const writer = await this.getWriter();
const changedAt = new Date().toISOString();
@ -240,9 +241,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
const { errors, docs_created: docsCreated } = await writer.bulk({
documentsToCreate: documents.map((doc) =>
transformToCreateSchema(changedAt, this.spaceId, authenticatedUser, {
// TODO: Update the LangChain Document Metadata type extension
metadata: {
kbResource: doc.metadata.kbResourcer ?? 'unknown',
kbResource: doc.metadata.kbResource ?? 'unknown',
required: doc.metadata.required ?? false,
source: doc.metadata.source ?? 'unknown',
},

View file

@ -9,7 +9,9 @@ import { Logger } from '@kbn/core/server';
import { DirectoryLoader } from 'langchain/document_loaders/fs/directory';
import { TextLoader } from 'langchain/document_loaders/fs/text';
import { resolve } from 'path';
import { Document } from 'langchain/document';
import { Metadata } from '@kbn/elastic-assistant-common';
import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store';
import { addRequiredKbResourceMetadata } from './add_required_kb_resource_metadata';
import { ESQL_RESOURCE } from '../../../routes/knowledge_base/constants';
@ -47,15 +49,15 @@ export const loadESQL = async (esStore: ElasticsearchStore, logger: Logger): Pro
true
);
const docs = await docsLoader.load();
const languageDocs = await languageLoader.load();
const docs = (await docsLoader.load()) as Array<Document<Metadata>>;
const languageDocs = (await languageLoader.load()) as Array<Document<Metadata>>;
const rawExampleQueries = await exampleQueriesLoader.load();
// Add additional metadata to the example queries that indicates they are required KB documents:
const requiredExampleQueries = addRequiredKbResourceMetadata({
docs: rawExampleQueries,
kbResource: ESQL_RESOURCE,
});
}) as Array<Document<Metadata>>;
logger.info(
`Loading ${docs.length} ES|QL docs, ${languageDocs.length} language docs, and ${requiredExampleQueries.length} example queries into the Knowledge Base`

View file

@ -25,6 +25,7 @@ import {
KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT,
KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT,
} from '../../telemetry/event_based_telemetry';
import { Metadata } from '@kbn/elastic-assistant-common';
jest.mock('uuid', () => ({
v4: jest.fn(),
@ -244,9 +245,9 @@ describe('ElasticsearchStore', () => {
],
});
const document = new Document({
const document = new Document<Metadata>({
pageContent: 'interesting stuff',
metadata: { source: '1' },
metadata: { kbResource: 'esql', required: false, source: '1' },
});
const docsInstalled = await esStore.addDocuments([document]);
@ -262,6 +263,8 @@ describe('ElasticsearchStore', () => {
},
{
metadata: {
kbResource: 'esql',
required: false,
source: '1',
},
text: 'interesting stuff',

View file

@ -17,6 +17,7 @@ import { Document } from 'langchain/document';
import { VectorStore } from '@langchain/core/vectorstores';
import * as uuid from 'uuid';
import { Metadata } from '@kbn/elastic-assistant-common';
import { transformError } from '@kbn/securitysolution-es-utils';
import { ElasticsearchEmbeddings } from '../embeddings/elasticsearch_embeddings';
import { FlattenedHit, getFlattenedHits } from './helpers/get_flattened_hits';
@ -105,7 +106,7 @@ export class ElasticsearchStore extends VectorStore {
* @returns Promise<string[]> of document IDs added to the store
*/
addDocuments = async (
documents: Document[],
documents: Array<Document<Metadata>>,
options?: Record<string, never>
): Promise<string[]> => {
// Code path for when `assistantKnowledgeBaseByDefault` FF is enabled
@ -145,7 +146,7 @@ export class ElasticsearchStore extends VectorStore {
};
addDocumentsViaDataClient = async (
documents: Document[],
documents: Array<Document<Metadata>>,
options?: Record<string, never>
): Promise<string[]> => {
if (!this.kbDataClient) {

View file

@ -5,11 +5,15 @@
* 2.0.
*/
import { KibanaRequest } from '@kbn/core-http-server';
import { IKibanaResponse, KibanaRequest, KibanaResponseFactory } from '@kbn/core-http-server';
import { Logger } from '@kbn/core/server';
import { Message, TraceData } from '@kbn/elastic-assistant-common';
import { ILicense } from '@kbn/licensing-plugin/server';
import { AwaitedProperties } from '@kbn/utility-types';
import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities';
import { MINIMUM_AI_ASSISTANT_LICENSE } from '../../common/constants';
import { ElasticAssistantRequestHandlerContext } from '../types';
import { buildResponse } from './utils';
interface GetPluginNameFromRequestParams {
request: KibanaRequest;
@ -87,3 +91,69 @@ export const hasAIAssistantLicense = (license: ILicense): boolean =>
export const UPGRADE_LICENSE_MESSAGE =
'Your license does not support AI Assistant. Please upgrade your license.';
interface PerformChecksParams {
authenticatedUser?: boolean;
capability?: AssistantFeatureKey;
context: AwaitedProperties<
Pick<ElasticAssistantRequestHandlerContext, 'elasticAssistant' | 'licensing' | 'core'>
>;
license?: boolean;
request: KibanaRequest;
response: KibanaResponseFactory;
}
/**
* Helper to perform checks for authenticated user, capability, and license. Perform all or one
* of the checks by providing relevant optional params. Check order is license, authenticated user,
* then capability.
*
* @param authenticatedUser - Whether to check for an authenticated user
* @param capability - Specific capability to check if enabled, e.g. `assistantModelEvaluation`
* @param context - Route context
* @param license - Whether to check for a valid license
* @param request - Route KibanaRequest
* @param response - Route KibanaResponseFactory
*/
export const performChecks = ({
authenticatedUser,
capability,
context,
license,
request,
response,
}: PerformChecksParams): IKibanaResponse | undefined => {
const assistantResponse = buildResponse(response);
if (license) {
if (!hasAIAssistantLicense(context.licensing.license)) {
return response.forbidden({
body: {
message: UPGRADE_LICENSE_MESSAGE,
},
});
}
}
if (authenticatedUser) {
if (context.elasticAssistant.getCurrentUser() == null) {
return assistantResponse.error({
body: `Authenticated user not found`,
statusCode: 401,
});
}
}
if (capability) {
const pluginName = getPluginNameFromRequest({
request,
defaultPluginName: DEFAULT_PLUGIN_NAME,
});
const registeredFeatures = context.elasticAssistant.getRegisteredFeatures(pluginName);
if (!registeredFeatures[capability]) {
return response.notFound();
}
}
return undefined;
};

View file

@ -0,0 +1,235 @@
/*
* 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 moment from 'moment';
import type { AuthenticatedUser, IKibanaResponse, KibanaResponseFactory } from '@kbn/core/server';
import { transformError } from '@kbn/securitysolution-es-utils';
import {
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION,
PerformKnowledgeBaseEntryBulkActionRequestBody,
API_VERSIONS,
KnowledgeBaseEntryBulkCrudActionResults,
KnowledgeBaseEntryBulkCrudActionResponse,
KnowledgeBaseEntryBulkCrudActionSummary,
PerformKnowledgeBaseEntryBulkActionResponse,
} from '@kbn/elastic-assistant-common';
import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
import { performChecks } from '../../helpers';
import { KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE } from '../../../../common/constants';
import { EsKnowledgeBaseEntrySchema } from '../../../ai_assistant_data_clients/knowledge_base/types';
import { ElasticAssistantPluginRouter } from '../../../types';
import { buildResponse } from '../../utils';
import { transformESSearchToKnowledgeBaseEntry } from '../../../ai_assistant_data_clients/knowledge_base/transforms';
import { transformToCreateSchema } from '../../../ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry';
export interface BulkOperationError {
message: string;
status?: number;
document: {
id: string;
name?: string;
};
}
export type BulkResponse = KnowledgeBaseEntryBulkCrudActionResults & {
errors?: BulkOperationError[];
};
export type BulkActionError = BulkOperationError | unknown;
const buildBulkResponse = (
response: KibanaResponseFactory,
{
errors = [],
updated = [],
created = [],
deleted = [],
skipped = [],
}: KnowledgeBaseEntryBulkCrudActionResults & { errors: BulkOperationError[] }
): IKibanaResponse<KnowledgeBaseEntryBulkCrudActionResponse> => {
const numSucceeded = updated.length + created.length + deleted.length;
const numSkipped = skipped.length;
const numFailed = errors.length;
const summary: KnowledgeBaseEntryBulkCrudActionSummary = {
failed: numFailed,
succeeded: numSucceeded,
skipped: numSkipped,
total: numSucceeded + numFailed + numSkipped,
};
const results: KnowledgeBaseEntryBulkCrudActionResults = {
updated,
created,
deleted,
skipped,
};
if (numFailed > 0) {
return response.custom<KnowledgeBaseEntryBulkCrudActionResponse>({
headers: { 'content-type': 'application/json' },
body: {
message: summary.succeeded > 0 ? 'Bulk edit partially failed' : 'Bulk edit failed',
attributes: {
errors: errors.map((e: BulkOperationError) => ({
statusCode: e.status ?? 500,
knowledgeBaseEntries: [{ id: e.document.id, name: '' }],
message: e.message,
})),
results,
summary,
},
},
statusCode: 500,
});
}
const responseBody: KnowledgeBaseEntryBulkCrudActionResponse = {
success: true,
knowledgeBaseEntriesCount: summary.total,
attributes: { results, summary },
};
return response.ok({ body: responseBody });
};
export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRouter) => {
router.versioned
.post({
access: 'internal',
path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION,
options: {
tags: ['access:elasticAssistant'],
timeout: {
idleSocket: moment.duration(15, 'minutes').asMilliseconds(),
},
},
})
.addVersion(
{
version: API_VERSIONS.internal.v1,
validate: {
request: {
body: buildRouteValidationWithZod(PerformKnowledgeBaseEntryBulkActionRequestBody),
},
},
},
async (
context,
request,
response
): Promise<IKibanaResponse<PerformKnowledgeBaseEntryBulkActionResponse>> => {
const assistantResponse = buildResponse(response);
try {
const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
const logger = ctx.elasticAssistant.logger;
// Perform license, authenticated user and FF checks
const checkResponse = performChecks({
authenticatedUser: true,
capability: 'assistantKnowledgeBaseByDefault',
context: ctx,
license: true,
request,
response,
});
if (checkResponse) {
return checkResponse;
}
logger.debug(
`Performing bulk action on Knowledge Base Entries:\n${JSON.stringify(request.body)}`
);
const { body } = request;
const operationsCount =
(body?.update ? body.update?.length : 0) +
(body?.create ? body.create?.length : 0) +
(body?.delete ? body.delete?.ids?.length ?? 0 : 0);
if (operationsCount > KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE) {
return assistantResponse.error({
body: `More than ${KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE} ids sent for bulk edit action.`,
statusCode: 400,
});
}
const abortController = new AbortController();
// subscribing to completed$, because it handles both cases when request was completed and aborted.
// when route is finished by timeout, aborted$ is not getting fired
request.events.completed$.subscribe(() => abortController.abort());
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(
false
);
const spaceId = ctx.elasticAssistant.getSpaceId();
// Authenticated user null check completed in `performChecks()` above
const authenticatedUser = ctx.elasticAssistant.getCurrentUser() as AuthenticatedUser;
if (body.create && body.create.length > 0) {
const result = await kbDataClient?.findDocuments<EsKnowledgeBaseEntrySchema>({
perPage: 100,
page: 1,
filter: `users:{ id: "${authenticatedUser?.profile_uid}" }`,
fields: [],
});
if (result?.data != null && result.total > 0) {
return assistantResponse.error({
statusCode: 409,
body: `Knowledge Base Entry id's: "${transformESSearchToKnowledgeBaseEntry(
result.data
)
.map((c) => c.id)
.join(',')}" already exists`,
});
}
}
const writer = await kbDataClient?.getWriter();
const changedAt = new Date().toISOString();
const {
errors,
docs_created: docsCreated,
docs_updated: docsUpdated,
docs_deleted: docsDeleted,
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
} = await writer!.bulk({
documentsToCreate: body.create?.map((c) =>
transformToCreateSchema(changedAt, spaceId, authenticatedUser, c)
),
documentsToDelete: body.delete?.ids,
documentsToUpdate: [], // TODO: Support bulk update
authenticatedUser,
});
const created =
docsCreated.length > 0
? await kbDataClient?.findDocuments<EsKnowledgeBaseEntrySchema>({
page: 1,
perPage: 100,
filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
})
: undefined;
return buildBulkResponse(response, {
// @ts-ignore-next-line TS2322
updated: docsUpdated,
created: created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : [],
deleted: docsDeleted ?? [],
errors,
});
} catch (err) {
const error = transformError(err);
return assistantResponse.error({
body: error.message,
statusCode: error.statusCode,
});
}
}
);
};

View file

@ -5,7 +5,7 @@
* 2.0.
*/
import type { IKibanaResponse } from '@kbn/core/server';
import { IKibanaResponse } from '@kbn/core/server';
import { transformError } from '@kbn/securitysolution-es-utils';
import {
API_VERSIONS,
@ -15,15 +15,17 @@ import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/
import {
KnowledgeBaseEntryCreateProps,
KnowledgeBaseEntryResponse,
Metadata,
} from '@kbn/elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen';
import type { Document } from 'langchain/document';
import { ElasticAssistantPluginRouter } from '../../../types';
import { buildResponse } from '../../utils';
import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../../helpers';
import { performChecks } from '../../helpers';
export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRouter): void => {
router.versioned
.post({
access: 'public',
access: 'internal',
path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL,
options: {
@ -32,7 +34,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
})
.addVersion(
{
version: API_VERSIONS.public.v1,
version: API_VERSIONS.internal.v1,
validate: {
request: {
body: buildRouteValidationWithZod(KnowledgeBaseEntryCreateProps),
@ -43,27 +45,40 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
const assistantResponse = buildResponse(response);
try {
const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
const license = ctx.licensing.license;
if (!hasAIAssistantLicense(license)) {
return response.forbidden({
body: {
message: UPGRADE_LICENSE_MESSAGE,
},
});
}
const logger = ctx.elasticAssistant.logger;
const authenticatedUser = ctx.elasticAssistant.getCurrentUser();
if (authenticatedUser == null) {
return assistantResponse.error({
body: `Authenticated user not found`,
statusCode: 401,
});
}
return assistantResponse.error({
body: `knowledge base entry was not created`,
statusCode: 400,
// Perform license, authenticated user and FF checks
const checkResponse = performChecks({
authenticatedUser: true,
capability: 'assistantKnowledgeBaseByDefault',
context: ctx,
license: true,
request,
response,
});
if (checkResponse) {
return checkResponse;
}
logger.debug(`Creating KB Entry:\n${JSON.stringify(request.body)}`);
const documents: Array<Document<Metadata>> = [
{
metadata: request.body.metadata,
pageContent: request.body.text,
},
];
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(
false
);
const createResponse = await kbDataClient?.addKnowledgeBaseDocuments({ documents });
if (createResponse == null) {
return assistantResponse.error({
body: `Knowledge Base Entry was not created`,
statusCode: 400,
});
}
return response.ok({ body: KnowledgeBaseEntryResponse.parse(createResponse[0]) });
} catch (err) {
const error = transformError(err as Error);
return assistantResponse.error({

View file

@ -0,0 +1,103 @@
/*
* 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 { IKibanaResponse } from '@kbn/core/server';
import { transformError } from '@kbn/securitysolution-es-utils';
import {
API_VERSIONS,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND,
FindKnowledgeBaseEntriesRequestQuery,
FindKnowledgeBaseEntriesResponse,
} from '@kbn/elastic-assistant-common';
import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
import { ElasticAssistantPluginRouter } from '../../../types';
import { buildResponse } from '../../utils';
import { performChecks } from '../../helpers';
import { transformESSearchToKnowledgeBaseEntry } from '../../../ai_assistant_data_clients/knowledge_base/transforms';
import { EsKnowledgeBaseEntrySchema } from '../../../ai_assistant_data_clients/knowledge_base/types';
export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRouter) => {
router.versioned
.get({
access: 'internal',
path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND,
options: {
tags: ['access:elasticAssistant'],
},
})
.addVersion(
{
version: API_VERSIONS.internal.v1,
validate: {
request: {
query: buildRouteValidationWithZod(FindKnowledgeBaseEntriesRequestQuery),
},
},
},
async (
context,
request,
response
): Promise<IKibanaResponse<FindKnowledgeBaseEntriesResponse>> => {
const assistantResponse = buildResponse(response);
try {
const { query } = request;
const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
// Perform license, authenticated user and FF checks
const checkResponse = performChecks({
authenticatedUser: true,
capability: 'assistantKnowledgeBaseByDefault',
context: ctx,
license: true,
request,
response,
});
if (checkResponse) {
return checkResponse;
}
const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(
false
);
const currentUser = ctx.elasticAssistant.getCurrentUser();
const additionalFilter = query.filter ? ` AND ${query.filter}` : '';
const result = await kbDataClient?.findDocuments<EsKnowledgeBaseEntrySchema>({
perPage: query.per_page,
page: query.page,
sortField: query.sort_field,
sortOrder: query.sort_order,
filter: `users:{ id: "${currentUser?.profile_uid}" }${additionalFilter}`, // TODO: Update filter to include non-user system entries
fields: query.fields,
});
if (result) {
return response.ok({
body: {
perPage: result.perPage,
page: result.page,
total: result.total,
data: transformESSearchToKnowledgeBaseEntry(result.data),
},
});
}
return response.ok({
body: { perPage: query.per_page, page: query.page, data: [], total: 0 },
});
} catch (err) {
const error = transformError(err);
return assistantResponse.error({
body: error.message,
statusCode: error.statusCode,
});
}
}
);
};

View file

@ -27,6 +27,9 @@ import { bulkPromptsRoute } from './prompts/bulk_actions_route';
import { findPromptsRoute } from './prompts/find_route';
import { bulkActionAnonymizationFieldsRoute } from './anonymization_fields/bulk_actions_route';
import { findAnonymizationFieldsRoute } from './anonymization_fields/find_route';
import { bulkActionKnowledgeBaseEntriesRoute } from './knowledge_base/entries/bulk_actions_route';
import { createKnowledgeBaseEntryRoute } from './knowledge_base/entries/create_route';
import { findKnowledgeBaseEntriesRoute } from './knowledge_base/entries/find_route';
export const registerRoutes = (
router: ElasticAssistantPluginRouter,
@ -49,11 +52,16 @@ export const registerRoutes = (
// User Conversations search
findUserConversationsRoute(router);
// Knowledge Base
// Knowledge Base Setup
deleteKnowledgeBaseRoute(router);
getKnowledgeBaseStatusRoute(router, getElserId);
postKnowledgeBaseRoute(router, getElserId);
// Knowledge Base Entries
findKnowledgeBaseEntriesRoute(router);
createKnowledgeBaseEntryRoute(router);
bulkActionKnowledgeBaseEntriesRoute(router);
// Actions Connector Execute (LLM Wrapper)
postActionsConnectorExecuteRoute(router, getElserId);