[8.x] [Security GenAI] Fix and un-skip Knowledge Base Integration Tests (#198861)

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')