Add query assist to query enhancements plugin

This PR picks

[Discover-next] add query assist to query enhancements plugin opensearch-project/OpenSearch-Dashboards#6895
[Discover-next] Address comments for search bar extensions and query assist opensearch-project/OpenSearch-Dashboards#6933
[Discover-next] Support data sources for query assist opensearch-project/OpenSearch-Dashboards#6972
adds query assist banner

    a callout that advertises query assist when user has at least 1 language with query assist configured and is not on a configured language

adds index selector for local cluster only

    This is a temporary solution given that in discover the index pattern
    selector will be removed. Before datasource and dataset selectors are
    added, query assist will rely on this index pattern selector to
    determine which index user wants to query.

Signed-off-by: Joshua Li <[email protected]>

.gitignore

@@ -1,2 +1,3 @@
 /build
-/target
+/target
+/coverage/

babel.config.js

@@ -0,0 +1,25 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// babelrc doesn't respect NODE_PATH anymore but using require does.
+// Alternative to install them locally in node_modules
+module.exports = function (api) {
+  // ensure env is test so that this config won't impact build or dev server
+  if (api.env('test')) {
+    return {
+      presets: [
+        require('@babel/preset-env', {
+          useBuiltIns: false,
+          targets: {
+            node: 'current',
+          },
+        }),
+        require('@babel/preset-react'),
+        require('@babel/preset-typescript'),
+      ],
+    };
+  }
+  return {};
+};

common/config.ts

@@ -7,6 +7,17 @@ import { schema, TypeOf } from '@osd/config-schema';
 
 export const configSchema = schema.object({
   enabled: schema.boolean({ defaultValue: true }),
+  queryAssist: schema.object({
+    supportedLanguages: schema.arrayOf(
+      schema.object({
+        language: schema.string(),
+        agentConfig: schema.string(),
+      }),
+      {
+        defaultValue: [{ language: 'PPL', agentConfig: 'os_query_assist_ppl' }],
+      }
+    ),
+  }),
 });
 
 export type ConfigSchema = TypeOf<typeof configSchema>;

common/constants.ts

@@ -17,6 +17,10 @@ export const API = {
   SEARCH: `${BASE_API}/search`,
   PPL_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.PPL}`,
   SQL_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.SQL}`,
