Skip to content

Commit

Permalink
[Security GenAI][BUG] Knowledge Base: Show only indices with `semanti…
Browse files Browse the repository at this point in the history
…c_text` fields (#198707)

## Summary

This is a fix the next issue:

> Index input should only list indices with semantic_text fields, not
all indices.

### Current behaviour

We show all available indices

<img width="1311" alt="Screenshot 2024-11-01 at 18 14 36"
src="https://github.com/user-attachments/assets/cf9d08fd-a809-4530-b653-d12b8e643e45">

### Behaviour after the fix

We show only indices with `semantic_text` fields

<img width="1311" alt="Screenshot 2024-11-01 at 18 08 29"
src="https://github.com/user-attachments/assets/864b5552-aece-4cc6-848a-8f73f88f55dc">

### Testing notes

Create some indices with `semantic_text` fields. For example, you can do
that via uploading and indexing a PDF file:

1. Navigate to Integrations page
2. Select "Upload a file"
3. Select and upload a PDF file
4. Press Import button
5. Switch to Advanced tab
6. Fill in "Index name"
7. Add additional field > Add semantic text field > Fill in form
  * Field: `attachment.content`
  * Copy to field: `content`
  * Inference endpoint: `elser_model_2`
8. Press Add button
9. Press Import button

### Checklist

Delete any items that are not applicable to this PR.

- [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
  • Loading branch information
e40pud authored Nov 5, 2024
1 parent 77b9d24 commit b122722
Show file tree
Hide file tree
Showing 23 changed files with 738 additions and 17 deletions.
2 changes: 2 additions & 0 deletions x-pack/packages/kbn-elastic-assistant-common/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND =
`${ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL}/_find` as const;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION =
`${ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL}/_bulk_action` as const;
export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_INDICES_URL =
`${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/knowledge_base/_indices` as const;

export const ELASTIC_AI_ASSISTANT_EVALUATE_URL =
`${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/evaluate` as const;
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ export * from './actions_connector/post_actions_connector_execute_route.gen';

// Knowledge Base Schemas
export * from './knowledge_base/crud_kb_route.gen';
export * from './knowledge_base/get_knowledge_base_indices_route.gen';
export * from './knowledge_base/entries/bulk_crud_knowledge_base_entries_route.gen';
export * from './knowledge_base/entries/common_attributes.gen';
export * from './knowledge_base/entries/crud_knowledge_base_entries_route.gen';
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
/*
* 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: Get Knowledge Base Indices API endpoints
* version: 1
*/

import { z } from '@kbn/zod';

export type GetKnowledgeBaseIndicesResponse = z.infer<typeof GetKnowledgeBaseIndicesResponse>;
export const GetKnowledgeBaseIndicesResponse = z.object({
/**
* List of indices with at least one field of a `sematic_text` type.
*/
indices: z.array(z.string()),
});
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
openapi: 3.0.0
info:
title: Get Knowledge Base Indices API endpoints
version: '1'
paths:
/internal/elastic_assistant/knowledge_base/_indices:
get:
x-codegen-enabled: true
x-labels: [ess, serverless]
operationId: GetKnowledgeBaseIndices
description: Gets Knowledge Base indices that have fields of a `sematic_text` type.
summary: Gets Knowledge Base indices that have fields of a `sematic_text` type.
tags:
- KnowledgeBase API
responses:
200:
description: Indicates a successful call.
content:
application/json:
schema:
type: object
properties:
indices:
type: array
description: List of indices with at least one field of a `sematic_text` type.
items:
type: string
required:
- indices
400:
description: Generic Error
content:
application/json:
schema:
type: object
properties:
statusCode:
type: number
error:
type: string
message:
type: string
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,12 @@

import { HttpSetup } from '@kbn/core-http-browser';

import { deleteKnowledgeBase, getKnowledgeBaseStatus, postKnowledgeBase } from './api';
import {
deleteKnowledgeBase,
getKnowledgeBaseIndices,
getKnowledgeBaseStatus,
postKnowledgeBase,
} from './api';

jest.mock('@kbn/core-http-browser');

Expand Down Expand Up @@ -95,4 +100,29 @@ describe('API tests', () => {
await expect(deleteKnowledgeBase(knowledgeBaseArgs)).resolves.toThrowError('simulated error');
});
});

