[Discover-next] Address comments for search bar extensions and query assist

plugins-extra/query_enhancements/common/index.ts

@@ -18,4 +18,6 @@ export const OPENSEARCH_DATACONNECTIONS_API = {
 
 export const JOBS_ENDPOINT_BASE = '/_plugins/_async_query';
 
+export const BASE_ML_COMMONS_URI = '/_plugins/_ml';
+
 export * from './utils';

plugins-extra/query_enhancements/public/plugin.tsx

@@ -52,7 +52,7 @@ export class QueryEnhancementsPlugin
               initialTo: moment().add(2, 'days').toISOString(),
             },
             showFilterBar: false,
-            extensions: [createQueryAssistExtension(core.http)],
+            extensions: [createQueryAssistExtension(core.http, 'PPL')],
           },
           fields: {
             visualizable: false,

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

@@ -1,8 +1,9 @@
 import { EuiCallOut, EuiCallOutProps } from '@elastic/eui';
+import { i18n } from '@osd/i18n';
 import React from 'react';
 
-type CalloutDismiss = Required<Pick<EuiCallOutProps, 'onDismiss'>>;
-interface QueryAssistCallOutProps extends CalloutDismiss {
+interface QueryAssistCallOutProps extends Required<Pick<EuiCallOutProps, 'onDismiss'>> {
+  language: string;
   type: QueryAssistCallOutType;
 }
 
@@ -14,10 +15,12 @@ export type QueryAssistCallOutType =
   | 'empty_index'
   | 'query_generated';
 
-const EmptyIndexCallOut: React.FC<CalloutDismiss> = (props) => (
+const EmptyIndexCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
     data-test-subj="query-assist-empty-index-callout"
-    title="Select a data source or index to ask a question."
+    title={i18n.translate('queryAssist.callOut.emptyIndex.title', {
+      defaultMessage: 'Select a data source or index to ask a question.',
+    })}
     size="s"
     color="warning"
     iconType="iInCircle"
@@ -26,10 +29,12 @@ const EmptyIndexCallOut: React.FC<CalloutDismiss> = (props) => (
   />
 );
 
-const ProhibitedQueryCallOut: React.FC<CalloutDismiss> = (props) => (
+const ProhibitedQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
     data-test-subj="query-assist-guard-callout"
-    title="I am unable to respond to this query. Try another question."
+    title={i18n.translate('queryAssist.callOut.prohibitedQuery.title', {
+      defaultMessage: 'I am unable to respond to this query. Try another question.',
+    })}
     size="s"
     color="danger"
     iconType="alert"
@@ -38,10 +43,13 @@ const ProhibitedQueryCallOut: React.FC<CalloutDismiss> = (props) => (
   />
 );
 
-const EmptyQueryCallOut: React.FC<CalloutDismiss> = (props) => (
+const EmptyQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
     data-test-subj="query-assist-empty-query-callout"
-    title="Enter a natural language question to automatically generate a query to view results."
+    title={i18n.translate('queryAssist.callOut.emptyQuery.title', {
+      defaultMessage:
+        'Enter a natural language question to automatically generate a query to view results.',
+    })}
     size="s"
     color="warning"
     iconType="iInCircle"
@@ -50,10 +58,12 @@ const EmptyQueryCallOut: React.FC<CalloutDismiss> = (props) => (
   />
 );
 
-const PPLGeneratedCallOut: React.FC<CalloutDismiss> = (props) => (
+const QueryGeneratedCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
-    data-test-subj="query-assist-ppl-callout"
-    title="PPL query generated"
+    data-test-subj="query-assist-query-generated-callout"
+    title={`${props.language} ${i18n.translate('queryAssist.callOut.queryGenerated.title', {
+      defaultMessage: `query generated. If there are any issues with the response, try adding more context to the question or a new question to submit.`,
+    })}`}
     size="s"
     color="success"
     iconType="check"
@@ -65,13 +75,13 @@ const PPLGeneratedCallOut: React.FC<CalloutDismiss> = (props) => (
 export const QueryAssistCallOut: React.FC<QueryAssistCallOutProps> = (props) => {
   switch (props.type) {
     case 'empty_query':
-      return <EmptyQueryCallOut onDismiss={props.onDismiss} />;
+      return <EmptyQueryCallOut {...props} />;
     case 'empty_index':
-      return <EmptyIndexCallOut onDismiss={props.onDismiss} />;
+      return <EmptyIndexCallOut {...props} />;
     case 'invalid_query':
-      return <ProhibitedQueryCallOut onDismiss={props.onDismiss} />;
+      return <ProhibitedQueryCallOut {...props} />;
     case 'query_generated':
-      return <PPLGeneratedCallOut onDismiss={props.onDismiss} />;
+      return <QueryGeneratedCallOut {...props} />;
     default:
       break;
   }

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

@@ -1,5 +1,5 @@
 import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui';
-import React, { SyntheticEvent, useEffect, useMemo, useRef, useState } from 'react';
+import React, { SyntheticEvent, 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';
@@ -11,6 +11,7 @@ import { QueryAssistInput } from './query_assist_input';
 import { QueryAssistSubmitButton } from './submit_button';
 
 interface QueryAssistInputProps {
+  language: string;
   dependencies: SearchBarExtensionDependencies;
 }
 
@@ -25,17 +26,12 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   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 selectedIndexPattern = props.dependencies.indexPatterns?.at(0);
+  const selectedIndex =
+    selectedIndexPattern &&
+    (typeof selectedIndexPattern === 'string' ? selectedIndexPattern : selectedIndexPattern.title);
   const previousQuestionRef = useRef<string>();
 
-  useEffect(() => {
-    mounted.current = true;
-    return () => {
-      mounted.current = false;
-    };
-  }, []);
-
   const onSubmit = async (e: SyntheticEvent) => {
     e.preventDefault();
     if (!inputRef.current?.value) {
@@ -52,10 +48,9 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
     const params = {
       question: inputRef.current.value,
       index: selectedIndex,
-      language: 'PPL',
+      language: props.language,
     };
     const { response, error } = await generateQuery(params);
-    if (!mounted.current) return;
     if (error) {
       if (error instanceof ProhibitedQueryError) {
         setCallOutType('invalid_query');
@@ -89,7 +84,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
           </EuiFlexItem>
         </EuiFlexGroup>
       </EuiFormRow>
-      <QueryAssistCallOut type={callOutType} onDismiss={dismissCallout} />
+      <QueryAssistCallOut language={props.language} type={callOutType} onDismiss={dismissCallout} />
     </EuiForm>
   );
 };

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

@@ -20,18 +20,49 @@ export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
   const [suggestionIndex, setSuggestionIndex] = useState<number | null>(null);
   const [value, setValue] = useState(props.initialValue ?? '');
 
-  const recentSearchSuggestions = useMemo(() => {
+  const sampleDataSuggestions = useMemo(() => {
+    switch (props.selectedIndex) {
+      case 'opensearch_dashboards_sample_data_ecommerce':
+        return [
+          'How many unique customers placed orders this week?',
+          'Count the number of orders grouped by manufacturer and category',
+          'find customers with first names like Eddie',
+        ];
+
+      case 'opensearch_dashboards_sample_data_logs':
+        return [
+          'Are there any errors in my logs?',
+          'How many requests were there grouped by response code last week?',
+          "What's the average request size by week?",
+        ];
+
+      case 'opensearch_dashboards_sample_data_flights':
+        return [
+          'how many flights were there this week grouped by destination country?',
+          'what were the longest flight delays this week?',
+          'what carriers have the furthest flights?',
+        ];
+
+      default:
+        return [];
+    }
+  }, [props.selectedIndex]);
+
+  const suggestions = useMemo(() => {
     if (!props.persistedLog) return [];
     return props.persistedLog
       .get()
-      .filter((recentSearch) => recentSearch.includes(value))
-      .map((recentSearch) => ({
+      .concat(sampleDataSuggestions)
+      .filter(
+        (suggestion, i, array) => array.indexOf(suggestion) === i && suggestion.includes(value)
+      )
+      .map((suggestion) => ({
         type: QuerySuggestionTypes.RecentSearch,
-        text: recentSearch,
+        text: suggestion,
         start: 0,
         end: value.length,
       }));
-  }, [props.persistedLog, value]);
+  }, [props.persistedLog, value, sampleDataSuggestions]);
 
   return (
     <EuiOutsideClickDetector onOutsideClick={() => setIsSuggestionsVisible(false)}>
@@ -54,7 +85,7 @@ export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
         <EuiPortal>
           <SuggestionsComponent
             show={isSuggestionsVisible}
-            suggestions={recentSearchSuggestions}
+            suggestions={suggestions}
             index={suggestionIndex}
             onClick={(suggestion) => {
               if (!props.inputRef.current) return;

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

@@ -13,7 +13,7 @@ export const QueryAssistSubmitButton: React.FC<SubmitButtonProps> = (props) => {
       isDisabled={props.isDisabled}
       size="s"
       type="submit"
-      aria-label="submit-question"
+      aria-label="Submit question to query assistant"
     />
   );
 };

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

@@ -5,15 +5,26 @@ import { QueryAssistParameters, QueryAssistResponse } from '../../../common/quer
 import { formatError } from '../utils';
 
 export const useGenerateQuery = () => {
+  const mounted = useRef(false);
   const [loading, setLoading] = useState(false);
   const abortControllerRef = useRef<AbortController>();
   const { services } = useOpenSearchDashboards<IDataPluginServices>();
 
-  useEffect(() => () => abortControllerRef.current?.abort(), []);
+  useEffect(() => {
+    mounted.current = true;
+    return () => {
+      mounted.current = false;
+      if (abortControllerRef.current) {
+        abortControllerRef.current.abort();
+        abortControllerRef.current = undefined;
+      }
+    };
+  }, []);
 
   const generateQuery = async (
     params: QueryAssistParameters
   ): Promise<{ response?: QueryAssistResponse; error?: Error }> => {
+    abortControllerRef.current?.abort();
     abortControllerRef.current = new AbortController();
     setLoading(true);
     try {
@@ -24,12 +35,13 @@ export const useGenerateQuery = () => {
           signal: abortControllerRef.current?.signal,
         }
       );
-      return { response };
+      if (mounted.current) return { response };
     } catch (error) {
-      return { error: formatError(error) };
+      if (mounted.current) return { error: formatError(error) };
     } finally {
-      setLoading(false);
+      if (mounted.current) setLoading(false);
     }
+    return {};
   };
 
   return { generateQuery, loading, abortControllerRef };

plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx

@@ -3,22 +3,27 @@ import { HttpSetup } from 'opensearch-dashboards/public';
 import { QueryAssistBar } from '../components';
 import { SearchBarExtensionConfig } from '../../../../../src/plugins/data/public/ui/search_bar_extensions';
 
-export const createQueryAssistExtension = (http: HttpSetup): SearchBarExtensionConfig => {
+export const createQueryAssistExtension = (
+  http: HttpSetup,
+  language: string
+): SearchBarExtensionConfig => {
   return {
-    id: 'query-assist-ppl',
+    id: 'query-assist',
     order: 1000,
     isEnabled: (() => {
       let agentConfigured: boolean;
       return async () => {
         if (agentConfigured === undefined) {
           agentConfigured = await http
-            .get<{ configured: boolean }>('/api/ql/query_assist/configured/PPL')
+            .get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`)
             .then((response) => response.configured)
             .catch(() => false);
         }
         return agentConfigured;
       };
     })(),
-    getComponent: (dependencies) => <QueryAssistBar dependencies={dependencies} />,
+    getComponent: (dependencies) => (
+      <QueryAssistBar language={language} dependencies={dependencies} />
+    ),
   };
 };

plugins-extra/query_enhancements/server/routes/query_assist/agents.ts

@@ -1,8 +1,8 @@
 import { ApiResponse } from '@opensearch-project/opensearch';
 import { RequestBody, TransportRequestPromise } from '@opensearch-project/opensearch/lib/Transport';
 import { RequestHandlerContext } from 'src/core/server';
+import { BASE_ML_COMMONS_URI } from '../../../common';
 
-const ML_COMMONS_API_PREFIX = '/_plugins/_ml';
 const AGENT_REQUEST_OPTIONS = {
   /**
    * It is time-consuming for LLM to generate final answer
@@ -30,7 +30,7 @@ export const getAgentIdByConfig = async (
   try {
     const response = (await client.transport.request({
       method: 'GET',
-      path: `${ML_COMMONS_API_PREFIX}/config/${configName}`,
+      path: `${BASE_ML_COMMONS_URI}/config/${configName}`,
     })) as ApiResponse<{ type: string; configuration: { agent_id?: string } }>;
 
     if (!response || response.body.configuration.agent_id === undefined) {
@@ -54,7 +54,7 @@ export const requestAgentByConfig = async (options: {
   return client.transport.request(
     {
       method: 'POST',
-      path: `${ML_COMMONS_API_PREFIX}/agents/${agentId}/_execute`,
+      path: `${BASE_ML_COMMONS_URI}/agents/${agentId}/_execute`,
       body,
     },
     AGENT_REQUEST_OPTIONS

src/plugins/data/public/ui/query_editor/_query_editor.scss

@@ -8,6 +8,13 @@
   width: 500px;
 }
 
+.osdQueryEditorHeader {
+  max-height: 400px;
+
+  // TODO fix styling: with "overflow: auto" the scroll bar appears although the content is below max-height
+  // overflow: auto;
+}
+
 @include euiBreakpoint("xs", "s") {
   .osdQueryEditor--withDatePicker {
     > :first-child {

src/plugins/data/public/ui/query_editor/query_editor.tsx

@@ -55,6 +55,7 @@
   className?: string;
   isInvalid?: boolean;
   queryEditorHeaderRef: React.RefObject<HTMLDivElement>;
+  queryEditorHeaderClassName?: string;
 }
 
 interface Props extends QueryEditorProps {
@@ -492,6 +493,11 @@
     // );
     const className = classNames(this.props.className);
 
+    const queryEditorHeaderClassName = classNames(
+      'osdQueryEditorHeader',
+      this.props.queryEditorHeaderClassName
+    );
+
     const queryLanguageSwitcher = (
       <QueryLanguageSwitcher
         language={this.props.query.language}
@@ -519,7 +525,7 @@
               </EuiFlexGroup>
             </EuiFlexItem>
             <EuiFlexItem onClick={this.onClickInput} grow={true}>
-              <div ref={this.props.queryEditorHeaderRef} />
+              <div ref={this.props.queryEditorHeaderRef} className={queryEditorHeaderClassName} />
               <CodeEditor
                 height={70}
                 languageId="xjson"

src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx

@@ -251,8 +251,8 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
     return (
       <SearchBarExtensions
         configs={props.queryEnhancements?.get(queryLanguage!)?.searchBar?.extensions}
-        dependencies={{ indexPatterns: props.indexPatterns }}
-        portalInsert={{ sibling: queryEditorHeaderRef.current, position: 'before' }}
+        portalContainer={queryEditorHeaderRef.current}
+        indexPatterns={props.indexPatterns}
       />
     );
   }