[Security Assistant] Knowledge base switch to use semantic_text (el…

…astic#197007) ## Summary - `text_expansion` is deprecated, use `semantic_text` instead - fix KB index entry form field options - explicitly create inference endpoint on KB setup if `assistantKnowledgeBaseByDefault` is true - when upgrade from v1 update KB ingest pipeline and remove unnecessary processor, but keep the pipeline for the backward compatibility - switch to use `doc` update for KB entries due to the limitations od `semantic_text` https://www.elastic.co/guide/en/elasticsearch/reference/current/semantic-text.html#update-script - split loading Security labs content into smaller chunks --------- Co-authored-by: kibanamachine <[email protected]> Co-authored-by: Pedro Jaramillo <[email protected]> (cherry picked from commit a8c54f2)
kibanamachine · Oct 30, 2024 · 4fe4ba0 · 4fe4ba0
1 parent 68a7d5e
commit 4fe4ba0
Show file tree

Hide file tree

Showing 33 changed files with 485 additions and 195 deletions.
diff --git a/packages/kbn-data-stream-adapter/src/field_maps/types.ts b/packages/kbn-data-stream-adapter/src/field_maps/types.ts
@@ -54,6 +54,8 @@ export type FieldMap<T extends string = string> = Record<
     scaling_factor?: number;
     dynamic?: boolean | 'strict';
     properties?: Record<string, { type: string }>;
+    inference_id?: string;
+    copy_to?: string;
   }
 >;
 

diff --git a/...ck/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/...ck/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts
@@ -81,4 +81,5 @@ export const ReadKnowledgeBaseResponse = z.object({
   is_setup_in_progress: z.boolean().optional(),
   pipeline_exists: z.boolean().optional(),
   security_labs_exists: z.boolean().optional(),
+  user_data_exists: z.boolean().optional(),
 });
diff --git a/...ckages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml b/...ckages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml
@@ -78,6 +78,8 @@ paths:
                     type: boolean
                   security_labs_exists:
                     type: boolean
+                  user_data_exists:
+                    type: boolean
         400:
           description: Generic Error
           content:

