Merge remote-tracking branch 'refs/remotes/origin/main' into insightW…

…ithRAG

# Conflicts:
#	CHANGELOG.md
#	common/constants/llm.ts
#	server/plugin.ts

CHANGELOG.md

@@ -3,8 +3,10 @@
 Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Unreleased
-- fix: make sure $schema always added to LLM generated vega json object([#252](https://github.com/opensearch-project/dashboards-assistant/pull/252))
+- fix: make sure $schema always added to LLM generated vega json object([252](https://github.com/opensearch-project/dashboards-assistant/pull/252))
+- feat: expose a general function for agent execution([268](https://github.com/opensearch-project/dashboards-assistant/pull/268))
 - Fix CVE-2024-4067 ([#269](https://github.com/opensearch-project/dashboards-assistant/pull/269))
+- feat: add a dashboards-assistant trigger in query editor([265](https://github.com/opensearch-project/dashboards-assistant/pull/265))
 
 ### 📈 Features/Enhancements
 

common/constants/llm.ts

@@ -24,6 +24,10 @@ export const TEXT2VIZ_API = {
   TEXT2VEGA: `${API_BASE}/text2vega`,
 };
 
+export const AGENT_API = {
+  EXECUTE: `${API_BASE}/agent/_execute`,
+};
+
 export const SUMMARY_ASSISTANT_API = {
   SUMMARIZE: `${API_BASE}/summary`,
   INSIGHT: `${API_BASE}/insight`,

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,11 +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');
 
@@ -72,6 +77,7 @@ export class AssistantPlugin
   incontextInsightRegistry: IncontextInsightRegistry | undefined;
   private dataSourceService: DataSourceService;
   private resetChatSubscription: Subscription | undefined;
+  private assistantService = new AssistantService();
 
   constructor(initializerContext: PluginInitializerContext) {
     this.config = initializerContext.config.get<ConfigSchema>();
@@ -82,6 +88,7 @@ export class AssistantPlugin
     core: CoreSetup<AssistantPluginStartDependencies>,
     setupDeps: AssistantPluginSetupDependencies
   ): AssistantSetup {
+    this.assistantService.setup();
     this.incontextInsightRegistry = new IncontextInsightRegistry();
     this.incontextInsightRegistry?.setIsEnabled(this.config.incontextInsight.enabled);
     setIncontextInsightRegistry(this.incontextInsightRegistry);
@@ -99,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,
@@ -129,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',
         }),
@@ -185,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) => {
@@ -200,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
       ),
@@ -212,19 +238,38 @@ 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(),
+      assistantClient: assistantServiceStart.client,
     };
   }
 
   public stop() {
     this.dataSourceService.stop();
+    this.assistantService.stop();
     this.resetChatSubscription?.unsubscribe();
   }
 }

public/services/assistant_client.ts

@@ -0,0 +1,35 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { API_BASE } from '../../common/constants/llm';
+import { HttpSetup } from '../../../../src/core/public';
+
+interface Options {
+  dataSourceId?: string;
+}
+
+export class AssistantClient {
+  constructor(private http: HttpSetup) {}
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  executeAgent = (agentId: string, parameters: Record<string, any>, options?: Options) => {
+    return this.http.fetch({
+      method: 'POST',
+      path: `${API_BASE}/agent/_execute`,
+      body: JSON.stringify(parameters),
+      query: { dataSourceId: options?.dataSourceId, agentId },
+    });
+  };
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  executeAgentByName = (agentName: string, parameters: Record<string, any>, options?: Options) => {
+    return this.http.fetch({
+      method: 'POST',
+      path: `${API_BASE}/agent/_execute`,
+      body: JSON.stringify(parameters),
+      query: { dataSourceId: options?.dataSourceId, agentName },
+    });
+  };
+}

public/services/assistant_service.ts

@@ -0,0 +1,26 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { HttpSetup } from '../../../../src/core/public';
+import { AssistantClient } from './assistant_client';
+
+export interface AssistantServiceStart {
+  client: AssistantClient;
+}
+
+export class AssistantService {
+  constructor() {}
+
+  setup() {}
+
+  start(http: HttpSetup): AssistantServiceStart {
+    const assistantClient = new AssistantClient(http);
+    return {
+      client: assistantClient,
+    };
+  }
+
+  stop() {}
+}

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

@@ -16,6 +16,8 @@ import {
 } from '../../../src/plugins/visualizations/public';
 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;
@@ -39,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 {
@@ -61,12 +65,14 @@ 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;
 }
 
 export interface AssistantStart {
   dataSource: DataSourceServiceContract;
+  assistantClient: AssistantClient;
 }
 
 export type StartServices = CoreStart &

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);
+};