[Discover-next] add query assist to query enhancements plugin (#6895)

it adds query assist specific logic in query enhancements plugin to show a UI above the PPL search bar if user has configured PPL agent.

Issues Resolved: #6820

* add query assist to query enhancements

Signed-off-by: Joshua Li <joshuali925@gmail.com>

* align language to uppercase

Signed-off-by: Joshua Li <joshuali925@gmail.com>

* pick PR 6167

Signed-off-by: Joshua Li <joshuali925@gmail.com>

* use useState hooks for query assist

There is a bug in data explorer `AppContainer` where memorized
`DiscoverCanvas` gets unmounted after `setQuery`. PR 6167 works around
it by memorizing `AppContainer`. As query assist is no longer being
unmounted, we can use proper hooks to persist state now.

Signed-off-by: Joshua Li <joshuali925@gmail.com>

* Revert "pick PR 6167"

This reverts commit acb0d41.

Wait for official 6167 to merge to avoid conflict

Signed-off-by: Joshua Li <joshuali925@gmail.com>

* address comments for PR 6894

Signed-off-by: Joshua Li <joshuali925@gmail.com>

---------

Signed-off-by: Joshua Li <joshuali925@gmail.com>

plugins-extra/query_enhancements/common/query_assist/index.ts

@@ -0,0 +1,14 @@
+import { TimeRange } from '../../../../src/plugins/data/common';
+
+export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrails triggered' };
+
+export interface QueryAssistResponse {
+  query: string;
+  timeRange?: TimeRange;
+}
+
+export interface QueryAssistParameters {
+  question: string;
+  index: string;
+  language: string;
+}

plugins-extra/query_enhancements/opensearch_dashboards.json

@@ -6,5 +6,5 @@
   "ui": true,
   "requiredPlugins": ["data"],
   "optionalPlugins": ["home"],
-  "requiredBundles": []
+  "requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"]
 }

plugins-extra/query_enhancements/public/plugin.ts → plugins-extra/query_enhancements/public/plugin.tsx

@@ -1,20 +1,29 @@
 import moment from 'moment';
 import { CoreSetup, CoreStart, Plugin } from '../../../src/core/public';
+import { IStorageWrapper, Storage } from '../../../src/plugins/opensearch_dashboards_utils/public';
+import { createQueryAssistExtension } from './query_assist';
+import { PPLQlSearchInterceptor } from './search/ppl_search_interceptor';
+import { SQLQlSearchInterceptor } from './search/sql_search_interceptor';
+import { setData, setStorage } from './services';
 import {
   QueryEnhancementsPluginSetup,
-  QueryEnhancementsPluginStart,
   QueryEnhancementsPluginSetupDependencies,
+  QueryEnhancementsPluginStart,
+  QueryEnhancementsPluginStartDependencies,
 } from './types';
-import { PPLQlSearchInterceptor } from './search/ppl_search_interceptor';
-import { SQLQlSearchInterceptor } from './search/sql_search_interceptor';
 
 export class QueryEnhancementsPlugin
   implements Plugin<QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart> {
+  private readonly storage: IStorageWrapper;
+
+  constructor() {
+    this.storage = new Storage(window.localStorage);
+  }
+
   public setup(
     core: CoreSetup,
     { data }: QueryEnhancementsPluginSetupDependencies
   ): QueryEnhancementsPluginSetup {
-
     const pplSearchInterceptor = new PPLQlSearchInterceptor({
       toasts: core.notifications.toasts,
       http: core.http,
@@ -43,6 +52,7 @@ export class QueryEnhancementsPlugin
               initialTo: moment().add(2, 'days').toISOString(),
             },
             showFilterBar: false,
+            extensions: [createQueryAssistExtension(core.http)],
           },
           fields: {
             visualizable: false,
@@ -75,7 +85,12 @@ export class QueryEnhancementsPlugin
     return {};
   }
 
-  public start(core: CoreStart): QueryEnhancementsPluginStart {
+  public start(
+    core: CoreStart,
+    deps: QueryEnhancementsPluginStartDependencies
+  ): QueryEnhancementsPluginStart {
+    setStorage(this.storage);
+    setData(deps.data);
     return {};
   }
 

plugins-extra/query_enhancements/public/query_assist/components/call_outs.tsx

@@ -0,0 +1,79 @@
+import { EuiCallOut, EuiCallOutProps } from '@elastic/eui';
+import React from 'react';
+
+type CalloutDismiss = Required<Pick<EuiCallOutProps, 'onDismiss'>>;
+interface QueryAssistCallOutProps extends CalloutDismiss {
+  type: QueryAssistCallOutType;
+}
+
+export type QueryAssistCallOutType =
+  | undefined
+  | 'invalid_query'
+  | 'prohibited_query'
+  | 'empty_query'
+  | 'empty_index'
+  | 'query_generated';
+
+const EmptyIndexCallOut: React.FC<CalloutDismiss> = (props) => (
+  <EuiCallOut
+    data-test-subj="query-assist-empty-index-callout"
+    title="Select a data source or index to ask a question."
+    size="s"
+    color="warning"
+    iconType="iInCircle"
+    dismissible
+    onDismiss={props.onDismiss}
+  />
+);
+
+const ProhibitedQueryCallOut: React.FC<CalloutDismiss> = (props) => (
+  <EuiCallOut
+    data-test-subj="query-assist-guard-callout"
+    title="I am unable to respond to this query. Try another question."
+    size="s"
+    color="danger"
+    iconType="alert"
+    dismissible
+    onDismiss={props.onDismiss}
+  />
+);
+
+const EmptyQueryCallOut: React.FC<CalloutDismiss> = (props) => (
+  <EuiCallOut
+    data-test-subj="query-assist-empty-query-callout"
+    title="Enter a natural language question to automatically generate a query to view results."
+    size="s"
+    color="warning"
+    iconType="iInCircle"
+    dismissible
+    onDismiss={props.onDismiss}
+  />
+);
+
+const PPLGeneratedCallOut: React.FC<CalloutDismiss> = (props) => (
+  <EuiCallOut
+    data-test-subj="query-assist-ppl-callout"
+    title="PPL query generated"
+    size="s"
+    color="success"
+    iconType="check"
+    dismissible
+    onDismiss={props.onDismiss}
+  />
+);
+
+export const QueryAssistCallOut: React.FC<QueryAssistCallOutProps> = (props) => {
+  switch (props.type) {
+    case 'empty_query':
+      return <EmptyQueryCallOut onDismiss={props.onDismiss} />;
+    case 'empty_index':
+      return <EmptyIndexCallOut onDismiss={props.onDismiss} />;
+    case 'invalid_query':
+      return <ProhibitedQueryCallOut onDismiss={props.onDismiss} />;
+    case 'query_generated':
+      return <PPLGeneratedCallOut onDismiss={props.onDismiss} />;
+    default:
+      break;
+  }
+  return null;
+};

plugins-extra/query_enhancements/public/query_assist/components/index.ts

@@ -0,0 +1 @@
+export { QueryAssistBar } from './query_assist_bar';

plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx

@@ -0,0 +1,95 @@
+import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui';
+import React, { SyntheticEvent, useEffect, useMemo, useRef, useState } from 'react';
+import { IDataPluginServices, PersistedLog } from '../../../../../src/plugins/data/public';
+import { SearchBarExtensionDependencies } from '../../../../../src/plugins/data/public/ui/search_bar_extensions/search_bar_extension';
+import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
+import { getStorage } from '../../services';
+import { useGenerateQuery } from '../hooks';
+import { getPersistedLog, ProhibitedQueryError } from '../utils';
+import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs';
+import { QueryAssistInput } from './query_assist_input';
+import { QueryAssistSubmitButton } from './submit_button';
+
+interface QueryAssistInputProps {
+  dependencies: SearchBarExtensionDependencies;
+}
+
+export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
+  const { services } = useOpenSearchDashboards<IDataPluginServices>();
+  const inputRef = useRef<HTMLInputElement>(null);
+  const storage = getStorage();
+  const persistedLog: PersistedLog = useMemo(
+    () => getPersistedLog(services.uiSettings, storage, 'query-assist'),
+    [services.uiSettings, storage]
+  );
+  const { generateQuery, loading } = useGenerateQuery();
+  const [callOutType, setCallOutType] = useState<QueryAssistCallOutType>();
+  const dismissCallout = () => setCallOutType(undefined);
+  const mounted = useRef(false);
+  const selectedIndex = props.dependencies.indexPatterns?.at(0)?.title;
+  const previousQuestionRef = useRef<string>();
+
+  useEffect(() => {
+    mounted.current = true;
+    return () => {
+      mounted.current = false;
+    };
+  }, []);
+
+  const onSubmit = async (e: SyntheticEvent) => {
+    e.preventDefault();
+    if (!inputRef.current?.value) {
+      setCallOutType('empty_query');
+      return;
+    }
+    if (!selectedIndex) {
+      setCallOutType('empty_index');
+      return;
+    }
+    dismissCallout();
+    previousQuestionRef.current = inputRef.current.value;
+    persistedLog.add(inputRef.current.value);
+    const params = {
+      question: inputRef.current.value,
+      index: selectedIndex,
+      language: 'PPL',
+    };
+    const { response, error } = await generateQuery(params);
+    if (!mounted.current) return;
+    if (error) {
+      if (error instanceof ProhibitedQueryError) {
+        setCallOutType('invalid_query');
+      } else {
+        services.notifications.toasts.addError(error, { title: 'Failed to generate results' });
+      }
+    } else if (response) {
+      services.data.query.queryString.setQuery({
+        query: response.query,
+        language: params.language,
+      });
+      if (response.timeRange) services.data.query.timefilter.timefilter.setTime(response.timeRange);
+      setCallOutType('query_generated');
+    }
+  };
+
+  return (
+    <EuiForm component="form" onSubmit={onSubmit}>
+      <EuiFormRow fullWidth>
+        <EuiFlexGroup gutterSize="s" responsive={false} alignItems="center">
+          <EuiFlexItem>
+            <QueryAssistInput
+              inputRef={inputRef}
+              persistedLog={persistedLog}
+              selectedIndex={selectedIndex}
+              previousQuestion={previousQuestionRef.current}
+            />
+          </EuiFlexItem>
+          <EuiFlexItem grow={false}>
+            <QueryAssistSubmitButton isDisabled={loading} />
+          </EuiFlexItem>
+        </EuiFlexGroup>
+      </EuiFormRow>
+      <QueryAssistCallOut type={callOutType} onDismiss={dismissCallout} />
+    </EuiForm>
+  );
+};

plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx

@@ -0,0 +1,75 @@
+import { EuiFieldText, EuiIcon, EuiOutsideClickDetector, EuiPortal } from '@elastic/eui';
+import React, { useMemo, useState } from 'react';
+import { PersistedLog, QuerySuggestionTypes } from '../../../../../src/plugins/data/public';
+import assistantLogo from '../../assets/query_assist_logo.svg';
+import { getData } from '../../services';
+
+interface QueryAssistInputProps {
+  inputRef: React.RefObject<HTMLInputElement>;
+  persistedLog: PersistedLog;
+  initialValue?: string;
+  selectedIndex?: string;
+  previousQuestion?: string;
+}
+
+export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
+  const {
+    ui: { SuggestionsComponent },
+  } = getData();
+  const [isSuggestionsVisible, setIsSuggestionsVisible] = useState(false);
+  const [suggestionIndex, setSuggestionIndex] = useState<number | null>(null);
+  const [value, setValue] = useState(props.initialValue ?? '');
+
+  const recentSearchSuggestions = useMemo(() => {
+    if (!props.persistedLog) return [];
+    return props.persistedLog
+      .get()
+      .filter((recentSearch) => recentSearch.includes(value))
+      .map((recentSearch) => ({
+        type: QuerySuggestionTypes.RecentSearch,
+        text: recentSearch,
+        start: 0,
+        end: value.length,
+      }));
+  }, [props.persistedLog, value]);
+
+  return (
+    <EuiOutsideClickDetector onOutsideClick={() => setIsSuggestionsVisible(false)}>
+      <div>
+        <EuiFieldText
+          inputRef={props.inputRef}
+          value={value}
+          onClick={() => setIsSuggestionsVisible(true)}
+          onChange={(e) => setValue(e.target.value)}
+          onKeyDown={() => setIsSuggestionsVisible(true)}
+          placeholder={
+            props.previousQuestion ||
+            (props.selectedIndex
+              ? `Ask a natural language question about ${props.selectedIndex} to generate a query`
+              : 'Select an index pattern to ask a question')
+          }
+          prepend={<EuiIcon type={assistantLogo} />}
+          fullWidth
+        />
+        <EuiPortal>
+          <SuggestionsComponent
+            show={isSuggestionsVisible}
+            suggestions={recentSearchSuggestions}
+            index={suggestionIndex}
+            onClick={(suggestion) => {
+              if (!props.inputRef.current) return;
+              setValue(suggestion.text);
+              setIsSuggestionsVisible(false);
+              setSuggestionIndex(null);
+              props.inputRef.current.focus();
+            }}
+            onMouseEnter={(i) => setSuggestionIndex(i)}
+            loadMore={() => {}}
+            queryBarRect={props.inputRef.current?.getBoundingClientRect()}
+            size="s"
+          />
+        </EuiPortal>
+      </div>
+    </EuiOutsideClickDetector>
+  );
+};

plugins-extra/query_enhancements/public/query_assist/components/submit_button.tsx

@@ -0,0 +1,19 @@
+import React from 'react';
+import { EuiButtonIcon } from '@elastic/eui';
+
+interface SubmitButtonProps {
+  isDisabled: boolean;
+}
+
+export const QueryAssistSubmitButton: React.FC<SubmitButtonProps> = (props) => {
+  return (
+    <EuiButtonIcon
+      iconType="returnKey"
+      display="base"
+      isDisabled={props.isDisabled}
+      size="s"
+      type="submit"
+      aria-label="submit-question"
+    />
+  );
+};

plugins-extra/query_enhancements/public/query_assist/hooks/index.ts

@@ -0,0 +1 @@
+export * from './use_generate';