[8.x] [Security GenAI] Fix and un-skip Knowledge Base Integration Tes…

…ts (elastic#198861) (elastic#198890)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[Security GenAI] Fix and un-skip Knowledge Base Integration Tests
(elastic#198861)](elastic#198861)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Ievgen
Sorokopud","email":"[email protected]"},"sourceCommit":{"committedDate":"2024-11-05T07:40:59Z","message":"[Security
GenAI] Fix and un-skip Knowledge Base Integration Tests (elastic#198861)\n\n##
Summary\r\n\r\nThis is a followup to
https://github.com/elastic/kibana/pull/198178\r\nwhere we skipped KB
integration tests. We enable it with this PR.\r\n\r\nSince it takes a
lot of time to setup all Security Labs docs, the idea\r\nis to skip
installing those docs when it is not needed. For these tests\r\nwe need
to make sure that inference endpoint is setup correctly - labs\r\ndocs
are not required in this case.\r\n\r\ncc
@stephmilovic","sha":"69c1e5a7dda2708773dfeed2314b0ef74f4537ee","branchLabelMapping":{"^v9.0.0$":"main","^v8.17.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","v9.0.0","Team:
SecuritySolution","Team:Security Generative
AI","backport:version","v8.17.0","v8.16.1"],"title":"[Security GenAI]
Fix and un-skip Knowledge Base Integration
Tests","number":198861,"url":"https://github.com/elastic/kibana/pull/198861","mergeCommit":{"message":"[Security
GenAI] Fix and un-skip Knowledge Base Integration Tests (elastic#198861)\n\n##
Summary\r\n\r\nThis is a followup to
https://github.com/elastic/kibana/pull/198178\r\nwhere we skipped KB
integration tests. We enable it with this PR.\r\n\r\nSince it takes a
lot of time to setup all Security Labs docs, the idea\r\nis to skip
installing those docs when it is not needed. For these tests\r\nwe need
to make sure that inference endpoint is setup correctly - labs\r\ndocs
are not required in this case.\r\n\r\ncc
@stephmilovic","sha":"69c1e5a7dda2708773dfeed2314b0ef74f4537ee"}},"sourceBranch":"main","suggestedTargetBranches":["8.x","8.16"],"targetPullRequestStates":[{"branch":"main","label":"v9.0.0","branchLabelMappingKey":"^v9.0.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/198861","number":198861,"mergeCommit":{"message":"[Security
GenAI] Fix and un-skip Knowledge Base Integration Tests (elastic#198861)\n\n##
Summary\r\n\r\nThis is a followup to
https://github.com/elastic/kibana/pull/198178\r\nwhere we skipped KB
integration tests. We enable it with this PR.\r\n\r\nSince it takes a
lot of time to setup all Security Labs docs, the idea\r\nis to skip
installing those docs when it is not needed. For these tests\r\nwe need
to make sure that inference endpoint is setup correctly - labs\r\ndocs
are not required in this case.\r\n\r\ncc
@stephmilovic","sha":"69c1e5a7dda2708773dfeed2314b0ef74f4537ee"}},{"branch":"8.x","label":"v8.17.0","branchLabelMappingKey":"^v8.17.0$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.16","label":"v8.16.1","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Ievgen Sorokopud <[email protected]>

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<typeof CreateKnowledgeBaseRequestQuery>;
 

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.

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<void> => {
     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...`);

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 } });

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<typeof MachineLearningProvider>;
 
-  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);

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<CreateKnowledgeBaseResponse> => {
   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')