mirror of
https://github.com/elastic/kibana.git
synced 2025-04-24 17:59:23 -04:00
[8.18] [Security Assistant] Fix initialization of Knowledge Base on undersized clusters (#212167) (#212809)
# Backport This will backport the following commits from `main` to `8.18`: - [[Security Assistant] Fix initialization of Knowledge Base on undersized clusters (#212167)](https://github.com/elastic/kibana/pull/212167) <!--- Backport version: 9.6.6 --> ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport) <!--BACKPORT [{"author":{"name":"Patryk Kopyciński","email":"contact@patrykkopycinski.com"},"sourceCommit":{"committedDate":"2025-02-28T20:42:04Z","message":"[Security Assistant] Fix initialization of Knowledge Base on undersized clusters (#212167)\n\n## Summary\n\nShow error to the user when trying to setup Knowledge base on undersized\ncluster\n\n<img width=\"1847\" alt=\"Zrzut ekranu 2025-02-26 o 19 03 43\"\nsrc=\"https://github.com/user-attachments/assets/a42d8560-aebb-410e-a364-7a27074f62fc\"\n/>\n\n---------\n\nCo-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>\nCo-authored-by: Garrett Spong <spong@users.noreply.github.com>\nCo-authored-by: Garrett Spong <garrett.spong@elastic.co>","sha":"b5caf904e775d32f8964dde8a407a3ca555b3f38","branchLabelMapping":{"^v9.1.0$":"main","^v8.19.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["bug","release_note:skip","v9.0.0","ci:cloud-deploy","ci:project-deploy-security","backport:version","v8.18.0","v9.1.0","v8.19.0"],"title":"[Security Assistant] Fix initialization of Knowledge Base on undersized clusters","number":212167,"url":"https://github.com/elastic/kibana/pull/212167","mergeCommit":{"message":"[Security Assistant] Fix initialization of Knowledge Base on undersized clusters (#212167)\n\n## Summary\n\nShow error to the user when trying to setup Knowledge base on undersized\ncluster\n\n<img width=\"1847\" alt=\"Zrzut ekranu 2025-02-26 o 19 03 43\"\nsrc=\"https://github.com/user-attachments/assets/a42d8560-aebb-410e-a364-7a27074f62fc\"\n/>\n\n---------\n\nCo-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>\nCo-authored-by: Garrett Spong <spong@users.noreply.github.com>\nCo-authored-by: Garrett Spong <garrett.spong@elastic.co>","sha":"b5caf904e775d32f8964dde8a407a3ca555b3f38"}},"sourceBranch":"main","suggestedTargetBranches":["9.0","8.18","8.x"],"targetPullRequestStates":[{"branch":"9.0","label":"v9.0.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.18","label":"v8.18.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v9.1.0","branchLabelMappingKey":"^v9.1.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/212167","number":212167,"mergeCommit":{"message":"[Security Assistant] Fix initialization of Knowledge Base on undersized clusters (#212167)\n\n## Summary\n\nShow error to the user when trying to setup Knowledge base on undersized\ncluster\n\n<img width=\"1847\" alt=\"Zrzut ekranu 2025-02-26 o 19 03 43\"\nsrc=\"https://github.com/user-attachments/assets/a42d8560-aebb-410e-a364-7a27074f62fc\"\n/>\n\n---------\n\nCo-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>\nCo-authored-by: Garrett Spong <spong@users.noreply.github.com>\nCo-authored-by: Garrett Spong <garrett.spong@elastic.co>","sha":"b5caf904e775d32f8964dde8a407a3ca555b3f38"}},{"branch":"8.x","label":"v8.19.0","branchLabelMappingKey":"^v8.19.0$","isSourceBranch":false,"state":"NOT_CREATED"}]}] BACKPORT--> Co-authored-by: Patryk Kopyciński <contact@patrykkopycinski.com>
This commit is contained in:
Kibana Machine
GitHub
parent
9f219c92f3
commit
2f02a2a050
11 changed files with 246 additions and 149 deletions
|
|
@ -79,7 +79,7 @@ export const useInvalidateKnowledgeBaseStatus = () => {
|
|||
|
||||
return useCallback(() => {
|
||||
queryClient.invalidateQueries(KNOWLEDGE_BASE_STATUS_QUERY_KEY, {
|
||||
refetchType: 'active',
|
||||
refetchType: 'all',
|
||||
});
|
||||
}, [queryClient]);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -75,6 +75,7 @@ const createKnowledgeBaseDataClientMock = () => {
|
|||
isSetupAvailable: jest.fn(),
|
||||
isUserDataExists: jest.fn(),
|
||||
setupKnowledgeBase: jest.fn(),
|
||||
getLoadedSecurityLabsDocsCount: jest.fn(),
|
||||
};
|
||||
return mocked;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@
|
|||
*/
|
||||
import { FieldMap } from '@kbn/data-stream-adapter';
|
||||
|
||||
export const ELSER_MODEL_2 = ['.elser_model_2', '.elser_model_2_linux-x86_64'];
|
||||
export const ASSISTANT_ELSER_INFERENCE_ID = 'elastic-security-ai-assistant-elser2';
|
||||
export const ELASTICSEARCH_ELSER_INFERENCE_ID = '.elser-2-elasticsearch';
|
||||
|
||||
|
|
|
|||
|
|
@ -4,12 +4,7 @@
|
|||
* 2.0; you may not use this file except in compliance with the Elastic License
|
||||
* 2.0.
|
||||
*/
|
||||
import {
|
||||
coreMock,
|
||||
elasticsearchServiceMock,
|
||||
loggingSystemMock,
|
||||
savedObjectsRepositoryMock,
|
||||
} from '@kbn/core/server/mocks';
|
||||
import { coreMock, elasticsearchServiceMock, loggingSystemMock } from '@kbn/core/server/mocks';
|
||||
import { AIAssistantKnowledgeBaseDataClient, KnowledgeBaseDataClientParams } from '.';
|
||||
import {
|
||||
getCreateKnowledgeBaseEntrySchemaMock,
|
||||
|
|
@ -19,6 +14,7 @@ import {
|
|||
import { authenticatedUser } from '../../__mocks__/user';
|
||||
import { IndexPatternsFetcher } from '@kbn/data-plugin/server';
|
||||
import type { MlPluginSetup } from '@kbn/ml-plugin/server';
|
||||
import { getMlNodeCount } from '@kbn/ml-plugin/server/lib/node_utils';
|
||||
import { mlPluginMock } from '@kbn/ml-plugin/public/mocks';
|
||||
import pRetry from 'p-retry';
|
||||
|
||||
|
|
@ -29,6 +25,10 @@ import {
|
|||
import { DynamicStructuredTool } from '@langchain/core/tools';
|
||||
import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock';
|
||||
import { KnowledgeBaseResource } from '@kbn/elastic-assistant-common';
|
||||
import { createTrainedModelsProviderMock } from '@kbn/ml-plugin/server/shared_services/providers/__mocks__/trained_models';
|
||||
import { ASSISTANT_ELSER_INFERENCE_ID } from './field_maps_configuration';
|
||||
|
||||
jest.mock('@kbn/ml-plugin/server/lib/node_utils');
|
||||
jest.mock('../../lib/langchain/content_loaders/security_labs_loader');
|
||||
jest.mock('p-retry');
|
||||
const date = '2023-03-28T22:27:28.159Z';
|
||||
|
|
@ -44,23 +44,16 @@ const telemetry = coreMock.createSetup().analytics;
|
|||
describe('AIAssistantKnowledgeBaseDataClient', () => {
|
||||
let mockOptions: KnowledgeBaseDataClientParams;
|
||||
let ml: MlPluginSetup;
|
||||
let savedObjectClient: ReturnType<typeof savedObjectsRepositoryMock.create>;
|
||||
const getElserId = jest.fn();
|
||||
const trainedModelsProvider = jest.fn();
|
||||
const installElasticModel = jest.fn();
|
||||
const mockLoadSecurityLabs = loadSecurityLabs as jest.Mock;
|
||||
const mockGetSecurityLabsDocsCount = getSecurityLabsDocsCount as jest.Mock;
|
||||
const mockGetIsKBSetupInProgress = jest.fn();
|
||||
const trainedModelsProviderMock = createTrainedModelsProviderMock()();
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
logger = loggingSystemMock.createLogger();
|
||||
savedObjectClient = savedObjectsRepositoryMock.create();
|
||||
mockLoadSecurityLabs.mockClear();
|
||||
ml = mlPluginMock.createSetupContract() as unknown as MlPluginSetup; // Missing SharedServices mock, so manually mocking trainedModelsProvider
|
||||
ml.trainedModelsProvider = trainedModelsProvider.mockImplementation(() => ({
|
||||
getELSER: jest.fn().mockImplementation(() => '.elser_model_2'),
|
||||
installElasticModel: installElasticModel.mockResolvedValue({}),
|
||||
}));
|
||||
mockOptions = {
|
||||
logger,
|
||||
elasticsearchClientPromise: Promise.resolve(esClientMock),
|
||||
|
|
@ -75,6 +68,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
setIsKBSetupInProgress: jest.fn().mockImplementation(() => {}),
|
||||
manageGlobalKnowledgeBaseAIAssistant: true,
|
||||
assistantDefaultInferenceEndpoint: false,
|
||||
trainedModelsProvider: trainedModelsProviderMock,
|
||||
};
|
||||
esClientMock.search.mockReturnValue(
|
||||
// @ts-expect-error not full response interface
|
||||
|
|
@ -130,7 +124,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
describe('isModelInstalled', () => {
|
||||
it('should check if ELSER model is installed and return true if fully_defined', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModels.mockResolvedValue({
|
||||
trainedModelsProviderMock.getTrainedModels.mockResolvedValue({
|
||||
count: 1,
|
||||
trained_model_configs: [
|
||||
{ fully_defined: true, model_id: '', tags: [], input: { field_names: ['content'] } },
|
||||
|
|
@ -138,7 +132,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
});
|
||||
const result = await client.isModelInstalled();
|
||||
expect(result).toBe(true);
|
||||
expect(esClientMock.ml.getTrainedModels).toHaveBeenCalledWith({
|
||||
expect(trainedModelsProviderMock.getTrainedModels).toHaveBeenCalledWith({
|
||||
model_id: 'elser-id',
|
||||
include: 'definition_status',
|
||||
});
|
||||
|
|
@ -146,7 +140,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
it('should return false if model is not fully defined', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModels.mockResolvedValue({
|
||||
trainedModelsProviderMock.getTrainedModels.mockResolvedValue({
|
||||
count: 0,
|
||||
trained_model_configs: [
|
||||
{ fully_defined: false, model_id: '', tags: [], input: { field_names: ['content'] } },
|
||||
|
|
@ -158,7 +152,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
it('should return false and log error if getting model details fails', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModels.mockRejectedValue(new Error('error happened'));
|
||||
trainedModelsProviderMock.getTrainedModels.mockRejectedValue(new Error('error happened'));
|
||||
const result = await client.isModelInstalled();
|
||||
expect(result).toBe(false);
|
||||
expect(logger.error).toHaveBeenCalled();
|
||||
|
|
@ -168,12 +162,12 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
describe('isInferenceEndpointExists', () => {
|
||||
it('returns true when the model is fully allocated and started in ESS', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trained_model_stats: [
|
||||
{
|
||||
deployment_stats: {
|
||||
state: 'started',
|
||||
// @ts-expect-error not full response interface
|
||||
deployment_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
allocation_status: { state: 'fully_allocated' },
|
||||
},
|
||||
},
|
||||
|
|
@ -187,11 +181,11 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
it('returns true when the model is started in serverless', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trained_model_stats: [
|
||||
{
|
||||
deployment_stats: {
|
||||
// @ts-expect-error not full response interface
|
||||
deployment_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
nodes: [{ routing_state: { routing_state: 'started' } }],
|
||||
},
|
||||
},
|
||||
|
|
@ -205,12 +199,12 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
it('returns false when the model is not fully allocated in ESS', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trained_model_stats: [
|
||||
{
|
||||
deployment_stats: {
|
||||
state: 'started',
|
||||
// @ts-expect-error not full response interface
|
||||
deployment_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
allocation_status: { state: 'partially_allocated' },
|
||||
},
|
||||
},
|
||||
|
|
@ -224,11 +218,11 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
it('returns false when the model is not started in serverless', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockResolvedValueOnce({
|
||||
trained_model_stats: [
|
||||
{
|
||||
deployment_stats: {
|
||||
// @ts-expect-error not full response interface
|
||||
deployment_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
nodes: [{ routing_state: { routing_state: 'stopped' } }],
|
||||
},
|
||||
},
|
||||
|
|
@ -242,7 +236,9 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
it('returns false when an error occurs during the check', async () => {
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
esClientMock.ml.getTrainedModelsStats.mockRejectedValueOnce(new Error('Mocked Error'));
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockRejectedValueOnce(
|
||||
new Error('Mocked Error')
|
||||
);
|
||||
|
||||
const result = await client.isInferenceEndpointExists();
|
||||
|
||||
|
|
@ -267,33 +263,49 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
|
||||
describe('setupKnowledgeBase', () => {
|
||||
it('should install, deploy, and load docs if not already done', async () => {
|
||||
(getMlNodeCount as jest.Mock).mockResolvedValue({ count: 1, lazyNodeCount: 0 });
|
||||
// @ts-expect-error not full response interface
|
||||
esClientMock.search.mockResolvedValue({});
|
||||
trainedModelsProviderMock.startTrainedModelDeployment.mockResolvedValue({});
|
||||
trainedModelsProviderMock.stopTrainedModelDeployment.mockResolvedValue({});
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockResolvedValue({
|
||||
trained_model_stats: [
|
||||
{
|
||||
deployment_stats: {
|
||||
state: 'started',
|
||||
deployment_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
allocation_status: {
|
||||
state: 'fully_allocated',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
await client.setupKnowledgeBase({ soClient: savedObjectClient });
|
||||
await client.setupKnowledgeBase({});
|
||||
|
||||
// install model
|
||||
expect(trainedModelsProvider).toHaveBeenCalledWith({}, savedObjectClient);
|
||||
expect(installElasticModel).toHaveBeenCalledWith('elser-id');
|
||||
expect(trainedModelsProviderMock.installElasticModel).toHaveBeenCalledWith('elser-id');
|
||||
|
||||
expect(loadSecurityLabs).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip installation and deployment if model is already installed and deployed', async () => {
|
||||
(getMlNodeCount as jest.Mock).mockResolvedValue({ count: 1, lazyNodeCount: 0 });
|
||||
mockGetSecurityLabsDocsCount.mockResolvedValue(1);
|
||||
esClientMock.ml.getTrainedModels.mockResolvedValue({
|
||||
trainedModelsProviderMock.getTrainedModels.mockResolvedValue({
|
||||
count: 1,
|
||||
trained_model_configs: [
|
||||
{ fully_defined: true, model_id: '', tags: [], input: { field_names: ['content'] } },
|
||||
],
|
||||
});
|
||||
esClientMock.ml.getTrainedModelsStats.mockResolvedValue({
|
||||
trainedModelsProviderMock.getTrainedModelsStats.mockResolvedValue({
|
||||
trained_model_stats: [
|
||||
{
|
||||
deployment_stats: {
|
||||
deployment_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
state: 'started',
|
||||
// @ts-expect-error not full response interface
|
||||
allocation_status: {
|
||||
state: 'fully_allocated',
|
||||
},
|
||||
|
|
@ -303,17 +315,17 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
});
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
|
||||
await client.setupKnowledgeBase({ soClient: savedObjectClient });
|
||||
await client.setupKnowledgeBase({});
|
||||
|
||||
expect(installElasticModel).not.toHaveBeenCalled();
|
||||
expect(esClientMock.ml.startTrainedModelDeployment).not.toHaveBeenCalled();
|
||||
expect(trainedModelsProviderMock.installElasticModel).not.toHaveBeenCalled();
|
||||
expect(trainedModelsProviderMock.startTrainedModelDeployment).not.toHaveBeenCalled();
|
||||
expect(loadSecurityLabs).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle errors during installation and deployment', async () => {
|
||||
// @ts-expect-error not full response interface
|
||||
esClientMock.search.mockResolvedValue({});
|
||||
esClientMock.ml.getTrainedModels.mockResolvedValue({
|
||||
trainedModelsProviderMock.getTrainedModels.mockResolvedValue({
|
||||
count: 0,
|
||||
trained_model_configs: [
|
||||
{ fully_defined: false, model_id: '', tags: [], input: { field_names: ['content'] } },
|
||||
|
|
@ -322,7 +334,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
|
|||
mockLoadSecurityLabs.mockRejectedValue(new Error('Installation error'));
|
||||
const client = new AIAssistantKnowledgeBaseDataClient(mockOptions);
|
||||
|
||||
await expect(client.setupKnowledgeBase({ soClient: savedObjectClient })).rejects.toThrow(
|
||||
await expect(client.setupKnowledgeBase({})).rejects.toThrow(
|
||||
'Error setting up Knowledge Base: Installation error'
|
||||
);
|
||||
expect(mockOptions.logger.error).toHaveBeenCalledWith(
|
||||
|
|
|
|||
|
|
@ -12,9 +12,7 @@ import {
|
|||
QueryDslQueryContainer,
|
||||
} from '@elastic/elasticsearch/lib/api/types';
|
||||
import type { MlPluginSetup } from '@kbn/ml-plugin/server';
|
||||
import type { KibanaRequest } from '@kbn/core-http-server';
|
||||
import { Document } from 'langchain/document';
|
||||
import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server';
|
||||
import {
|
||||
DocumentEntryType,
|
||||
DocumentEntry,
|
||||
|
|
@ -27,9 +25,16 @@ import {
|
|||
} from '@kbn/elastic-assistant-common';
|
||||
import pRetry from 'p-retry';
|
||||
import { StructuredTool } from '@langchain/core/tools';
|
||||
import { AnalyticsServiceSetup, AuditLogger, ElasticsearchClient } from '@kbn/core/server';
|
||||
import {
|
||||
AnalyticsServiceSetup,
|
||||
AuditLogger,
|
||||
ElasticsearchClient,
|
||||
IScopedClusterClient,
|
||||
} from '@kbn/core/server';
|
||||
import { IndexPatternsFetcher } from '@kbn/data-views-plugin/server';
|
||||
import { map } from 'lodash';
|
||||
import type { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers';
|
||||
import { getMlNodeCount } from '@kbn/ml-plugin/server/lib/node_utils';
|
||||
import { AIAssistantDataClient, AIAssistantDataClientParams } from '..';
|
||||
import { GetElser } from '../../types';
|
||||
import {
|
||||
|
|
@ -63,6 +68,7 @@ import {
|
|||
import {
|
||||
ASSISTANT_ELSER_INFERENCE_ID,
|
||||
ELASTICSEARCH_ELSER_INFERENCE_ID,
|
||||
ELSER_MODEL_2,
|
||||
} from './field_maps_configuration';
|
||||
import { BulkOperationError } from '../../lib/data_stream/documents_data_writer';
|
||||
import { AUDIT_OUTCOME, KnowledgeBaseAuditAction, knowledgeBaseAuditEvent } from './audit_events';
|
||||
|
|
@ -79,11 +85,12 @@ export interface GetAIAssistantKnowledgeBaseDataClientParams {
|
|||
export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams {
|
||||
ml: MlPluginSetup;
|
||||
getElserId: GetElser;
|
||||
getIsKBSetupInProgress: () => boolean;
|
||||
getIsKBSetupInProgress: (spaceId: string) => boolean;
|
||||
ingestPipelineResourceName: string;
|
||||
setIsKBSetupInProgress: (isInProgress: boolean) => void;
|
||||
setIsKBSetupInProgress: (spaceId: string, isInProgress: boolean) => void;
|
||||
manageGlobalKnowledgeBaseAIAssistant: boolean;
|
||||
assistantDefaultInferenceEndpoint: boolean;
|
||||
trainedModelsProvider: ReturnType<TrainedModelsProvider['trainedModelsProvider']>;
|
||||
}
|
||||
export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
||||
constructor(public readonly options: KnowledgeBaseDataClientParams) {
|
||||
|
|
@ -91,7 +98,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
}
|
||||
|
||||
public get isSetupInProgress() {
|
||||
return this.options.getIsKBSetupInProgress();
|
||||
return this.options.getIsKBSetupInProgress(this.spaceId);
|
||||
}
|
||||
/**
|
||||
* Returns whether setup of the Knowledge Base can be performed (essentially an ML features check)
|
||||
|
|
@ -112,18 +119,13 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
|
||||
/**
|
||||
* Downloads and installs ELSER model if not already installed
|
||||
*
|
||||
* @param soClient SavedObjectsClientContract for installing ELSER so that ML SO's are in sync
|
||||
*/
|
||||
private installModel = async ({ soClient }: { soClient: SavedObjectsClientContract }) => {
|
||||
private installModel = async () => {
|
||||
const elserId = await this.options.getElserId();
|
||||
this.options.logger.debug(`Installing ELSER model '${elserId}'...`);
|
||||
|
||||
try {
|
||||
await this.options.ml
|
||||
// TODO: Potentially plumb soClient through DataClient from pluginStart
|
||||
.trainedModelsProvider({} as KibanaRequest, soClient)
|
||||
.installElasticModel(elserId);
|
||||
await this.options.trainedModelsProvider.installElasticModel(elserId);
|
||||
} catch (error) {
|
||||
this.options.logger.error(`Error installing ELSER model '${elserId}':\n${error}`);
|
||||
}
|
||||
|
|
@ -139,8 +141,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
this.options.logger.debug(`Checking if ELSER model '${elserId}' is installed...`);
|
||||
|
||||
try {
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
const getResponse = await esClient.ml.getTrainedModels({
|
||||
const getResponse = await this.options.trainedModelsProvider.getTrainedModels({
|
||||
model_id: elserId,
|
||||
include: 'definition_status',
|
||||
});
|
||||
|
|
@ -156,7 +157,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
};
|
||||
|
||||
public getInferenceEndpointId = async () => {
|
||||
if (!this.options.assistantDefaultInferenceEndpoint) {
|
||||
const elserId = await this.options.getElserId();
|
||||
if (!this.options.assistantDefaultInferenceEndpoint || !ELSER_MODEL_2.includes(elserId)) {
|
||||
return ASSISTANT_ELSER_INFERENCE_ID;
|
||||
}
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
|
|
@ -186,6 +188,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
* @returns Promise<boolean> indicating whether the model is deployed
|
||||
*/
|
||||
public isInferenceEndpointExists = async (inferenceEndpointId?: string): Promise<boolean> => {
|
||||
const elserId = await this.options.getElserId();
|
||||
const inferenceId = inferenceEndpointId || (await this.getInferenceEndpointId());
|
||||
|
||||
try {
|
||||
|
|
@ -195,13 +198,19 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
inference_id: inferenceId,
|
||||
task_type: 'sparse_embedding',
|
||||
}));
|
||||
|
||||
if (!inferenceExists) {
|
||||
return false;
|
||||
}
|
||||
const elserId = await this.options.getElserId();
|
||||
const getResponse = await esClient.ml.getTrainedModelsStats({
|
||||
model_id: elserId,
|
||||
});
|
||||
|
||||
let getResponse;
|
||||
try {
|
||||
getResponse = await this.options.trainedModelsProvider.getTrainedModelsStats({
|
||||
model_id: elserId,
|
||||
});
|
||||
} catch (e) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986
|
||||
const isReadyESS = (stats: MlTrainedModelStats) =>
|
||||
|
|
@ -213,9 +222,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
(node) => node.routing_state.routing_state === 'started'
|
||||
);
|
||||
|
||||
return !!getResponse.trained_model_stats?.some(
|
||||
(stats) => isReadyESS(stats) || isReadyServerless(stats)
|
||||
);
|
||||
return !!getResponse.trained_model_stats
|
||||
.filter((stats) => stats.deployment_stats?.deployment_id === inferenceId)
|
||||
?.some((stats) => isReadyESS(stats) || isReadyServerless(stats));
|
||||
} catch (error) {
|
||||
this.options.logger.debug(
|
||||
`Error checking if Inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} exists: ${error}`
|
||||
|
|
@ -224,31 +233,62 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
}
|
||||
};
|
||||
|
||||
private dryRunTrainedModelDeployment = async () => {
|
||||
const elserId = await this.options.getElserId();
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
|
||||
try {
|
||||
// As there is no better way to check if the model is deployed, we try to start the model
|
||||
// deployment and throw an error if it fails
|
||||
const dryRunId = await esClient.ml.startTrainedModelDeployment({
|
||||
model_id: elserId,
|
||||
wait_for: 'fully_allocated',
|
||||
});
|
||||
this.options.logger.debug(`Dry run for ELSER model '${elserId}' successfully deployed!`);
|
||||
|
||||
await this.options.trainedModelsProvider.stopTrainedModelDeployment({
|
||||
model_id: elserId,
|
||||
deployment_id: dryRunId.assignment.task_parameters.deployment_id,
|
||||
});
|
||||
this.options.logger.debug(`Dry run for ELSER model '${elserId}' successfully stopped!`);
|
||||
} catch (e) {
|
||||
this.options.logger.error(`Dry run error starting trained model deployment: ${e.message}`);
|
||||
throw new Error(`${e.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
private deleteInferenceEndpoint = async () => {
|
||||
const elserId = await this.options.getElserId();
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
|
||||
try {
|
||||
await esClient.inference.delete({
|
||||
inference_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
// it's being used in the mapping so we need to force delete
|
||||
force: true,
|
||||
});
|
||||
this.options.logger.debug(`Deleted existing inference endpoint for ELSER model '${elserId}'`);
|
||||
} catch (error) {
|
||||
this.options.logger.error(
|
||||
`Error deleting inference endpoint ${ASSISTANT_ELSER_INFERENCE_ID} for ELSER model '${elserId}':\n${error}`
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
public createInferenceEndpoint = async () => {
|
||||
const elserId = await this.options.getElserId();
|
||||
this.options.logger.debug(`Deploying ELSER model '${elserId}'...`);
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
const inferenceId = await this.getInferenceEndpointId();
|
||||
const inferenceExists = await this.isInferenceEndpointExists(inferenceId);
|
||||
|
||||
// Don't try to create the inference endpoint for ELASTICSEARCH_ELSER_INFERENCE_ID
|
||||
if (inferenceId === ASSISTANT_ELSER_INFERENCE_ID) {
|
||||
if (inferenceExists) {
|
||||
try {
|
||||
await esClient.inference.delete({
|
||||
inference_id: ASSISTANT_ELSER_INFERENCE_ID,
|
||||
// it's being used in the mapping so we need to force delete
|
||||
force: true,
|
||||
});
|
||||
this.options.logger.debug(
|
||||
`Deleted existing inference endpoint for ELSER model '${elserId}'`
|
||||
);
|
||||
} catch (error) {
|
||||
this.options.logger.error(
|
||||
`Error deleting inference endpoint for ELSER model '${elserId}':\n${error}`
|
||||
);
|
||||
}
|
||||
}
|
||||
await this.deleteInferenceEndpoint();
|
||||
|
||||
await pRetry(async () => this.dryRunTrainedModelDeployment(), {
|
||||
minTimeout: 10000,
|
||||
maxTimeout: 10000,
|
||||
retries: 1,
|
||||
});
|
||||
|
||||
try {
|
||||
await esClient.inference.put({
|
||||
|
|
@ -270,7 +310,14 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
});
|
||||
|
||||
// await for the model to be deployed
|
||||
await this.isInferenceEndpointExists(inferenceId);
|
||||
const inferenceEndpointExists = await this.isInferenceEndpointExists(
|
||||
ASSISTANT_ELSER_INFERENCE_ID
|
||||
);
|
||||
if (!inferenceEndpointExists) {
|
||||
throw new Error(
|
||||
`Inference endpoint for ELSER model '${elserId}' was not deployed successfully`
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
this.options.logger.error(
|
||||
`Error creating inference endpoint for ELSER model '${elserId}':\n${error}`
|
||||
|
|
@ -279,6 +326,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
`Error creating inference endpoint for ELSER model '${elserId}':\n${error}`
|
||||
);
|
||||
}
|
||||
} else {
|
||||
await this.dryRunTrainedModelDeployment();
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -295,24 +344,30 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
* @returns Promise<void>
|
||||
*/
|
||||
public setupKnowledgeBase = async ({
|
||||
soClient,
|
||||
ignoreSecurityLabs = false,
|
||||
}: {
|
||||
soClient: SavedObjectsClientContract;
|
||||
ignoreSecurityLabs?: boolean;
|
||||
}): Promise<void> => {
|
||||
if (this.options.getIsKBSetupInProgress()) {
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
|
||||
if (this.options.getIsKBSetupInProgress(this.spaceId)) {
|
||||
this.options.logger.debug('Knowledge Base setup already in progress');
|
||||
return;
|
||||
}
|
||||
|
||||
this.options.logger.debug('Checking if ML nodes are available...');
|
||||
const mlNodesCount = await getMlNodeCount({ asInternalUser: esClient } as IScopedClusterClient);
|
||||
|
||||
if (mlNodesCount.count === 0 && mlNodesCount.lazyNodeCount === 0) {
|
||||
throw new Error('No ML nodes available');
|
||||
}
|
||||
|
||||
this.options.logger.debug('Starting Knowledge Base setup...');
|
||||
this.options.setIsKBSetupInProgress(true);
|
||||
this.options.setIsKBSetupInProgress(this.spaceId, true);
|
||||
const elserId = await this.options.getElserId();
|
||||
|
||||
// Delete legacy ESQL knowledge base docs if they exist, and silence the error if they do not
|
||||
try {
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
const legacyESQL = await esClient.deleteByQuery({
|
||||
index: this.indexTemplateAndPattern.alias,
|
||||
query: {
|
||||
|
|
@ -326,28 +381,21 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
`Removed ${legacyESQL?.total} ESQL knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.`
|
||||
);
|
||||
}
|
||||
// Delete any existing Security Labs content
|
||||
const securityLabsDocs = await esClient.deleteByQuery({
|
||||
index: this.indexTemplateAndPattern.alias,
|
||||
query: {
|
||||
bool: {
|
||||
must: [{ terms: { kb_resource: [SECURITY_LABS_RESOURCE] } }],
|
||||
},
|
||||
},
|
||||
});
|
||||
if (securityLabsDocs?.total) {
|
||||
this.options.logger.info(
|
||||
`Removed ${securityLabsDocs?.total} Security Labs knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.`
|
||||
);
|
||||
}
|
||||
} catch (e) {
|
||||
this.options.logger.info('No legacy ESQL or Security Labs knowledge base docs to delete');
|
||||
}
|
||||
|
||||
try {
|
||||
/*
|
||||
#1 Check if ELSER model is downloaded
|
||||
#2 Check if inference endpoint is deployed
|
||||
#3 Dry run ELSER model deployment if not already deployed
|
||||
#4 Create inference endpoint if not deployed / delete and create inference endpoint if model was not deployed
|
||||
#5 Load Security Labs docs
|
||||
*/
|
||||
const isInstalled = await this.isModelInstalled();
|
||||
if (!isInstalled) {
|
||||
await this.installModel({ soClient });
|
||||
await this.installModel();
|
||||
await pRetry(
|
||||
async () =>
|
||||
(await this.isModelInstalled())
|
||||
|
|
@ -373,11 +421,29 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
);
|
||||
}
|
||||
|
||||
this.options.logger.debug(`Checking if Knowledge Base docs have been loaded...`);
|
||||
|
||||
if (!ignoreSecurityLabs) {
|
||||
this.options.logger.debug(`Checking if Knowledge Base docs have been loaded...`);
|
||||
|
||||
const labsDocsLoaded = await this.isSecurityLabsDocsLoaded();
|
||||
if (!labsDocsLoaded) {
|
||||
// Delete any existing Security Labs content
|
||||
const securityLabsDocs = await (
|
||||
await this.options.elasticsearchClientPromise
|
||||
).deleteByQuery({
|
||||
index: this.indexTemplateAndPattern.alias,
|
||||
query: {
|
||||
bool: {
|
||||
must: [{ terms: { kb_resource: [SECURITY_LABS_RESOURCE] } }],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (securityLabsDocs?.total) {
|
||||
this.options.logger.info(
|
||||
`Removed ${securityLabsDocs?.total} Security Labs knowledge base docs from knowledge base data stream: ${this.indexTemplateAndPattern.alias}.`
|
||||
);
|
||||
}
|
||||
|
||||
this.options.logger.debug(`Loading Security Labs KB docs...`);
|
||||
await loadSecurityLabs(this, this.options.logger);
|
||||
} else {
|
||||
|
|
@ -385,11 +451,11 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
}
|
||||
}
|
||||
} catch (e) {
|
||||
this.options.setIsKBSetupInProgress(false);
|
||||
this.options.setIsKBSetupInProgress(this.spaceId, false);
|
||||
this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`);
|
||||
throw new Error(`Error setting up Knowledge Base: ${e.message}`);
|
||||
} finally {
|
||||
this.options.setIsKBSetupInProgress(false);
|
||||
this.options.setIsKBSetupInProgress(this.spaceId, false);
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -493,9 +559,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
};
|
||||
|
||||
/**
|
||||
* Returns if allSecurity Labs KB docs have been loaded
|
||||
* Returns loaded Security Labs KB docs count
|
||||
*/
|
||||
public isSecurityLabsDocsLoaded = async (): Promise<boolean> => {
|
||||
public getLoadedSecurityLabsDocsCount = async (): Promise<number> => {
|
||||
const user = this.options.currentUser;
|
||||
if (user == null) {
|
||||
throw new Error(
|
||||
|
|
@ -503,8 +569,6 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
);
|
||||
}
|
||||
|
||||
const expectedDocsCount = await getSecurityLabsDocsCount({ logger: this.options.logger });
|
||||
|
||||
const esClient = await this.options.elasticsearchClientPromise;
|
||||
|
||||
try {
|
||||
|
|
@ -521,7 +585,27 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
track_total_hits: true,
|
||||
});
|
||||
|
||||
const existingDocs = (result.hits?.total as SearchTotalHits).value;
|
||||
return (result.hits?.total as SearchTotalHits).value;
|
||||
} catch (e) {
|
||||
this.options.logger.info(`Error checking if Security Labs docs are loaded: ${e.message}`);
|
||||
return 0;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Returns if allSecurity Labs KB docs have been loaded
|
||||
*/
|
||||
public isSecurityLabsDocsLoaded = async (): Promise<boolean> => {
|
||||
const user = this.options.currentUser;
|
||||
if (user == null) {
|
||||
throw new Error(
|
||||
'Authenticated user not found! Ensure kbDataClient was initialized from a request.'
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
const expectedDocsCount = await getSecurityLabsDocsCount({ logger: this.options.logger });
|
||||
const existingDocs = await this.getLoadedSecurityLabsDocsCount();
|
||||
|
||||
if (existingDocs !== expectedDocsCount) {
|
||||
this.options.logger.debug(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import {
|
|||
IndicesIndexSettings,
|
||||
} from '@elastic/elasticsearch/lib/api/types';
|
||||
import { omit } from 'lodash';
|
||||
import { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers';
|
||||
import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration';
|
||||
import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration';
|
||||
import { getDefaultAnonymizationFields } from '../../common/anonymization';
|
||||
|
|
@ -101,7 +102,7 @@ export class AIAssistantService {
|
|||
private defendInsightsDataStream: DataStreamSpacesAdapter;
|
||||
private resourceInitializationHelper: ResourceInstallationHelper;
|
||||
private initPromise: Promise<InitializationPromise>;
|
||||
private isKBSetupInProgress: boolean = false;
|
||||
private isKBSetupInProgress: Map<string, boolean> = new Map();
|
||||
private hasInitializedV2KnowledgeBase: boolean = false;
|
||||
private productDocManager?: ProductDocBaseStartContract['management'];
|
||||
// Temporary 'feature flag' to determine if we should initialize the new knowledge base mappings
|
||||
|
|
@ -161,12 +162,12 @@ export class AIAssistantService {
|
|||
return this.initialized;
|
||||
}
|
||||
|
||||
public getIsKBSetupInProgress() {
|
||||
return this.isKBSetupInProgress;
|
||||
public getIsKBSetupInProgress(spaceId: string) {
|
||||
return this.isKBSetupInProgress.get(spaceId) ?? false;
|
||||
}
|
||||
|
||||
public setIsKBSetupInProgress(isInProgress: boolean) {
|
||||
this.isKBSetupInProgress = isInProgress;
|
||||
public setIsKBSetupInProgress(spaceId: string, isInProgress: boolean) {
|
||||
this.isKBSetupInProgress.set(spaceId, isInProgress);
|
||||
}
|
||||
|
||||
private createDataStream: CreateDataStream = ({
|
||||
|
|
@ -488,7 +489,10 @@ export class AIAssistantService {
|
|||
}
|
||||
|
||||
public async createAIAssistantKnowledgeBaseDataClient(
|
||||
opts: CreateAIAssistantClientParams & GetAIAssistantKnowledgeBaseDataClientParams
|
||||
opts: CreateAIAssistantClientParams &
|
||||
GetAIAssistantKnowledgeBaseDataClientParams & {
|
||||
trainedModelsProvider: ReturnType<TrainedModelsProvider['trainedModelsProvider']>;
|
||||
}
|
||||
): Promise<AIAssistantKnowledgeBaseDataClient | null> {
|
||||
// If modelIdOverride is set, swap getElserId(), and ensure the pipeline is re-created with the correct model
|
||||
if (opts?.modelIdOverride != null) {
|
||||
|
|
@ -525,6 +529,7 @@ export class AIAssistantService {
|
|||
spaceId: opts.spaceId,
|
||||
manageGlobalKnowledgeBaseAIAssistant: opts.manageGlobalKnowledgeBaseAIAssistant ?? false,
|
||||
assistantDefaultInferenceEndpoint: this.assistantDefaultInferenceEndpoint,
|
||||
trainedModelsProvider: opts.trainedModelsProvider,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ describe('Get Knowledge Base Status Route', () => {
|
|||
isSetupInProgress: false,
|
||||
isSecurityLabsDocsLoaded: jest.fn().mockResolvedValue(true),
|
||||
isUserDataExists: jest.fn().mockResolvedValue(true),
|
||||
getLoadedSecurityLabsDocsCount: jest.fn().mockResolvedValue(0),
|
||||
});
|
||||
|
||||
getKnowledgeBaseStatusRoute(server.router);
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ import {
|
|||
API_VERSIONS,
|
||||
ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL,
|
||||
ReadKnowledgeBaseRequestParams,
|
||||
ReadKnowledgeBaseResponse,
|
||||
} from '@kbn/elastic-assistant-common';
|
||||
import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
|
||||
import { KibanaRequest } from '@kbn/core/server';
|
||||
|
|
@ -57,32 +56,24 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter
|
|||
|
||||
const indexExists = true; // Installed at startup, always true
|
||||
const pipelineExists = true; // Installed at startup, always true
|
||||
const modelExists = await kbDataClient.isModelInstalled();
|
||||
const setupAvailable = await kbDataClient.isSetupAvailable();
|
||||
const isInferenceEndpointExists = await kbDataClient.isInferenceEndpointExists();
|
||||
const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded();
|
||||
const loadedSecurityLabsDocsCount = await kbDataClient.getLoadedSecurityLabsDocsCount();
|
||||
const userDataExists = await kbDataClient.isUserDataExists();
|
||||
|
||||
const body: ReadKnowledgeBaseResponse = {
|
||||
elser_exists: modelExists,
|
||||
index_exists: indexExists,
|
||||
is_setup_in_progress: kbDataClient.isSetupInProgress,
|
||||
is_setup_available: setupAvailable,
|
||||
pipeline_exists: pipelineExists,
|
||||
};
|
||||
|
||||
if (indexExists && isInferenceEndpointExists) {
|
||||
const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded();
|
||||
const userDataExists = await kbDataClient.isUserDataExists();
|
||||
|
||||
return response.ok({
|
||||
body: {
|
||||
...body,
|
||||
security_labs_exists: securityLabsExists,
|
||||
user_data_exists: userDataExists,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
return response.ok({ body });
|
||||
return response.ok({
|
||||
body: {
|
||||
elser_exists: isInferenceEndpointExists,
|
||||
index_exists: indexExists,
|
||||
is_setup_in_progress: kbDataClient.isSetupInProgress,
|
||||
is_setup_available: setupAvailable,
|
||||
security_labs_exists: securityLabsExists,
|
||||
// If user data exists, we should have at least one document in the Security Labs index
|
||||
user_data_exists: userDataExists || !!loadedSecurityLabsDocsCount,
|
||||
pipeline_exists: pipelineExists,
|
||||
},
|
||||
});
|
||||
} catch (err) {
|
||||
logger.error(err);
|
||||
const error = transformError(err);
|
||||
|
|
|
|||
|
|
@ -55,8 +55,6 @@ export const postKnowledgeBaseRoute = (router: ElasticAssistantPluginRouter) =>
|
|||
const resp = buildResponse(response);
|
||||
const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
|
||||
const assistantContext = ctx.elasticAssistant;
|
||||
const core = ctx.core;
|
||||
const soClient = core.savedObjects.getClient();
|
||||
const ignoreSecurityLabs = request.query.ignoreSecurityLabs;
|
||||
|
||||
try {
|
||||
|
|
@ -69,7 +67,6 @@ export const postKnowledgeBaseRoute = (router: ElasticAssistantPluginRouter) =>
|
|||
}
|
||||
|
||||
await knowledgeBaseDataClient.setupKnowledgeBase({
|
||||
soClient,
|
||||
ignoreSecurityLabs,
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ export class RequestContextFactory implements IRequestContextFactory {
|
|||
request: KibanaRequest
|
||||
): Promise<ElasticAssistantApiRequestHandlerContext> {
|
||||
const { options } = this;
|
||||
const { core } = options;
|
||||
const { core, plugins } = options;
|
||||
|
||||
const [coreStart, startPlugins] = await core.getStartServices();
|
||||
const coreContext = await context.core;
|
||||
|
|
@ -77,6 +77,8 @@ export class RequestContextFactory implements IRequestContextFactory {
|
|||
return contextUser;
|
||||
};
|
||||
|
||||
const savedObjectsClient = coreStart.savedObjects.getScopedClient(request);
|
||||
|
||||
return {
|
||||
core: coreContext,
|
||||
|
||||
|
|
@ -99,7 +101,7 @@ export class RequestContextFactory implements IRequestContextFactory {
|
|||
},
|
||||
llmTasks: startPlugins.llmTasks,
|
||||
inference: startPlugins.inference,
|
||||
savedObjectsClient: coreStart.savedObjects.getScopedClient(request),
|
||||
savedObjectsClient,
|
||||
telemetry: core.analytics,
|
||||
|
||||
// Note: modelIdOverride is used here to enable setting up the KB using a different ELSER model, which
|
||||
|
|
@ -121,6 +123,11 @@ export class RequestContextFactory implements IRequestContextFactory {
|
|||
modelIdOverride: params?.modelIdOverride,
|
||||
manageGlobalKnowledgeBaseAIAssistant:
|
||||
securitySolutionAssistant.manageGlobalKnowledgeBaseAIAssistant as boolean,
|
||||
// uses internal user to interact with ML API
|
||||
trainedModelsProvider: plugins.ml.trainedModelsProvider(
|
||||
{} as KibanaRequest,
|
||||
coreStart.savedObjects.createInternalRepository()
|
||||
),
|
||||
});
|
||||
}),
|
||||
|
||||
|
|
|
|||
|
|
@ -54,12 +54,10 @@ export default ({ getService }: FtrProviderContext) => {
|
|||
it('should create a new document entry for the current user', async () => {
|
||||
const entry = await createEntry({ supertest, log, entry: documentEntry });
|
||||
|
||||
const expectedDocumentEntry = {
|
||||
expect(removeServerGeneratedProperties(entry)).toMatchObject({
|
||||
...documentEntry,
|
||||
users: [{ name: 'elastic' }],
|
||||
};
|
||||
|
||||
expect(removeServerGeneratedProperties(entry)).toEqual(expectedDocumentEntry);
|
||||
});
|
||||
});
|
||||
|
||||
it('should create a new index entry for the current user', async () => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue