From 69c1e5a7dda2708773dfeed2314b0ef74f4537ee Mon Sep 17 00:00:00 2001 From: Ievgen Sorokopud Date: Tue, 5 Nov 2024 08:40:59 +0100 Subject: [PATCH] [Security GenAI] Fix and un-skip Knowledge Base Integration Tests (#198861) ## Summary This is a followup to https://github.com/elastic/kibana/pull/198178 where we skipped KB integration tests. We enable it with this PR. Since it takes a lot of time to setup all Security Labs docs, the idea is to skip installing those docs when it is not needed. For these tests we need to make sure that inference endpoint is setup correctly - labs docs are not required in this case. cc @stephmilovic --- .../impl/schemas/knowledge_base/crud_kb_route.gen.ts | 5 +++++ .../impl/schemas/knowledge_base/crud_kb_route.schema.yaml | 7 +++++++ .../ai_assistant_data_clients/knowledge_base/index.ts | 4 +++- .../server/routes/knowledge_base/post_knowledge_base.ts | 2 ++ .../entries/trial_license_complete_tier/entries.ts | 2 +- .../genai/knowledge_base/entries/utils/helpers.ts | 5 ++++- 6 files changed, 22 insertions(+), 3 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index aad215021da81..4f03dbe0b1343 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -15,6 +15,7 @@ */ import { z } from '@kbn/zod'; +import { BooleanFromString } from '@kbn/zod-helpers'; /** * AI assistant KnowledgeBase. @@ -33,6 +34,10 @@ export const CreateKnowledgeBaseRequestQuery = z.object({ * Optional ELSER modelId to use when setting up the Knowledge Base */ modelId: z.string().optional(), + /** + * Indicates whether we should or should not install Security Labs docs when setting up the Knowledge Base + */ + ignoreSecurityLabs: BooleanFromString.optional().default(false), }); export type CreateKnowledgeBaseRequestQueryInput = z.input; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml index 0e0f1e9267916..b4c16189e2387 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml @@ -24,6 +24,13 @@ paths: required: false schema: type: string + - name: ignoreSecurityLabs + in: query + description: Indicates whether we should or should not install Security Labs docs when setting up the Knowledge Base + required: false + schema: + type: boolean + default: false responses: 200: description: Indicates a successful call. diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index f985095661f3e..333fbb796ddd9 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -269,9 +269,11 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { public setupKnowledgeBase = async ({ soClient, v2KnowledgeBaseEnabled = true, + ignoreSecurityLabs = false, }: { soClient: SavedObjectsClientContract; v2KnowledgeBaseEnabled?: boolean; + ignoreSecurityLabs?: boolean; }): Promise => { if (this.options.getIsKBSetupInProgress()) { this.options.logger.debug('Knowledge Base setup already in progress'); @@ -366,7 +368,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { this.options.logger.debug(`Checking if Knowledge Base docs have been loaded...`); - if (v2KnowledgeBaseEnabled) { + if (v2KnowledgeBaseEnabled && !ignoreSecurityLabs) { const labsDocsLoaded = await this.isSecurityLabsDocsLoaded(); if (!labsDocsLoaded) { this.options.logger.debug(`Loading Security Labs KB docs...`); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 96317da303ac1..23604886e4a52 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -60,6 +60,7 @@ export const postKnowledgeBaseRoute = (router: ElasticAssistantPluginRouter) => // Only allow modelId override if FF is enabled as this will re-write the ingest pipeline and break any previous KB entries // This is only really needed for API integration tests const modelIdOverride = v2KnowledgeBaseEnabled ? request.query.modelId : undefined; + const ignoreSecurityLabs = request.query.ignoreSecurityLabs; try { const knowledgeBaseDataClient = @@ -74,6 +75,7 @@ export const postKnowledgeBaseRoute = (router: ElasticAssistantPluginRouter) => await knowledgeBaseDataClient.setupKnowledgeBase({ soClient, v2KnowledgeBaseEnabled, + ignoreSecurityLabs, }); return response.ok({ body: { success: true } }); diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts index a2e1c2c08be2c..2ecb368c2ba7b 100644 --- a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts @@ -32,7 +32,7 @@ export default ({ getService }: FtrProviderContext) => { const es = getService('es'); const ml = getService('ml') as ReturnType; - describe.skip('@ess Basic Security AI Assistant Knowledge Base Entries', () => { + describe('@ess Basic Security AI Assistant Knowledge Base Entries', () => { before(async () => { await installTinyElser(ml); await setupKnowledgeBase(supertest, log); diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/helpers.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/helpers.ts index 62174da6bce4c..cdbb1cca4de24 100644 --- a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/helpers.ts +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/helpers.ts @@ -64,7 +64,10 @@ export const setupKnowledgeBase = async ( namespace?: string ): Promise => { const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); - const route = routeWithNamespace(`${path}?modelId=pt_tiny_elser`, namespace); + const route = routeWithNamespace( + `${path}?modelId=pt_tiny_elser&ignoreSecurityLabs=true`, + namespace + ); const response = await supertest .post(route) .set('kbn-xsrf', 'true')