Add query assist to query enhancements plugin

.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);
   }
 
@@ -84,6 +90,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();
+  });
+});