add source selector for vega visualization

Signed-off-by: Yulong Ruan <ruanyl@amazon.com>

src/plugins/vis_type_vega/public/components/vega_vis_editor.tsx

@@ -28,7 +28,7 @@
  * under the License.
  */
 
-import React, { useCallback, useEffect } from 'react';
+import React, { useCallback, useEffect, useRef } from 'react';
 import { EuiCodeEditor } from '@elastic/eui';
 import compactStringify from 'json-stringify-pretty-compact';
 import hjson from 'hjson';
@@ -76,6 +76,9 @@ function format(
 }
 
 function VegaVisEditor({ stateParams, setValue }: VisOptionsProps<VisParams>) {
+  const setValueRef = useRef(setValue);
+  setValueRef.current = setValue;
+
   const onChange = useCallback(
     (value: string) => {
       setValue('spec', value);
@@ -104,15 +107,15 @@ function VegaVisEditor({ stateParams, setValue }: VisOptionsProps<VisParams>) {
             }),
           });
         } else {
-          setValue('spec', JSON.stringify(result, null, 4));
+          setValueRef.current('spec', JSON.stringify(result, null, 4));
         }
       }
     });
 
     return () => {
       subscription.unsubscribe();
     };
-  }, [setValue]);
+  }, []);
 
   return (
     <div className="vgaEditor">

src/plugins/vis_type_vega/public/text_to_vega.ts

@@ -29,7 +29,7 @@ Just reply with the json based Vega-Lite object, do not include any other conten
 };
 
 export class Text2Vega {
-  input$ = new BehaviorSubject({ input: '' });
+  input$ = new BehaviorSubject({ prompt: '', index: '' });
   result$: Observable<Record<string, any> | { error: any }>;
   status$ = new BehaviorSubject<'RUNNING' | 'STOPPED'>('STOPPED');
   http: HttpSetup;
@@ -38,19 +38,19 @@ export class Text2Vega {
     this.http = http;
     this.result$ = this.input$
       .pipe(
-        filter((v) => v.input.length > 0),
+        filter((v) => v.prompt.length > 0),
         debounceTime(200),
         tap(() => this.status$.next('RUNNING'))
       )
       .pipe(
         switchMap((v) =>
-          of(v.input).pipe(
+          of(v).pipe(
             // text to ppl
             switchMap(async (value) => {
-              const pplQuestion = value.split('//')[0];
-              const ppl = await this.text2ppl(pplQuestion);
+              const pplQuestion = value.prompt.split('//')[0];
+              const ppl = await this.text2ppl(pplQuestion, value.index);
               return {
-                input: value,
+                ...value,
                 ppl,
               };
             }),
@@ -64,7 +64,7 @@ export class Text2Vega {
             }),
             // call llm to generate vega
             switchMap(async (value) => {
-              const prompt = createPrompt(value.input, value.ppl, value.sample);
+              const prompt = createPrompt(value.prompt, value.ppl, value.sample);
               const result = await this.text2vega(prompt);
               result.data = {
                 url: {
@@ -89,18 +89,18 @@ export class Text2Vega {
     return result;
   }
 
-  async text2ppl(query: string) {
+  async text2ppl(query: string, index: string) {
     const pplResponse = await this.http.post('/api/llm/text2ppl', {
       body: JSON.stringify({
         question: query,
-        index: 'opensearch_dashboards_sample_data_logs',
+        index: index,
       }),
     });
     const result = JSON.parse(pplResponse.body.inference_results[0].output[0].result);
     return result.ppl;
   }
 
-  invoke(value: { input: string }) {
+  invoke(value: { prompt: string; index: string }) {
     this.input$.next(value);
   }
 

src/plugins/visualize/public/application/components/source_selector.tsx

@@ -0,0 +1,99 @@
+import React, { useCallback, useMemo, useState, useEffect } from 'react';
+import { i18n } from '@osd/i18n';
+
+import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
+import {
+  DataSource,
+  DataSourceGroup,
+  DataSourceSelectable,
+  DataSourceOption,
+} from '../../../../data/public';
+import { VisualizeServices } from '../types';
+
+export const SourceSelector = ({
+  selectedSourceId,
+  onChange,
+}: {
+  selectedSourceId: string;
+  onChange: (ds: DataSourceOption) => void;
+}) => {
+  const {
+    services: {
+      data: { dataSources },
+      notifications: { toasts },
+    },
+  } = useOpenSearchDashboards<VisualizeServices>();
+  const [currentDataSources, setCurrentDataSources] = useState<DataSource[]>([]);
+  const [dataSourceOptions, setDataSourceOptions] = useState<DataSourceGroup[]>([]);
+
+  const selectedSources = useMemo(() => {
+    if (selectedSourceId) {
+      for (const group of dataSourceOptions) {
+        for (const item of group.options) {
+          if (item.value === selectedSourceId) {
+            return [item];
+          }
+        }
+      }
+    }
+    return [];
+  }, [selectedSourceId, dataSourceOptions]);
+
+  useEffect(() => {
+    if (
+      !selectedSourceId &&
+      dataSourceOptions.length > 0 &&
+      dataSourceOptions[0].options.length > 0
+    ) {
+      onChange(dataSourceOptions[0].options[0]);
+    }
+  }, [selectedSourceId, dataSourceOptions]);
+
+  useEffect(() => {
+    const subscription = dataSources.dataSourceService
+      .getDataSources$()
+      .subscribe((currentDataSources) => {
+        setCurrentDataSources(Object.values(currentDataSources));
+      });
+
+    return () => {
+      subscription.unsubscribe();
+    };
+  }, [dataSources]);
+
+  const onDataSourceSelect = useCallback(
+    (selectedDataSources: DataSourceOption[]) => {
+      onChange(selectedDataSources[0]);
+    },
+    [onChange]
+  );
+
+  const handleGetDataSetError = useCallback(
+    () => (error: Error) => {
+      toasts.addError(error, {
+        title:
+          i18n.translate('visualize.vega.failedToGetDataSetErrorDescription', {
+            defaultMessage: 'Failed to get data set: ',
+          }) + (error.message || error.name),
+      });
+    },
+    [toasts]
+  );
+
+  const memorizedReload = useCallback(() => {
+    dataSources.dataSourceService.reload();
+  }, [dataSources.dataSourceService]);
+
+  return (
+    <DataSourceSelectable
+      dataSources={currentDataSources}
+      dataSourceOptionList={dataSourceOptions}
+      setDataSourceOptionList={setDataSourceOptions}
+      onDataSourceSelect={onDataSourceSelect}
+      selectedSources={selectedSources}
+      onGetDataSetError={handleGetDataSetError}
+      onRefresh={memorizedReload}
+      fullWidth
+    />
+  );
+};

src/plugins/visualize/public/application/components/visualize_top_nav.tsx

@@ -51,8 +51,9 @@ import {
 } from '../types';
 import { APP_NAME } from '../visualize_constants';
 import { getTopNavConfig } from '../utils';
-import type { IndexPattern } from '../../../../data/public';
+import type { IndexPattern, DataSourceOption } from '../../../../data/public';
 import chatLogo from './query_assistant_logo.svg';
+import { SourceSelector } from './source_selector';
 
 interface VisualizeTopNavProps {
   currentAppState: VisualizeAppState;
@@ -145,6 +146,7 @@ const TopNav = ({
   const [indexPatterns, setIndexPatterns] = useState<IndexPattern[]>(
     vis.data.indexPattern ? [vis.data.indexPattern] : []
   );
+  const [selectedSource, setSelectedSource] = useState<DataSourceOption>();
   const showDatePicker = () => {
     // tsvb loads without an indexPattern initially (TODO investigate).
     // hide timefilter only if timeFieldName is explicitly undefined.
@@ -254,7 +256,11 @@ const TopNav = ({
 
   const indexName = 'opensearch_dashboards_sample_data_logs';
   const onGenerate = async () => {
-    (window as any).llm.text2vega.invoke({ input: value });
+    if (selectedSource) {
+      (window as any).llm.text2vega.invoke({ prompt: value, index: selectedSource.label });
+    } else {
+      services.notifications.toasts.addWarning('Please select a index');
+    }
   };
 
   return isChromeVisible ? (
@@ -273,6 +279,12 @@ const TopNav = ({
           alignItems="center"
           style={{ maxHeight: '100px' }}
         >
+          <EuiFlexItem grow={3}>
+            <SourceSelector
+              selectedSourceId={selectedSource?.value ?? ''}
+              onChange={(ds) => setSelectedSource(ds)}
+            />
+          </EuiFlexItem>
           <EuiFlexItem grow={9}>
             <EuiInputPopover
               input={