[Security GenAI] Fix and un-skip Knowledge Base Integration Tests (el…

…astic#198861)

## Summary

This is a followup to elastic#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

(cherry picked from commit 69c1e5a)

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