+  QUERY_ASSIST: {
+    LANGUAGES: `${BASE_API}/assist/languages`,
+    GENERATE: `${BASE_API}/assist/generate`,
+  },
 };
 
 export const URI = {
@@ -34,3 +38,5 @@ export const OPENSEARCH_API = {
 };
 
 export const UI_SETTINGS = {};
+
+export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrails triggered' };

common/query_assist/index.ts

@@ -0,0 +1 @@
+export { QueryAssistParameters, QueryAssistResponse } from './types';

common/query_assist/types.ts

@@ -0,0 +1,14 @@
+import { TimeRange } from '../../../../src/plugins/data/common';
+
+export interface QueryAssistResponse {
+  query: string;
+  timeRange?: TimeRange;
+}
+
+export interface QueryAssistParameters {
+  question: string;
+  index: string;
+  language: string;
+  // for MDS
+  dataSourceId?: string;
+}

opensearch_dashboards.json

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

package.json

@@ -6,6 +6,7 @@
   "scripts": {
     "build": "yarn plugin-helpers build",
     "plugin-helpers": "node ../../scripts/plugin_helpers",
+    "test": "../../node_modules/.bin/jest --config ./test/jest.config.js",
     "osd": "node ../../scripts/osd"
   },
   "dependencies": {},

public/index.ts

@@ -3,12 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { PluginInitializerContext } from '../../../src/core/public';
 import './index.scss';
-
 import { QueryEnhancementsPlugin } from './plugin';
 
-export function plugin() {
-  return new QueryEnhancementsPlugin();
+export function plugin(initializerContext: PluginInitializerContext) {
+  return new QueryEnhancementsPlugin(initializerContext);
 }
 
 export { QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart } from './types';

public/plugin.tsx

@@ -4,8 +4,10 @@
  */
 
 import moment from 'moment';
-import { CoreSetup, CoreStart, Plugin } from '../../../src/core/public';
+import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../src/core/public';
 import { IStorageWrapper, Storage } from '../../../src/plugins/opensearch_dashboards_utils/public';
+import { ConfigSchema } from '../common/config';
+import { createQueryAssistExtension } from './query_assist';
 import { PPLSearchInterceptor, SQLSearchInterceptor } from './search';
 import { setData, setStorage } from './services';
 import {
@@ -15,11 +17,15 @@ import {
   QueryEnhancementsPluginStartDependencies,
 } from './types';
 
+export type PublicConfig = Pick<ConfigSchema, 'queryAssist'>;
+
 export class QueryEnhancementsPlugin
   implements Plugin<QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart> {
   private readonly storage: IStorageWrapper;
+  private readonly config: PublicConfig;
 
-  constructor() {
+  constructor(initializerContext: PluginInitializerContext) {
+    this.config = initializerContext.config.get<PublicConfig>();
     this.storage = new Storage(window.localStorage);
   }
 
@@ -86,6 +92,12 @@ export class QueryEnhancementsPlugin
       },
     });
 
+    data.__enhance({
+      ui: {
+        queryEditorExtension: createQueryAssistExtension(core.http, this.config),
+      },
+    });
+
     return {};
   }
 

public/query_assist/components/__snapshots__/call_outs.test.tsx.snap

@@ -0,0 +1,163 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`CallOuts spec should display empty_index call out 1`] = `
+<div>
+  <div
+    class="euiCallOut euiCallOut--warning euiCallOut--small"
+    data-test-subj="query-assist-empty-index-callout"
+  >
+    <div
+      class="euiCallOutHeader"
+    >
+      <svg
+        aria-hidden="true"
+        class="euiIcon euiIcon--medium euiIcon--inherit euiCallOutHeader__icon"
+        focusable="false"
+        height="16"
+        role="img"
+        viewBox="0 0 16 16"
+        width="16"
+        xmlns="http://www.w3.org/2000/svg"
+      >
+        <path
+          d="M7.5 11.508 7.468 8H6.25V7h2.401l.03 3.508H9.8v1H7.5Zm-.25-6.202a.83.83 0 0 1 .207-.577c.137-.153.334-.229.59-.229.256 0 .454.076.594.23.14.152.209.345.209.576 0 .228-.07.417-.21.568-.14.15-.337.226-.593.226-.256 0-.453-.075-.59-.226a.81.81 0 0 1-.207-.568ZM8 13A5 5 0 1 0 8 3a5 5 0 0 0 0 10Zm0 1A6 6 0 1 1 8 2a6 6 0 0 1 0 12Z"
+          fill-rule="evenodd"
+        />
+      </svg>
+      <span
+        class="euiCallOutHeader__title"
+      >
+        <span>
+          Select a data source or index to ask a question.
+        </span>
+      </span>
+    </div>
+  </div>
+</div>
+`;
+
+exports[`CallOuts spec should display empty_query call out 1`] = `
+<div>
+  <div
+    class="euiCallOut euiCallOut--warning euiCallOut--small"
+    data-test-subj="query-assist-empty-query-callout"
+  >
+    <div
+      class="euiCallOutHeader"
+    >
+      <svg
+        aria-hidden="true"
+        class="euiIcon euiIcon--medium euiIcon--inherit euiIcon-isLoading euiCallOutHeader__icon"
+        focusable="false"
+        height="16"
+        role="img"
+        viewBox="0 0 16 16"
+        width="16"
+        xmlns="http://www.w3.org/2000/svg"
+      >
+        <path
+          d="M5.277 10.088c.02.014.04.03.057.047.582.55 1.134.812 1.666.812.586 0 1.84-.293 3.713-.88L9 6.212V2H7v4.212l-1.723 3.876Zm-.438.987L3.539 14h8.922l-1.32-2.969C9.096 11.677 7.733 12 7 12c-.74 0-1.463-.315-2.161-.925ZM6 2H5V1h6v1h-1v4l3.375 7.594A1 1 0 0 1 12.461 15H3.54a1 1 0 0 1-.914-1.406L6 6V2Z"
+        />
+      </svg>
+      <span
+        class="euiCallOutHeader__title"
+      >
+        <span>
+          Enter a natural language question to automatically generate a query to view results.
+        </span>
+      </span>
+    </div>
+  </div>
+</div>
+`;
+
+exports[`CallOuts spec should display invalid_query call out 1`] = `
+<div>
+  <div
+    class="euiCallOut euiCallOut--danger euiCallOut--small"
+    data-test-subj="query-assist-guard-callout"
+  >
+    <div
+      class="euiCallOutHeader"
+    >
+      <svg
+        aria-hidden="true"
+        class="euiIcon euiIcon--medium euiIcon--inherit euiIcon-isLoading euiCallOutHeader__icon"
+        focusable="false"
+        height="16"
+        role="img"
+        viewBox="0 0 16 16"
+        width="16"
+        xmlns="http://www.w3.org/2000/svg"
+      >
+        <path
+          d="M5.277 10.088c.02.014.04.03.057.047.582.55 1.134.812 1.666.812.586 0 1.84-.293 3.713-.88L9 6.212V2H7v4.212l-1.723 3.876Zm-.438.987L3.539 14h8.922l-1.32-2.969C9.096 11.677 7.733 12 7 12c-.74 0-1.463-.315-2.161-.925ZM6 2H5V1h6v1h-1v4l3.375 7.594A1 1 0 0 1 12.461 15H3.54a1 1 0 0 1-.914-1.406L6 6V2Z"
+        />
+      </svg>
+      <span
+        class="euiCallOutHeader__title"
+      >
+        <span>
+          I am unable to respond to this query. Try another question.
+        </span>
+      </span>
+    </div>
+  </div>
+</div>
+`;
+
+exports[`CallOuts spec should display query_generated call out 1`] = `
+<div>
+  <div
+    class="euiCallOut euiCallOut--success euiCallOut--small"
+    data-test-subj="query-assist-query-generated-callout"
+  >
+    <div
+      class="euiCallOutHeader"
+    >
+      <svg
+        aria-hidden="true"
+        class="euiIcon euiIcon--medium euiIcon--inherit euiIcon-isLoading euiCallOutHeader__icon"
+        focusable="false"
+        height="16"
+        role="img"
+        viewBox="0 0 16 16"
+        width="16"
+        xmlns="http://www.w3.org/2000/svg"
+      >
+        <path
+          d="M5.277 10.088c.02.014.04.03.057.047.582.55 1.134.812 1.666.812.586 0 1.84-.293 3.713-.88L9 6.212V2H7v4.212l-1.723 3.876Zm-.438.987L3.539 14h8.922l-1.32-2.969C9.096 11.677 7.733 12 7 12c-.74 0-1.463-.315-2.161-.925ZM6 2H5V1h6v1h-1v4l3.375 7.594A1 1 0 0 1 12.461 15H3.54a1 1 0 0 1-.914-1.406L6 6V2Z"
+        />
+      </svg>
+      <span
+        class="euiCallOutHeader__title"
+      >
+        <span>
+          test lang query generated. If there are any issues with the response, try adding more context to the question or a new question to submit.
+        </span>
+      </span>
+    </div>
+    <button
+      aria-label="dismissible_icon"
+      class="euiButtonIcon euiButtonIcon--primary euiButtonIcon--empty euiButtonIcon--xSmall euiCallOut__closeIcon"
+      data-test-subj="closeCallOutButton"
+      type="button"
+    >
+      <svg
+        aria-hidden="true"
+        class="euiIcon euiIcon--medium euiIcon--inherit euiIcon-isLoading euiButtonIcon__icon"
+        focusable="false"
+        height="16"
+        role="img"
+        viewBox="0 0 16 16"
+        width="16"
+        xmlns="http://www.w3.org/2000/svg"
+      >
+        <path
+          d="M5.277 10.088c.02.014.04.03.057.047.582.55 1.134.812 1.666.812.586 0 1.84-.293 3.713-.88L9 6.212V2H7v4.212l-1.723 3.876Zm-.438.987L3.539 14h8.922l-1.32-2.969C9.096 11.677 7.733 12 7 12c-.74 0-1.463-.315-2.161-.925ZM6 2H5V1h6v1h-1v4l3.375 7.594A1 1 0 0 1 12.461 15H3.54a1 1 0 0 1-.914-1.406L6 6V2Z"
+        />
+      </svg>
+    </button>
+  </div>
+</div>
+`;

