feat: added an AI assistant trigger button to query editor (#265) (#278)

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

# Conflicts:
#	CHANGELOG.md

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
(cherry picked from commit 6cce85d)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

opensearch_dashboards.json

@@ -11,7 +11,8 @@
     "opensearchDashboardsReact",
     "opensearchDashboardsUtils",
     "visualizations",
-    "savedObjects"
+    "savedObjects",
+    "uiActions"
   ],
   "optionalPlugins": [
     "dataSource",

public/components/ui_action_context_menu.tsx

@@ -0,0 +1,56 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React, { useState, useRef } from 'react';
+import useAsync from 'react-use/lib/useAsync';
+import { EuiButtonIcon, EuiContextMenu, EuiPopover } from '@elastic/eui';
+
+import { buildContextMenuForActions } from '../../../../src/plugins/ui_actions/public';
+import { AI_ASSISTANT_QUERY_EDITOR_TRIGGER } from '../ui_triggers';
+import { getUiActions } from '../services';
+import assistantTriggerIcon from '../assets/assistant_trigger.svg';
+
+export const ActionContextMenu = () => {
+  const uiActions = getUiActions();
+  const actionsRef = useRef(uiActions.getTriggerActions(AI_ASSISTANT_QUERY_EDITOR_TRIGGER));
+  const [open, setOpen] = useState(false);
+
+  const panels = useAsync(
+    () =>
+      buildContextMenuForActions({
+        actions: actionsRef.current.map((action) => ({
+          action,
+          context: {},
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          trigger: AI_ASSISTANT_QUERY_EDITOR_TRIGGER as any,
+        })),
+        closeMenu: () => setOpen(false),
+      }),
+    []
+  );
+
+  if (actionsRef.current.length === 0) {
+    return null;
+  }
+
+  return (
+    <EuiPopover
+      button={
+        <EuiButtonIcon
+          aria-label="AI assistant trigger button"
+          size="s"
+          iconType={assistantTriggerIcon}
+          onClick={() => setOpen(!open)}
+        />
+      }
+      isOpen={open}
+      panelPaddingSize="none"
+      anchorPosition="downRight"
+      closePopover={() => setOpen(false)}
+    >
+      <EuiContextMenu size="s" initialPanelId={'mainMenu'} panels={panels.value} />
+    </EuiPopover>
+  );
+};

public/index.scss

@@ -179,3 +179,7 @@ button.llm-chat-error-refresh-button.llm-chat-error-refresh-button {
     display: none;
   }
 }
+
+.osdQueryEditorExtensionComponent__assistant-query-actions {
+  margin-left: auto;
+}

public/plugin.tsx

@@ -6,7 +6,7 @@
 import { i18n } from '@osd/i18n';
 import { EuiLoadingSpinner } from '@elastic/eui';
 import React, { lazy, Suspense } from 'react';
-import { Subscription } from 'rxjs';
+import { of, Subscription } from 'rxjs';
 import {
   AppMountParameters,
   AppNavLinkStatus,
@@ -39,12 +39,16 @@ import {
   setNotifications,
   setIncontextInsightRegistry,
   setConfigSchema,
+  setUiActions,
 } from './services';
 import { ConfigSchema } from '../common/types/config';
 import { DataSourceService } from './services/data_source_service';
 import { ASSISTANT_API, DEFAULT_USER_NAME } from '../common/constants/llm';
 import { IncontextInsightProps } from './components/incontext_insight';
 import { AssistantService } from './services/assistant_service';
+import { ActionContextMenu } from './components/ui_action_context_menu';
+import { AI_ASSISTANT_QUERY_EDITOR_TRIGGER, bootstrap } from './ui_triggers';
+import { TEXT2VIZ_APP_ID } from './text2viz';
 
 export const [getCoreStart, setCoreStart] = createGetterSetter<CoreStart>('CoreStart');
 
@@ -102,6 +106,9 @@ export class AssistantPlugin
       return account;
     };
 
