Skip to content

Commit

Permalink
[Security Assistant] Adds client hooks and internal routes for managi…
Browse files Browse the repository at this point in the history
…ng Knowledge Base Entries (elastic#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
elastic/security-docs#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>
  • Loading branch information
2 people authored and seanrathier committed Jun 21, 2024
1 parent b91e4bf commit dbc2363
Show file tree
Hide file tree
Showing 40 changed files with 1,158 additions and 285 deletions.
5 changes: 3 additions & 2 deletions x-pack/packages/kbn-elastic-assistant-common/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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`;
Original file line number Diff line number Diff line change
Expand Up @@ -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
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand All @@ -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
>;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -101,9 +101,3 @@ components:
- 'allowed'
- 'field'
- 'updated_at'

SortOrder:
type: string
enum:
- 'asc'
- 'desc'
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Original file line number Diff line number Diff line change
Expand Up @@ -28,3 +28,9 @@ components:
type: string
description: User name

SortOrder:
type: string
enum:
- 'asc'
- 'desc'

Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand All @@ -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(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -188,9 +188,3 @@ components:
- 'is_default'
- 'title'
- 'updated_at'

SortOrder:
type: string
enum:
- 'asc'
- 'desc'
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Original file line number Diff line number Diff line change
Expand Up @@ -121,4 +121,3 @@ components:
type: string
description: Knowledge Base Entry content


Original file line number Diff line number Diff line change
@@ -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),
});
Original file line number Diff line number Diff line change
@@ -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'
Original file line number Diff line number Diff line change
Expand Up @@ -17,18 +17,14 @@
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>;
export const FindPromptsSortField = z.enum(['created_at', 'is_default', 'name', 'updated_at']);
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(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -100,9 +100,3 @@ components:
- 'is_default'
- 'name'
- 'updated_at'

SortOrder:
type: string
enum:
- 'asc'
- 'desc'
Loading

0 comments on commit dbc2363

Please sign in to comment.