public/query_assist/components/call_outs.test.tsx

@@ -0,0 +1,46 @@
+import { render } from '@testing-library/react';
+import React, { ComponentProps } from 'react';
+import { QueryAssistCallOut } from './call_outs';
+
+type Props = ComponentProps<typeof QueryAssistCallOut>;
+
+const renderCallOut = (overrideProps: Partial<Props> = {}) => {
+  const props: Props = Object.assign<Props, Partial<Props>>(
+    {
+      type: 'empty_query',
+      language: 'test lang',
+      onDismiss: jest.fn(),
+    },
+    overrideProps
+  );
+  const component = render(<QueryAssistCallOut {...props} />);
+  return { component, props: props as jest.MockedObjectDeep<Props> };
+};
+
+describe('CallOuts spec', () => {
+  it('should display nothing if type is invalid', () => {
+    // @ts-expect-error testing invalid type
+    const { component } = renderCallOut({ type: '' });
+    expect(component.container).toBeEmptyDOMElement();
+  });
+
+  it('should display empty_query call out', () => {
+    const { component } = renderCallOut({ type: 'empty_query' });
+    expect(component.container).toMatchSnapshot();
+  });
+
+  it('should display empty_index call out', () => {
+    const { component } = renderCallOut({ type: 'empty_index' });
+    expect(component.container).toMatchSnapshot();
+  });
+
+  it('should display invalid_query call out', () => {
+    const { component } = renderCallOut({ type: 'invalid_query' });
+    expect(component.container).toMatchSnapshot();
+  });
+
+  it('should display query_generated call out', () => {
+    const { component } = renderCallOut({ type: 'query_generated' });
+    expect(component.container).toMatchSnapshot();
+  });
+});