[Security Assistant] Add field type badge to Index Entry field suggestions

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx

@@ -432,6 +432,7 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
               entry={selectedEntry as IndexEntry}
               originalEntry={originalEntry as IndexEntry}
               dataViews={dataViews}
+              http={http}
               setEntry={
                 setSelectedEntry as React.Dispatch<React.SetStateAction<Partial<IndexEntry>>>
               }

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.test.tsx

@@ -6,13 +6,16 @@
  */
 
 import React from 'react';
-import userEvent from '@testing-library/user-event';
 import { render, fireEvent, waitFor, within } from '@testing-library/react';
 import { IndexEntryEditor } from './index_entry_editor';
 import type { DataViewsContract } from '@kbn/data-views-plugin/public';
 import type { IndexEntry } from '@kbn/elastic-assistant-common';
 import * as i18n from './translations';
 import { I18nProvider } from '@kbn/i18n-react';
+import { useIndexMappings } from './use_index_mappings';
+import { HttpSetup } from '@kbn/core-http-browser';
+
+jest.mock('./use_index_mappings');
 
 const Wrapper = ({ children }: { children?: React.ReactNode }) => (
   <I18nProvider>{children}</I18nProvider>
@@ -21,17 +24,15 @@ const Wrapper = ({ children }: { children?: React.ReactNode }) => (
 describe('IndexEntryEditor', () => {
   const mockSetEntry = jest.fn();
   const mockDataViews = {
-    getFieldsForWildcard: jest.fn().mockResolvedValue([
-      { name: 'field-1', esTypes: ['semantic_text'] },
-      { name: 'field-2', esTypes: ['text'] },
-      { name: 'field-3', esTypes: ['semantic_text'] },
-    ]),
     getExistingIndices: jest.fn().mockResolvedValue(['index-1']),
     getIndices: jest.fn().mockResolvedValue([
       { name: 'index-1', attributes: ['open'] },
       { name: 'index-2', attributes: ['open'] },
     ]),
   } as unknown as DataViewsContract;
+  const http = {
+    get: jest.fn(),
+  } as unknown as HttpSetup;
 
   const defaultProps = {
     dataViews: mockDataViews,
@@ -46,10 +47,38 @@ describe('IndexEntryEditor', () => {
       queryDescription: 'Test Query Description',
       users: [],
     } as unknown as IndexEntry,
+    http,
   };
 
   beforeEach(() => {
     jest.clearAllMocks();
+    (useIndexMappings as jest.Mock).mockImplementation(({ indexName }) => {
+      if (indexName === 'index-1') {
+        return {
+          data: {
+            mappings: {
+              properties: {
+                'field-1-text': { type: 'text' },
+                'field-1-keyword': { type: 'keyword' },
+              },
+            },
+          },
+        };
+      }
+      if (indexName === 'index-2') {
+        return {
+          data: {
+            mappings: {
+              properties: {
+                'field-2-text': { type: 'text' },
+                'field-2-semantic': { type: 'semantic_text' },
+              },
+            },
+          },
+        };
+      }
+      return { data: undefined };
+    });
   });
 
   it('renders the form fields with initial values', async () => {
@@ -131,31 +160,50 @@ describe('IndexEntryEditor', () => {
   });
 
   it('fetches field options based on selected index and updates on selection', async () => {
-    const { getByTestId, getAllByTestId } = render(<IndexEntryEditor {...defaultProps} />, {
-      wrapper: Wrapper,
-    });
+    const { getByTestId, queryByTestId, rerender } = render(
+      <IndexEntryEditor {...defaultProps} />,
+      {
+        wrapper: Wrapper,
+      }
+    );
 
     await waitFor(() => {
-      expect(mockDataViews.getFieldsForWildcard).toHaveBeenCalledWith({
-        pattern: 'index-1',
+      expect(useIndexMappings).toHaveBeenCalledWith({
+        http,
+        indexName: 'index-1',
       });
     });
 
-    await waitFor(async () => {
-      fireEvent.click(getByTestId('index-combobox'));
-      fireEvent.click(getAllByTestId('comboBoxToggleListButton')[0]);
+    // Open field dropdown and check options for index-1
+    fireEvent.click(within(getByTestId('entry-combobox')).getByTestId('comboBoxToggleListButton'));
+    await waitFor(() => {
+      expect(getByTestId('field-option-field-1-text')).toBeInTheDocument();
+      expect(queryByTestId('field-option-field-1-keyword')).not.toBeInTheDocument();
     });
-    fireEvent.click(getByTestId('index-2'));
+
+    // Close the dropdown before re-rendering
+    fireEvent.click(within(getByTestId('entry-combobox')).getByTestId('comboBoxToggleListButton'));
+
+    // Change index to index-2
+    const newEntry = { ...defaultProps.entry, index: 'index-2', field: '' };
+    rerender(<IndexEntryEditor {...defaultProps} entry={newEntry} />);
 
     await waitFor(() => {
-      fireEvent.click(getByTestId('entry-combobox'));
+      expect(useIndexMappings).toHaveBeenCalledWith({
+        http,
+        indexName: 'index-2',
+      });
     });
 
-    await userEvent.type(
-      within(getByTestId('entry-combobox')).getByTestId('comboBoxSearchInput'),
-      'field-3'
-    );
+    // Open field dropdown and check options for index-2
+    fireEvent.click(within(getByTestId('entry-combobox')).getByTestId('comboBoxToggleListButton'));
+    await waitFor(() => {
+      expect(getByTestId('field-option-field-2-text')).toBeInTheDocument();
+      expect(getByTestId('field-option-field-2-semantic')).toBeInTheDocument();
+    });
 
+    // Select a new field
+    fireEvent.click(getByTestId('field-option-field-2-text'));
     await waitFor(() => {
       expect(mockSetEntry).toHaveBeenCalledWith(expect.any(Function));
     });

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx

@@ -7,6 +7,7 @@
 
 import type { EuiComboBoxOptionOption } from '@elastic/eui';
 import {
+  EuiBadge,
   EuiComboBox,
   EuiFieldText,
   EuiForm,
@@ -22,10 +23,13 @@ import React, { useCallback, useMemo } from 'react';
 import type { IndexEntry } from '@kbn/elastic-assistant-common';
 import type { DataViewsContract } from '@kbn/data-views-plugin/public';
 import { FormattedMessage } from '@kbn/i18n-react';
+import { HttpSetup } from '@kbn/core-http-browser';
+import { useIndexMappings } from './use_index_mappings';
 import * as i18n from './translations';
 import { isGlobalEntry } from './helpers';
 
 interface Props {
+  http: HttpSetup;
   dataViews: DataViewsContract;
   entry?: IndexEntry;
   originalEntry?: IndexEntry;
@@ -35,7 +39,7 @@ interface Props {
 }
 
 export const IndexEntryEditor: React.FC<Props> = React.memo<Props>(
-  ({ dataViews, entry, setEntry, hasManageGlobalKnowledgeBase, originalEntry, docLink }) => {
+  ({ http, dataViews, entry, setEntry, hasManageGlobalKnowledgeBase, originalEntry, docLink }) => {
     const privateUsers = useMemo(() => {
       const originalUsers = originalEntry?.users;
       if (originalEntry && !isGlobalEntry(originalEntry)) {
@@ -122,36 +126,33 @@ export const IndexEntryEditor: React.FC<Props> = React.memo<Props>(
       return !(await dataViews.getExistingIndices([entry.index])).length;
     }, [entry?.index]);
 
-    const indexFields = useAsync(
-      async () =>
-        dataViews.getFieldsForWildcard({
-          pattern: entry?.index ?? '',
-        }),
-      [entry?.index]
-    );
+    const { data: mappingData } = useIndexMappings({
+      http,
+      indexName: entry?.index ?? '',
+    });
 
     const fieldOptions = useMemo(
       () =>
-        Array.isArray(indexFields?.value)
-          ? indexFields.value
-              .filter((field) => field.esTypes?.includes('text'))
-              .map((field) => ({
-                'data-test-subj': field.name,
-                label: field.name,
-                value: field.name,
-              }))
-          : [],
-      [indexFields?.value]
+        Object.entries(mappingData?.mappings?.properties ?? {})
+          .filter(([, m]) => m.type === 'text' || m.type === 'semantic_text')
+          .map(([name, details]) => ({
+            'data-test-subj': `field-option-${name}`,
+            label: name,
+            value: name,
+            append: <EuiBadge color={'hollow'}>{details.type}</EuiBadge>,
+          })),
+      [mappingData]
     );
 
     const outputFieldOptions = useMemo(
       () =>
-        indexFields?.value?.map((field) => ({
-          'data-test-subj': field.name,
-          label: field.name,
-          value: field.name,
-        })) ?? [],
-      [indexFields?.value]
+        Object.entries(mappingData?.mappings?.properties ?? {}).map(([name, details]) => ({
+          'data-test-subj': `output-field-option-${name}`,
+          label: name,
+          value: name,
+          append: <EuiBadge color={'hollow'}>{details.type}</EuiBadge>,
+        })),
+      [mappingData]
     );
 
     const onCreateIndexOption = (searchValue: string) => {

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_index_mappings.ts

@@ -0,0 +1,68 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { UseQueryResult } from '@tanstack/react-query';
+import type { HttpSetup, IHttpFetchError, ResponseErrorBody } from '@kbn/core-http-browser';
+import type { IToasts } from '@kbn/core-notifications-browser';
+import { i18n } from '@kbn/i18n';
+import type { MappingPropertyBase } from '@elastic/elasticsearch/lib/api/types';
+import { useQuery } from '@tanstack/react-query';
+
+const KNOWLEDGE_BASE_MAPPINGS_QUERY_KEY = ['elastic-assistant', 'knowledge-base-mappings'];
+
+export interface Mappings {
+  mappings: {
+    properties: MappingPropertyBase['properties'];
+  };
+}
+
+export interface UseIndexMappingParams {
+  http: HttpSetup;
+  indexName: string;
+  toasts?: IToasts;
+}
+
+/**
+ * Hook for getting index mappings
+ *
+ * @param {Object} options - The options object.
+ * @param {HttpSetup} [options.http] - HttpSetup
+ * @param {String} [options.indexName] - String
+ * @param {IToasts} [options.toasts] - IToasts
+ *
+ * @returns {useQuery} hook for getting mappings for a given index
+ */
+export const useIndexMappings = ({
+  indexName,
+  http,
+  toasts,
+}: UseIndexMappingParams): UseQueryResult<Mappings, IHttpFetchError> => {
+  return useQuery(
+    [KNOWLEDGE_BASE_MAPPINGS_QUERY_KEY, indexName],
+    async ({ signal }) => {
+      return http.fetch<Mappings>(
+        `/api/index_management/mapping/${encodeURIComponent(indexName)}`,
+        { signal }
+      );
+    },
+    {
+      enabled: !!indexName,
+      onError: (error: IHttpFetchError<ResponseErrorBody>) => {
+        if (error.name !== 'AbortError') {
+          toasts?.addError(
+            error.body && error.body.message ? new Error(error.body.message) : error,
+            {
+              title: i18n.translate('xpack.elasticAssistant.knowledgeBase.mappingsError', {
+                defaultMessage: 'Error fetching Knowledge Base mappings',
+              }),
+            }
+          );
+        }
+      },
+    }
+  );
+};