describe('getKnowledgeBaseIndices', () => {
it('calls the knowledge base API when correct resource path', async () => {
await getKnowledgeBaseIndices({ http: mockHttp });

expect(mockHttp.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/_indices',
{
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(getKnowledgeBaseIndices({ http: mockHttp })).resolves.toThrowError(
'simulated error'
);
});
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@ import {
CreateKnowledgeBaseResponse,
DeleteKnowledgeBaseRequestParams,
DeleteKnowledgeBaseResponse,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_INDICES_URL,
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL,
GetKnowledgeBaseIndicesResponse,
ReadKnowledgeBaseRequestParams,
ReadKnowledgeBaseResponse,
} from '@kbn/elastic-assistant-common';
Expand Down Expand Up @@ -108,3 +110,32 @@ export const deleteKnowledgeBase = async ({
return error as IHttpFetchError;
}
};

/**
* API call for getting indices that have fields of `semantic_text` type.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {AbortSignal} [options.signal] - AbortSignal
*
* @returns {Promise<GetKnowledgeBaseIndicesResponse | IHttpFetchError>}
*/
export const getKnowledgeBaseIndices = async ({
http,
signal,
}: {
http: HttpSetup;
signal?: AbortSignal | undefined;
}): Promise<GetKnowledgeBaseIndicesResponse | IHttpFetchError> => {
try {
const response = await http.fetch(ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_INDICES_URL, {
method: 'GET',
signal,
version: API_VERSIONS.internal.v1,
});

return response as GetKnowledgeBaseIndicesResponse;
} catch (error) {
return error as IHttpFetchError;
}
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
/*
* 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 { act, renderHook } from '@testing-library/react-hooks';
import {
useKnowledgeBaseIndices,
UseKnowledgeBaseIndicesParams,
} from './use_knowledge_base_indices';
import { getKnowledgeBaseIndices as _getKnowledgeBaseIndices } from './api';

const getKnowledgeBaseIndicesMock = _getKnowledgeBaseIndices as jest.Mock;

jest.mock('./api', () => {
const actual = jest.requireActual('./api');
return {
...actual,
getKnowledgeBaseIndices: jest.fn((...args) => actual.getKnowledgeBaseIndices(...args)),
};
});

jest.mock('@tanstack/react-query', () => ({
useQuery: jest.fn().mockImplementation(async (queryKey, fn, opts) => {
try {
const res = await fn({});
return Promise.resolve(res);
} catch (e) {
opts.onError(e);
}
}),
}));

const indicesResponse = ['index-1', 'index-2', 'index-3'];

const http = {
fetch: jest.fn().mockResolvedValue(indicesResponse),
};
const toasts = {
addError: jest.fn(),
};
const defaultProps = { http, toasts } as unknown as UseKnowledgeBaseIndicesParams;
describe('useKnowledgeBaseIndices', () => {
beforeEach(() => {
jest.clearAllMocks();
});

it('should call api to get knowledge base indices', async () => {
await act(async () => {
const { waitForNextUpdate } = renderHook(() => useKnowledgeBaseIndices(defaultProps));
await waitForNextUpdate();

expect(defaultProps.http.fetch).toHaveBeenCalledWith(
'/internal/elastic_assistant/knowledge_base/_indices',
{
method: 'GET',
signal: undefined,
version: '1',
}
);
expect(toasts.addError).not.toHaveBeenCalled();
});
});

it('should return indices response', async () => {
await act(async () => {
const { result, waitForNextUpdate } = renderHook(() => useKnowledgeBaseIndices(defaultProps));
await waitForNextUpdate();

await expect(result.current).resolves.toStrictEqual(indicesResponse);
});
});

it('should display error toast when api throws error', async () => {
getKnowledgeBaseIndicesMock.mockRejectedValue(new Error('this is an error'));
await act(async () => {
const { waitForNextUpdate } = renderHook(() => useKnowledgeBaseIndices(defaultProps));
await waitForNextUpdate();

expect(toasts.addError).toHaveBeenCalled();
});
});
});
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

import 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 { GetKnowledgeBaseIndicesResponse } from '@kbn/elastic-assistant-common';
import { getKnowledgeBaseIndices } from './api';

const KNOWLEDGE_BASE_INDICES_QUERY_KEY = ['elastic-assistant', 'knowledge-base-indices'];

export interface UseKnowledgeBaseIndicesParams {
http: HttpSetup;
toasts?: IToasts;
}

/**
* Hook for getting indices that have fields of `semantic_text` type.
*
* @param {Object} options - The options object.
* @param {HttpSetup} options.http - HttpSetup
* @param {IToasts} [options.toasts] - IToasts
*
* @returns {useQuery} hook for getting indices that have fields of `semantic_text` type
*/
export const useKnowledgeBaseIndices = ({
http,
toasts,
}: UseKnowledgeBaseIndicesParams): UseQueryResult<
GetKnowledgeBaseIndicesResponse,
IHttpFetchError
> => {
return useQuery(
KNOWLEDGE_BASE_INDICES_QUERY_KEY,
async ({ signal }) => {
return getKnowledgeBaseIndices({ http, 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.indicesError', {
defaultMessage: 'Error fetching Knowledge Base Indices',
}),
}
);
}
},
}
);
};
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import { MOCK_QUICK_PROMPTS } from '../../mock/quick_prompt';
import { useAssistantContext } from '../../..';
import { I18nProvider } from '@kbn/i18n-react';
import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
import { useKnowledgeBaseIndices } from '../../assistant/api/knowledge_base/use_knowledge_base_indices';

const mockContext = {
basePromptContexts: MOCK_QUICK_PROMPTS,
Expand All @@ -44,13 +45,13 @@ jest.mock('../../assistant/api/knowledge_base/entries/use_update_knowledge_base_
jest.mock('../../assistant/api/knowledge_base/entries/use_delete_knowledge_base_entries');

jest.mock('../../assistant/settings/use_settings_updater/use_settings_updater');
jest.mock('../../assistant/api/knowledge_base/use_knowledge_base_indices');
jest.mock('../../assistant/api/knowledge_base/use_knowledge_base_status');
jest.mock('../../assistant/api/knowledge_base/entries/use_knowledge_base_entries');
jest.mock(
'../../assistant/common/components/assistant_settings_management/flyout/use_flyout_modal_visibility'
);
const mockDataViews = {
getIndices: jest.fn().mockResolvedValue([{ name: 'index-1' }, { name: 'index-2' }]),
getFieldsForWildcard: jest.fn().mockResolvedValue([
{ name: 'field-1', esTypes: ['semantic_text'] },
{ name: 'field-2', esTypes: ['text'] },
Expand Down Expand Up @@ -148,6 +149,9 @@ describe('KnowledgeBaseSettingsManagement', () => {
},
isFetched: true,
});
(useKnowledgeBaseIndices as jest.Mock).mockReturnValue({
data: { indices: ['index-1', 'index-2'] },
});
(useKnowledgeBaseEntries as jest.Mock).mockReturnValue({
data: { data: mockData },
isFetching: false,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,7 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
/>
) : (
<IndexEntryEditor
http={http}
entry={selectedEntry as IndexEntry}
originalEntry={originalEntry as IndexEntry}
dataViews={dataViews}
Expand Down
Loading

0 comments on commit b122722

Please sign in to comment.