Merge pull request elastic#15 from yansavitski/ai-playground-indices-…

…sources Ai playground indices sources
Samiul-TheSoccerFan · Feb 20, 2024 · beebf84 · beebf84
2 parents 270a41e + 4e34ed1
commit beebf84
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 52 deletions.
diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
@@ -35,6 +35,7 @@ import { SourcesPanelSidebar } from './sources_panel/sources_panel_sidebar';
 import { TelegramIcon } from './telegram_icon';
 
 import { transformFromChatMessages } from '../utils/transformToMessages';
+import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 
 export const Chat = () => {
   const { euiTheme } = useEuiTheme();
@@ -52,7 +53,7 @@ export const Chat = () => {
       {
         data: {
           prompt: data[ChatFormFields.prompt],
-          indices: 'workplace_index',
+          indices: data[ChatFormFields.indices].join(),
           api_key: data[ChatFormFields.openAIKey],
           citations: data[ChatFormFields.citations].toString(),
         },
@@ -183,7 +184,22 @@ export const Chat = () => {
             )}
           />
 
-          <SourcesPanelSidebar />
+          <Controller
+            name={ChatFormFields.indices}
+            control={control}
+            defaultValue={[]}
+            render={({ field }) => (
+              <SourcesPanelSidebar
+                selectedIndices={field.value}
+                addIndex={(newIndex: IndexName) => {
+                  field.onChange([...field.value, newIndex]);
+                }}
+                removeIndex={(index: IndexName) => {
+                  field.onChange(field.value.filter((indexName) => indexName !== index));
+                }}
+              />
+            )}
+          />
         </EuiFlexItem>
       </EuiFlexGroup>
     </EuiForm>

diff --git a/packages/kbn-ai-playground/components/sources_panel/add_indices_field.tsx b/packages/kbn-ai-playground/components/sources_panel/add_indices_field.tsx
@@ -6,24 +6,32 @@
  * Side Public License, v 1.
  */
 
-import { EuiFormRow, EuiSuperSelect } from '@elastic/eui';
+import { EuiComboBox, EuiFormRow } from '@elastic/eui';
 import React, { useState } from 'react';
 import { i18n } from '@kbn/i18n';
 import { useQueryIndices } from '../../hooks/useQueryIndices';
+import { EuiComboBoxOptionOption } from '@elastic/eui/src/components/combo_box/types';
+import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 
 interface AddIndicesFieldProps {
-  indices: string[];
-  selectedIndices: string[];
-  onIndexSelect: (index: string) => void;
+  selectedIndices: IndexName[];
+  onIndexSelect: (index: IndexName) => void;
 }
 
 export const AddIndicesField: React.FC<AddIndicesFieldProps> = ({
   selectedIndices,
-  indices,
   onIndexSelect,
 }) => {
   const [query, setQuery] = useState<string>('');
-  const { options, isLoading } = useQueryIndices(query);
+  const { indices, isLoading } = useQueryIndices(query);
+  const handleChange = (value: Array<EuiComboBoxOptionOption<IndexName>>) => {
+    if (value?.[0]?.label) {
+      onIndexSelect(value[0].label);
+    }
+  };
+  const handleSearchChange = (searchValue: string) => {
+    setQuery(searchValue);
+  };
 
   return (
     <EuiFormRow
@@ -33,18 +41,21 @@ export const AddIndicesField: React.FC<AddIndicesFieldProps> = ({
       })}
       labelType="legend"
     >
-      <EuiSuperSelect
+      <EuiComboBox
+        singleSelection={{ asPlainText: true }}
         placeholder={i18n.translate('aiPlayground.sources.addIndex.placeholder', {
           defaultMessage: 'Select new data source',
         })}
+        async
+        isLoading={isLoading}
+        onChange={handleChange}
+        onSearchChange={handleSearchChange}
         fullWidth
         options={indices.map((index) => ({
-          value: index,
-          inputDisplay: index,
+          label: index,
           disabled: selectedIndices.includes(index),
         }))}
-        onChange={onIndexSelect}
-        hasDividers
+        isClearable={false}
       />
     </EuiFormRow>
   );

diff --git a/packages/kbn-ai-playground/components/sources_panel/sources_panel_empty_prompt.tsx b/packages/kbn-ai-playground/components/sources_panel/sources_panel_empty_prompt.tsx
@@ -5,27 +5,19 @@
  * 2.0.
  */
 
-import {
-  EuiPanel,
-  EuiTitle,
-  EuiText,
-  EuiLink,
-  EuiFlexGroup,
-  EuiFlexItem,
-  EuiSpacer,
-} from '@elastic/eui';
+import { EuiPanel, EuiTitle, EuiText, EuiFlexGroup, EuiFlexItem, EuiSpacer } from '@elastic/eui';
 import React from 'react';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { AddIndicesField } from './add_indices_field';
 import { IndicesTable } from './indices_table';
+import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 
 export const SourcesPanelEmptyPrompt: React.FC = () => {
-  const indices = ['search-index', 'search-books'];
-  const [selectedIndices, setSelectedIndices] = React.useState<string[]>([]);
-  const addIndex = (newIndex: string) => {
+  const [selectedIndices, setSelectedIndices] = React.useState<IndexName[]>([]);
+  const addIndex = (newIndex: IndexName) => {
     setSelectedIndices([...selectedIndices, newIndex]);
   };
-  const removeIndex = (index: string) => {
+  const removeIndex = (index: IndexName) => {
     setSelectedIndices(selectedIndices.filter((indexName) => indexName !== index));
   };
 
@@ -59,11 +51,7 @@ export const SourcesPanelEmptyPrompt: React.FC = () => {
         )}
 
         <EuiFlexItem>
-          <AddIndicesField
-            selectedIndices={selectedIndices}
-            indices={indices}
-            onIndexSelect={addIndex}
-          />
+          <AddIndicesField selectedIndices={selectedIndices} onIndexSelect={addIndex} />
         </EuiFlexItem>
       </EuiFlexGroup>
     </EuiPanel>

diff --git a/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx b/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
@@ -17,19 +17,20 @@ import {
 } from '@elastic/eui';
 import { IndicesList } from './indices_list';
 import { AddIndicesField } from './add_indices_field';
+import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 
-interface SourcesPanelSidebarProps {}
+interface SourcesPanelSidebarProps {
+  selectedIndices: IndexName[];
+  addIndex: (newIndex: IndexName) => void;
+  removeIndex: (index: IndexName) => void;
+}
 
-export const SourcesPanelSidebar: React.FC<SourcesPanelSidebarProps> = () => {
+export const SourcesPanelSidebar: React.FC<SourcesPanelSidebarProps> = ({
+  selectedIndices,
+  addIndex,
+  removeIndex,
+}) => {
   const accordionId = useGeneratedHtmlId({ prefix: 'sourceAccordion' });
-  const indices = ['search-index', 'search-books'];
-  const [selectedIndices, setSelectedIndices] = React.useState<string[]>([]);
-  const addIndex = (newIndex: string) => {
-    setSelectedIndices([...selectedIndices, newIndex]);
-  };
-  const removeIndex = (index: string) => {
-    setSelectedIndices(selectedIndices.filter((indexName) => indexName !== index));
-  };
 
   return (
     <EuiAccordion
@@ -56,11 +57,7 @@ export const SourcesPanelSidebar: React.FC<SourcesPanelSidebarProps> = () => {
         </EuiFlexItem>
 
         <EuiFlexItem>
-          <AddIndicesField
-            selectedIndices={selectedIndices}
-            indices={indices}
-            onIndexSelect={addIndex}
-          />
+          <AddIndicesField selectedIndices={selectedIndices} onIndexSelect={addIndex} />
         </EuiFlexItem>
       </EuiFlexGroup>
     </EuiAccordion>

diff --git a/packages/kbn-ai-playground/hooks/useQueryIndices.tsx b/packages/kbn-ai-playground/hooks/useQueryIndices.tsx
@@ -9,8 +9,11 @@
 import { useQuery } from '@tanstack/react-query';
 import { useKibana } from '@kbn/kibana-react-plugin/public';
 import { AIPlaygroundPluginStartDeps, ElasticsearchIndex } from '../types';
+import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 
-export const useQueryIndices = (query: string = '') => {
+export const useQueryIndices = (
+  query: string = ''
+): { indices: IndexName[]; isLoading: boolean } => {
   const { services } = useKibana<AIPlaygroundPluginStartDeps>();
 
   const { data, isLoading } = useQuery({
@@ -28,12 +31,9 @@ export const useQueryIndices = (query: string = '') => {
         },
       });
 
-      return response.indices.map((index) => ({
-        label: index.name,
-        key: index.name,
-      }));
+      return response.indices.map((index) => index.name);
     },
   });
 
-  return { options: data || [], isLoading };
+  return { indices: data || [], isLoading };
 };
diff --git a/packages/kbn-ai-playground/types.ts b/packages/kbn-ai-playground/types.ts
@@ -44,13 +44,15 @@ export enum ChatFormFields {
   citations = 'citations',
   prompt = 'prompt',
   openAIKey = 'api_key',
+  indices = 'indices',
 }
 
 export interface ChatForm {
   [ChatFormFields.question]: string;
   [ChatFormFields.prompt]: string;
   [ChatFormFields.citations]: boolean;
   [ChatFormFields.openAIKey]: string;
+  [ChatFormFields.indices]: IndexName[];
 }
 
 export interface AIPlaygroundPluginStartDeps {