diff --git a/...elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts b/...elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts
@@ -24,6 +24,7 @@ export interface UseKnowledgeBaseEntriesParams {
   signal?: AbortSignal | undefined;
   toasts?: IToasts;
   enabled?: boolean; // For disabling if FF is off
+  isRefetching?: boolean; // For enabling polling
 }
 
 const defaultQuery: FindKnowledgeBaseEntriesRequestQuery = {
@@ -56,6 +57,7 @@ export const useKnowledgeBaseEntries = ({
   signal,
   toasts,
   enabled = false,
+  isRefetching = false,
 }: UseKnowledgeBaseEntriesParams) =>
   useQuery(
     KNOWLEDGE_BASE_ENTRY_QUERY_KEY,
@@ -73,6 +75,7 @@ export const useKnowledgeBaseEntries = ({
       enabled,
       keepPreviousData: true,
       initialData: { page: 1, perPage: 100, total: 0, data: [] },
+      refetchInterval: isRefetching ? 30000 : false,
       onError: (error: IHttpFetchError<ResponseErrorBody>) => {
         if (error.name !== 'AbortError') {
           toasts?.addError(error, {

diff --git a/...bn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx b/...bn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.test.tsx
@@ -34,6 +34,7 @@ const statusResponse = {
   elser_exists: true,
   index_exists: true,
   pipeline_exists: true,
+  security_labs_exists: true,
 };
 
 const http = {

diff --git a/...ges/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/...ges/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx
@@ -45,6 +45,8 @@ export const useKnowledgeBaseStatus = ({
     {
       retry: false,
       keepPreviousData: true,
+      // Polling interval for Knowledge Base setup in progress
+      refetchInterval: (data) => (data?.is_setup_in_progress ? 30000 : false),
       // Deprecated, hoist to `queryCache` w/in `QueryClient. See: https://stackoverflow.com/a/76961109
       onError: (error: IHttpFetchError<ResponseErrorBody>) => {
         if (error.name !== 'AbortError') {
@@ -86,12 +88,12 @@ export const useInvalidateKnowledgeBaseStatus = () => {
  *
  * @param kbStatus ReadKnowledgeBaseResponse
  */
-export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undefined): boolean => {
-  return (
-    (kbStatus?.elser_exists &&
-      kbStatus?.security_labs_exists &&
-      kbStatus?.index_exists &&
-      kbStatus?.pipeline_exists) ??
-    false
-  );
-};
+export const isKnowledgeBaseSetup = (kbStatus: ReadKnowledgeBaseResponse | undefined): boolean =>
+  (kbStatus?.elser_exists &&
+    kbStatus?.index_exists &&
+    kbStatus?.pipeline_exists &&
+    // Allows to use UI while importing Security Labs docs
+    (kbStatus?.security_labs_exists ||
+      kbStatus?.is_setup_in_progress ||
+      kbStatus?.user_data_exists)) ??
+  false;
diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx
@@ -53,9 +53,9 @@ export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
     const isSecurityLabsEnabled = kbStatus?.security_labs_exists ?? false;
     const isKnowledgeBaseSetup =
       (isElserEnabled &&
-        isSecurityLabsEnabled &&
         kbStatus?.index_exists &&
-        kbStatus?.pipeline_exists) ??
+        kbStatus?.pipeline_exists &&
+        (isSecurityLabsEnabled || kbStatus?.user_data_exists)) ??
       false;
     const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false;
     const isSetupAvailable = kbStatus?.is_setup_available ?? false;

diff --git a/...es/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/...es/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx
@@ -160,6 +160,7 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
     http,
     toasts,
     enabled: enableKnowledgeBaseByDefault,
+    isRefetching: kbStatus?.is_setup_in_progress,
   });
 
   // Flyout Save/Cancel Actions
@@ -190,13 +191,15 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
         indices.push(entry.index);
       }
     });
-    return dataViews.getExistingIndices(indices);
+
+    return indices.length ? dataViews.getExistingIndices(indices) : Promise.resolve([]);
   }, [entries.data]);
 
   const { getColumns } = useKnowledgeBaseTable();
   const columns = useMemo(
     () =>
       getColumns({
+        isKbSetupInProgress: kbStatus?.is_setup_in_progress ?? false,
         existingIndices,
         isDeleteEnabled: (entry: KnowledgeBaseEntryResponse) => {
           return (
@@ -219,7 +222,14 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
           openFlyout();
         },
       }),
-    [entries.data, existingIndices, getColumns, hasManageGlobalKnowledgeBase, openFlyout]
+    [
+      entries.data,
+      existingIndices,
+      getColumns,
+      hasManageGlobalKnowledgeBase,
+      kbStatus?.is_setup_in_progress,
+      openFlyout,
+    ]
   );
 
   // Refresh button

diff --git a/...c-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/...c-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx
@@ -117,7 +117,7 @@ export const IndexEntryEditor: React.FC<Props> = React.memo(
         dataViews.getFieldsForWildcard({
           pattern: entry?.index ?? '',
         }),
-      []
+      [entry?.index]
     );
 
     const fieldOptions = useMemo(

diff --git a/...-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts b/...-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts
@@ -372,3 +372,10 @@ export const MISSING_INDEX_TOOLTIP_CONTENT = i18n.translate(
       'The index assigned to this knowledge base entry is unavailable. Check the permissions on the configured index, or that the index has not been deleted. You can update the index to be used for this knowledge entry, or delete the entry entirely.',
   }
 );
+
+export const SECURITY_LABS_NOT_FULLY_LOADED = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.securityLabsNotFullyLoadedTooltipContent',
+  {
+    defaultMessage: 'Security Labs content is not fully loaded. Click to reload.',
+  }
+);
diff --git a/...stant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx b/...stant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx
@@ -11,6 +11,7 @@ import {
   EuiBasicTableColumn,
   EuiIcon,
   EuiText,
+  EuiLoadingSpinner,
   EuiToolTip,
 } from '@elastic/eui';
 import { css } from '@emotion/react';
@@ -29,11 +30,16 @@ import * as i18n from './translations';
 import { BadgesColumn } from '../../assistant/common/components/assistant_settings_management/badges';
 import { useInlineActions } from '../../assistant/common/components/assistant_settings_management/inline_actions';
 import { isSystemEntry } from './helpers';
+import { SetupKnowledgeBaseButton } from '../setup_knowledge_base_button';
 
 const AuthorColumn = ({ entry }: { entry: KnowledgeBaseEntryResponse }) => {
   const { userProfileService } = useAssistantContext();
 
   const userProfile = useAsync(async () => {
+    if (isSystemEntry(entry) || entry.createdBy === 'unknown') {
+      return;
+    }
+
     const profile = await userProfileService?.bulkGet<{ avatar: UserProfileAvatarData }>({
       uids: new Set([entry.createdBy]),
       dataPath: 'avatar',
@@ -45,7 +51,7 @@ const AuthorColumn = ({ entry }: { entry: KnowledgeBaseEntryResponse }) => {
     () => userProfile?.value?.username ?? 'Unknown',
     [userProfile?.value?.username]
   );
-  const userAvatar = userProfile.value?.avatar;
+  const userAvatar = userProfile?.value?.avatar;
   const badgeItem = isSystemEntry(entry) ? 'Elastic' : userName;
   const userImage = isSystemEntry(entry) ? (
     <EuiIcon
@@ -142,12 +148,14 @@ export const useKnowledgeBaseTable = () => {
       isEditEnabled,
       onDeleteActionClicked,
       onEditActionClicked,
+      isKbSetupInProgress,
     }: {
       existingIndices?: string[];
       isDeleteEnabled: (entry: KnowledgeBaseEntryResponse) => boolean;
       isEditEnabled: (entry: KnowledgeBaseEntryResponse) => boolean;
       onDeleteActionClicked: (entry: KnowledgeBaseEntryResponse) => void;
       onEditActionClicked: (entry: KnowledgeBaseEntryResponse) => void;
+      isKbSetupInProgress: boolean;
     }): Array<EuiBasicTableColumn<KnowledgeBaseEntryResponse>> => {
       return [
         {
@@ -180,11 +188,27 @@ export const useKnowledgeBaseTable = () => {
         {
           name: i18n.COLUMN_ENTRIES,
           render: (entry: KnowledgeBaseEntryResponse) => {
-            return isSystemEntry(entry)
-              ? entry.text
-              : entry.type === DocumentEntryType.value
-              ? '1'
-              : '-';
+            return isSystemEntry(entry) ? (
+              <>
+                {`${entry.text}`}
+                {isKbSetupInProgress ? (
+                  <EuiLoadingSpinner
+                    size="m"
+                    css={css`
+                      margin-left: 8px;
+                    `}
+                  />
+                ) : (
+                  <EuiToolTip content={i18n.SECURITY_LABS_NOT_FULLY_LOADED}>
+                    <SetupKnowledgeBaseButton display="refresh" />
+                  </EuiToolTip>
+                )}
+              </>
+            ) : entry.type === DocumentEntryType.value ? (
+              '1'
+            ) : (
+              '-'
+            );
           },
         },
         {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx
@@ -6,15 +6,16 @@
  */
 
 import React, { useCallback } from 'react';
-import { EuiButton, EuiButtonEmpty, EuiToolTip } from '@elastic/eui';
+import { EuiButton, EuiButtonIcon, EuiButtonEmpty, EuiToolTip } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 
+import { css } from '@emotion/react';
 import { useAssistantContext } from '../..';
 import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base';
 import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status';
 
 interface Props {
-  display?: 'mini';
+  display?: 'mini' | 'refresh';
 }
 
 /**
@@ -48,6 +49,23 @@ export const SetupKnowledgeBaseButton: React.FC<Props> = React.memo(({ display }
       })
     : undefined;
 
+  if (display === 'refresh') {
+    return (
+      <EuiButtonIcon
+        color="primary"
+        data-test-subj="setup-knowledge-base-button"
+        disabled={!kbStatus?.is_setup_available}
+        isLoading={isSetupInProgress}
+        iconType="refresh"
+        onClick={onInstallKnowledgeBase}
+        size="xs"
+        css={css`
+          margin-left: 8px;
+        `}
+      />
+    );
+  }
+
   return (
     <EuiToolTip position={'bottom'} content={toolTipContent}>
       {display === 'mini' ? (

diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/msearch_query.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/msearch_query.ts
@@ -34,12 +34,10 @@ export const mSearchQueryBody: MsearchQueryBody = {
           ],
           must: [
             {
-              text_expansion: {
-                'vector.tokens': {
-                  model_id: '.elser_model_2',
-                  model_text:
-                    'Generate an ESQL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.',
-                },
+              semantic: {
+                field: 'semantic_text',
+                query:
+                  'Generate an ESQL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.',
               },
             },
           ],

diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/vector_search_query.ts
@@ -26,12 +26,10 @@ export const mockVectorSearchQuery: QueryDslQueryContainer = {
     ],
     must: [
       {
-        text_expansion: {
-          'vector.tokens': {
-            model_id: '.elser_model_2',
-            model_text:
-              'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.',
-          },
+        semantic: {
+          field: 'semantic_text',
+          query:
+            'Generate an ES|QL query that will count the number of connections made to external IP addresses, broken down by user. If the count is greater than 100 for a specific user, add a new field called "follow_up" that contains a value of "true", otherwise, it should contain "false". The user names should also be enriched with their respective group names.',
         },
       },
     ],

diff --git a/...lugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts b/...lugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts
@@ -99,7 +99,8 @@ export const getUpdateScript = ({
   isPatch?: boolean;
 }) => {
   return {
-    source: `
+    script: {
+      source: `
     if (params.assignEmpty == true || params.containsKey('allowed')) {
       ctx._source.allowed = params.allowed;
     }
@@ -108,11 +109,12 @@ export const getUpdateScript = ({
     }
     ctx._source.updated_at = params.updated_at;
   `,
-    lang: 'painless',
-    params: {
-      ...anonymizationField, // when assigning undefined in painless, it will remove property and wil set it to null
-      // for patch we don't want to remove unspecified value in payload
-      assignEmpty: !(isPatch ?? true),
+      lang: 'painless',
+      params: {
+        ...anonymizationField, // when assigning undefined in painless, it will remove property and wil set it to null
+        // for patch we don't want to remove unspecified value in payload
+        assignEmpty: !(isPatch ?? true),
+      },
     },
   };
 };
diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/helpers.ts
@@ -15,7 +15,8 @@ export const getUpdateScript = ({
   isPatch?: boolean;
 }) => {
   return {
-    source: `
+    script: {
+      source: `
     if (params.assignEmpty == true || params.containsKey('api_config')) {
       if (ctx._source.api_config != null) {
         if (params.assignEmpty == true || params.api_config.containsKey('connector_id')) {
@@ -70,11 +71,12 @@ export const getUpdateScript = ({
     }
     ctx._source.updated_at = params.updated_at;
   `,
-    lang: 'painless',
-    params: {
-      ...conversation, // when assigning undefined in painless, it will remove property and wil set it to null
-      // for patch we don't want to remove unspecified value in payload
-      assignEmpty: !(isPatch ?? true),
+      lang: 'painless',
+      params: {
+        ...conversation, // when assigning undefined in painless, it will remove property and wil set it to null
+        // for patch we don't want to remove unspecified value in payload
+        assignEmpty: !(isPatch ?? true),
+      },
     },
   };
 };
diff --git a/...s/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts b/...s/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.ts
@@ -76,7 +76,7 @@ export const updateConversation = async ({
         },
       },
       refresh: true,
-      script: getUpdateScript({ conversation: params, isPatch }),
+      script: getUpdateScript({ conversation: params, isPatch }).script,
     });
 
     if (response.failures && response.failures.length > 0) {