T2viz ux improvements (#330)

* fix(t2viz): UI fix and improvements

1. Only preselect index pattern that is enabled
2. Display an error toast when PPL query returns no result
3. Disable nature language input if no index pattern was selected

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

---------

Signed-off-by: Yulong Ruan <ruanyl@amazon.com>
(cherry picked from commit 3cdd0e2)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

common/types/config.ts

@@ -26,7 +26,8 @@ export const configSchema = schema.object({
     enabled: schema.boolean({ defaultValue: false }),
   }),
   branding: schema.object({
-    label: schema.string({ defaultValue: '' }),
+    label: schema.maybe(schema.string()),
+    logo: schema.maybe(schema.string()),
   }),
 });
 

public/components/ui_action_context_menu.tsx

@@ -31,6 +31,7 @@ export const ActionContextMenu = (props: Props) => {
           trigger: AI_ASSISTANT_QUERY_EDITOR_TRIGGER as any,
         })),
         closeMenu: () => setOpen(false),
+        title: '',
       }),
     []
   );

public/components/visualization/source_selector.tsx

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useCallback, useMemo, useState, useEffect } from 'react';
+import React, { useCallback, useMemo, useState, useEffect, useRef } from 'react';
 import { i18n } from '@osd/i18n';
 
 import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
@@ -38,6 +38,8 @@ export const SourceSelector = ({
   } = useOpenSearchDashboards<StartServices>();
   const [currentDataSources, setCurrentDataSources] = useState<DataSource[]>([]);
   const [dataSourceOptions, setDataSourceOptions] = useState<DataSourceGroup[]>([]);
+  const onChangeRef = useRef(onChange);
+  onChangeRef.current = onChange;
 
   const selectedSources = useMemo(() => {
     if (selectedSourceId) {
@@ -52,13 +54,20 @@ export const SourceSelector = ({
     return [];
   }, [selectedSourceId, dataSourceOptions]);
 
+  /**
+   * When initialized, select the first non-disabled option
+   */
   useEffect(() => {
     if (
       !selectedSourceId &&
       dataSourceOptions.length > 0 &&
       dataSourceOptions[0].options.length > 0
     ) {
-      onChange(dataSourceOptions[0].options[0]);
+      const options = dataSourceOptions[0].options;
+      const selectedOption = options.find((o) => !o.disabled);
+      if (selectedOption) {
+        onChangeRef.current(selectedOption);
+      }
     }
   }, [selectedSourceId, dataSourceOptions]);
 
@@ -81,7 +90,7 @@ export const SourceSelector = ({
 
   const onSetDataSourceOptions = useCallback(
     async (options: DataSourceGroup[]) => {
-      // Only support opensearch default data source
+      // Only support OpenSearch default data source
       const indexPatternOptions = options.find(
         (item) => item.groupType === DEFAULT_DATA_SOURCE_TYPE
       );
@@ -117,28 +126,30 @@ export const SourceSelector = ({
        * Check each data source to see if text to vega agents are configured or not
        * If not configured, disable the corresponding index pattern from the selection list
        */
-      Object.keys(dataSourceIdToIndexPatternIds).forEach(async (key) => {
-        const res = await assistantService.client.agentConfigExists(
-          [
-            TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
-            TEXT2VEGA_WITH_INSTRUCTIONS_AGENT_CONFIG_ID,
-            TEXT2PPL_AGENT_CONFIG_ID,
-          ],
-          {
-            dataSourceId: key !== 'DEFAULT' ? key : undefined,
-          }
-        );
-        if (!res.exists) {
-          dataSourceIdToIndexPatternIds[key].forEach((indexPatternId) => {
-            indexPatternOptions.options.forEach((option) => {
-              if (option.value === indexPatternId) {
-                option.disabled = true;
-              }
+      const updateIndexPatternPromises = Object.keys(dataSourceIdToIndexPatternIds).map(
+        async (key) => {
+          const res = await assistantService.client.agentConfigExists(
+            [
+              TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
+              TEXT2VEGA_WITH_INSTRUCTIONS_AGENT_CONFIG_ID,
+              TEXT2PPL_AGENT_CONFIG_ID,
+            ],
+            {
+              dataSourceId: key !== 'DEFAULT' ? key : undefined,
+            }
+          );
+          if (!res.exists) {
+            dataSourceIdToIndexPatternIds[key].forEach((indexPatternId) => {
+              indexPatternOptions.options.forEach((option) => {
+                if (option.value === indexPatternId) {
+                  option.disabled = true;
+                }
+              });
             });
-          });
+          }
         }
-      });
-
+      );
+      await Promise.allSettled(updateIndexPatternPromises);
       setDataSourceOptions([indexPatternOptions]);
     },
     [currentDataSources]

public/components/visualization/text2vega.ts

@@ -63,6 +63,9 @@ export class Text2Vega {
                 )
                 // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 .toPromise<any>();
+              if (res.rawResponse.total === 0) {
+                throw new Error(`There is no result with the generated query: '${value.ppl}'.`);
+              }
               return { ...value, sample: res.rawResponse };
             }),
             // call llm to generate vega

public/components/visualization/text2viz.tsx

@@ -67,6 +67,7 @@ export const Text2Viz = () => {
       data,
       uiSettings,
       savedObjects,
+      config,
     },
   } = useOpenSearchDashboards<StartServices>();
 
@@ -106,7 +107,7 @@ export const Text2Viz = () => {
         if (result.error) {
           notifications.toasts.addError(result.error, {
             title: i18n.translate('dashboardAssistant.feature.text2viz.error', {
-              defaultMessage: 'Error while executing text to vega',
+              defaultMessage: 'Error while executing text to visualization',
             }),
           });
         } else {
@@ -198,7 +199,7 @@ export const Text2Viz = () => {
     });
 
     setSubmitting(false);
-  }, [selectedSource, input, status]);
+  }, [selectedSource, input, status, notifications.toasts]);
 
   /**
    * Display the save visualization dialog to persist the current generated visualization
@@ -342,16 +343,17 @@ export const Text2Viz = () => {
             onChange={(e) => setInput(e.target.value)}
             fullWidth
             compressed
-            prepend={<EuiIcon type={chatIcon} />}
+            prepend={<EuiIcon type={config.branding.logo || chatIcon} />}
             placeholder="Generate visualization with a natural language question."
             onKeyDown={(e) => e.key === 'Enter' && onSubmit()}
+            disabled={!selectedSource}
           />
         </EuiFlexItem>
         <EuiFlexItem grow={false}>
           <EuiButtonIcon
             aria-label="submit"
             onClick={onSubmit}
-            isDisabled={loading || input.trim().length === 0}
+            isDisabled={loading || input.trim().length === 0 || !selectedSource}
             display="base"
             size="s"
             iconType="returnKey"

public/plugin.tsx

@@ -200,6 +200,7 @@ export class AssistantPlugin
               ...pluginsStart,
               ...coreStart,
               setHeaderActionMenu: params.setHeaderActionMenu,
+              config: this.config,
             });
 
             return () => {

public/types.ts

@@ -22,6 +22,7 @@ import { ExpressionsSetup, ExpressionsStart } from '../../../src/plugins/express
 import { SavedObjectsStart } from '../../../src/plugins/saved_objects/public';
 
 import { UsageCollectionSetup } from '../../../src/plugins/usage_collection/public';
+import { ConfigSchema } from '../common/types/config';
 
 export interface RenderProps {
   props: MessageContentProps;
@@ -95,6 +96,7 @@ export interface AssistantStart {
 export type StartServices = CoreStart &
   Omit<AssistantPluginStartDependencies, 'savedObjects'> & {
     setHeaderActionMenu: AppMountParameters['setHeaderActionMenu'];
+    config: ConfigSchema;
   };
 
 export interface UserAccount {