+    // setup ui trigger
+    bootstrap(setupDeps.uiActions);
+
     const dataSourceSetupResult = this.dataSourceService.setup({
       uiSettings: core.uiSettings,
       dataSourceManagement: setupDeps.dataSourceManagement,
@@ -132,7 +139,7 @@ export class AssistantPlugin
       });
 
       core.application.register({
-        id: 'text2viz',
+        id: TEXT2VIZ_APP_ID,
         title: i18n.translate('dashboardAssistant.feature.text2viz', {
           defaultMessage: 'Natural language previewer',
         }),
@@ -188,6 +195,19 @@ export class AssistantPlugin
       setupChat();
     }
 
+    setupDeps.data.__enhance({
+      editor: {
+        queryEditorExtension: {
+          id: 'assistant-query-actions',
+          order: 2000,
+          isEnabled$: () => of(true),
+          getComponent: () => {
+            return <ActionContextMenu />;
+          },
+        },
+      },
+    });
+
     return {
       dataSource: dataSourceSetupResult,
       registerMessageRenderer: (contentType, render) => {
@@ -203,6 +223,9 @@ export class AssistantPlugin
       chatEnabled: () => this.config.chat.enabled,
       nextEnabled: () => this.config.next.enabled,
       assistantActions,
+      assistantTriggers: {
+        AI_ASSISTANT_QUERY_EDITOR_TRIGGER,
+      },
       registerIncontextInsight: this.incontextInsightRegistry.register.bind(
         this.incontextInsightRegistry
       ),
@@ -215,12 +238,28 @@ export class AssistantPlugin
     };
   }
 
-  public start(core: CoreStart): AssistantStart {
+  public start(
+    core: CoreStart,
+    { data, uiActions }: AssistantPluginStartDependencies
+  ): AssistantStart {
     const assistantServiceStart = this.assistantService.start(core.http);
     setCoreStart(core);
     setChrome(core.chrome);
     setNotifications(core.notifications);
     setConfigSchema(this.config);
+    setUiActions(uiActions);
+
+    if (this.config.next.enabled) {
+      uiActions.addTriggerAction(AI_ASSISTANT_QUERY_EDITOR_TRIGGER, {
+        id: 'assistant_generate_visualization_action',
+        order: 1,
+        getDisplayName: () => 'Generate visualization',
+        getIconType: () => 'visLine' as const,
+        execute: async () => {
+          core.application.navigateToApp(TEXT2VIZ_APP_ID);
+        },
+      });
+    }
 
     return {
       dataSource: this.dataSourceService.start(),

public/services/index.ts

@@ -4,6 +4,7 @@
  */
 
 import { createGetterSetter } from '../../../../src/plugins/opensearch_dashboards_utils/public';
+import { UiActionsSetup, UiActionsStart } from '../../../../src/plugins/ui_actions/public';
 import { ChromeStart, NotificationsStart } from '../../../../src/core/public';
 import { IncontextInsightRegistry } from './incontext_insight';
 import { ConfigSchema } from '../../common/types/config';
@@ -24,4 +25,6 @@ export const [getNotifications, setNotifications] = createGetterSetter<Notificat
 
 export const [getConfigSchema, setConfigSchema] = createGetterSetter<ConfigSchema>('ConfigSchema');
 
+export const [getUiActions, setUiActions] = createGetterSetter<UiActionsStart>('uiActions');
+
 export { DataSourceService, DataSourceServiceContract } from './data_source_service';

public/text2viz.tsx

@@ -10,6 +10,8 @@ import { Text2Viz } from './components/visualization/text2viz';
 import { OpenSearchDashboardsContextProvider } from '../../../src/plugins/opensearch_dashboards_react/public';
 import { StartServices } from './types';
 
+export const TEXT2VIZ_APP_ID = 'text2viz';
+
 export const renderText2VizApp = (params: AppMountParameters, services: StartServices) => {
   ReactDOM.render(
     <OpenSearchDashboardsContextProvider services={services}>

public/types.ts

@@ -17,6 +17,7 @@ import {
 import { DataPublicPluginSetup, DataPublicPluginStart } from '../../../src/plugins/data/public';
 import { AppMountParameters, CoreStart } from '../../../src/core/public';
 import { AssistantClient } from './services/assistant_client';
+import { UiActionsSetup, UiActionsStart } from '../../../src/plugins/ui_actions/public';
 
 export interface RenderProps {
   props: MessageContentProps;
@@ -40,13 +41,15 @@ export interface AssistantPluginStartDependencies {
   visualizations: VisualizationsStart;
   embeddable: EmbeddableStart;
   dashboard: DashboardStart;
+  uiActions: UiActionsStart;
 }
 
 export interface AssistantPluginSetupDependencies {
   data: DataPublicPluginSetup;
   visualizations: VisualizationsSetup;
   embeddable: EmbeddableSetup;
   dataSourceManagement?: DataSourceManagementPluginSetup;
+  uiActions: UiActionsSetup;
 }
 
 export interface AssistantSetup {
@@ -62,6 +65,7 @@ export interface AssistantSetup {
    */
   nextEnabled: () => boolean;
   assistantActions: Omit<AssistantActions, 'executeAction'>;
+  assistantTriggers: { AI_ASSISTANT_QUERY_EDITOR_TRIGGER: string };
   registerIncontextInsight: IncontextInsightRegistry['register'];
   renderIncontextInsight: (component: React.ReactNode) => React.ReactNode;
 }

public/ui_triggers.ts

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Trigger, UiActionsSetup } from '../../../src/plugins/ui_actions/public';
+
+export const AI_ASSISTANT_QUERY_EDITOR_TRIGGER = 'AI_ASSISTANT_QUERY_EDITOR_TRIGGER';
+
+declare module '../../../src/plugins/ui_actions/public' {
+  export interface TriggerContextMapping {
+    [AI_ASSISTANT_QUERY_EDITOR_TRIGGER]: {};
+  }
+}
+
+const aiAssistantTrigger: Trigger<'AI_ASSISTANT_QUERY_EDITOR_TRIGGER'> = {
+  id: AI_ASSISTANT_QUERY_EDITOR_TRIGGER,
+};
+
+export const bootstrap = (uiActions: UiActionsSetup) => {
+  uiActions.registerTrigger(aiAssistantTrigger);
+};