extract supported languages to config

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

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

@@ -1,18 +1 @@
-import { TimeRange } from '../../../../src/plugins/data/common';
-
-export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrails triggered' };
-
-export const SUPPORTED_LANGUAGES = ['PPL'] as const;
-
-export interface QueryAssistResponse {
-  query: string;
-  timeRange?: TimeRange;
-}
-
-export interface QueryAssistParameters {
-  question: string;
-  index: string;
-  language: string;
-  // for MDS
-  dataSourceId?: string;
-}
+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;
+}

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,9 @@
  */
 
 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';
@@ -16,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);
   }
 
@@ -87,7 +92,7 @@ export class QueryEnhancementsPlugin
 
     data.__enhance({
       ui: {
-        queryEditorExtension: createQueryAssistExtension(core.http),
+        queryEditorExtension: createQueryAssistExtension(core.http, this.config),
       },
     });
 

public/query_assist/components/query_assist_banner.tsx

@@ -9,13 +9,16 @@ import {
 } from '@elastic/eui';
 import { FormattedMessage } from '@osd/i18n/react';
 import React, { useState } from 'react';
-import { SUPPORTED_LANGUAGES } from '../../../common/query_assist';
 import assistantMark from '../../assets/query_assist_mark.svg';
 import { getStorage } from '../../services';
 
 const BANNER_STORAGE_KEY = 'queryAssist:banner:show';
 
-export const QueryAssistBanner: React.FC = () => {
+interface QueryAssistBannerProps {
+  languages: string[];
+}
+
+export const QueryAssistBanner: React.FC<QueryAssistBannerProps> = (props) => {
   const storage = getStorage();
   const [showCallOut, _setShowCallOut] = useState(true);
   const setShowCallOut: typeof _setShowCallOut = (show) => {
@@ -50,7 +53,7 @@ export const QueryAssistBanner: React.FC = () => {
                 <FormattedMessage
                   id="queryAssist.banner.title.suffix"
                   defaultMessage="Natural Language Query Generation for {languages}"
-                  values={{ languages: SUPPORTED_LANGUAGES.join(', ') }}
+                  values={{ languages: props.languages.join(', ') }}
                 />
               </EuiLink>
             </EuiTextColor>

public/query_assist/hooks/use_generate.ts

@@ -1,6 +1,7 @@
 import { useEffect, useRef, useState } from 'react';
 import { IDataPluginServices } from '../../../../../src/plugins/data/public';
 import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
+import { API } from '../../../common';
 import { QueryAssistParameters, QueryAssistResponse } from '../../../common/query_assist';
 import { formatError } from '../utils';
 
@@ -28,13 +29,10 @@ export const useGenerateQuery = () => {
     abortControllerRef.current = new AbortController();
     setLoading(true);
     try {
-      const response = await services.http.post<QueryAssistResponse>(
-        '/api/ql/query_assist/generate',
-        {
-          body: JSON.stringify(params),
-          signal: abortControllerRef.current?.signal,
-        }
-      );
+      const response = await services.http.post<QueryAssistResponse>(API.QUERY_ASSIST.GENERATE, {
+        body: JSON.stringify(params),
+        signal: abortControllerRef.current?.signal,
+      });
       if (mounted.current) return { response };
     } catch (error) {
       if (mounted.current) return { error: formatError(error) };

public/query_assist/utils/create_extension.tsx

@@ -3,6 +3,8 @@ import React, { useEffect, useState } from 'react';
 import { getMdsDataSourceId } from '.';
 import { QueryEditorExtensionConfig } from '../../../../../src/plugins/data/public/ui/query_editor';
 import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension';
+import { API } from '../../../common';
+import { PublicConfig } from '../../plugin';
 import { getData } from '../../services';
 import { QueryAssistBar } from '../components';
 import { QueryAssistBanner } from '../components/query_assist_banner';
@@ -29,7 +31,7 @@ const getAvailableLanguages = async (
   if (cached !== undefined) return cached;
 
   const languages = await http
-    .get<{ configuredLanguages: string[] }>('/api/ql/query_assist/configured_languages', {
+    .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
       query: { dataSourceId },
     })
     .then((response) => response.configuredLanguages)
@@ -38,7 +40,10 @@ const getAvailableLanguages = async (
   return languages;
 };
 
-export const createQueryAssistExtension = (http: HttpSetup): QueryEditorExtensionConfig => {
+export const createQueryAssistExtension = (
+  http: HttpSetup,
+  config: PublicConfig
+): QueryEditorExtensionConfig => {
   return {
     id: 'query-assist',
     order: 1000,
@@ -62,7 +67,9 @@ export const createQueryAssistExtension = (http: HttpSetup): QueryEditorExtensio
       // advertise query assist if user is not on a supported language.
       return (
         <QueryAssistWrapper dependencies={dependencies} http={http} invert>
-          <QueryAssistBanner />
+          <QueryAssistBanner
+            languages={config.queryAssist.supportedLanguages.map((conf) => conf.language)}
+          />
         </QueryAssistWrapper>
       );
     },

public/query_assist/utils/errors.ts

@@ -1,5 +1,5 @@
 import { ResponseError } from '@opensearch-project/opensearch/lib/errors';
-import { ERROR_DETAILS } from '../../../common/query_assist';
+import { ERROR_DETAILS } from '../../../common';
 
 export class ProhibitedQueryError extends Error {
   constructor(message?: string) {

server/index.ts

@@ -8,7 +8,9 @@ import { QueryEnhancementsPlugin } from './plugin';
 import { configSchema, ConfigSchema } from '../common/config';
 
 export const config: PluginConfigDescriptor<ConfigSchema> = {
-  exposeToBrowser: {},
+  exposeToBrowser: {
+    queryAssist: true,
+  },
   schema: configSchema,
 };
 

server/plugin.ts

@@ -4,6 +4,7 @@
  */
 
 import { Observable } from 'rxjs';
+import { first } from 'rxjs/operators';
 import {
   CoreSetup,
   CoreStart,
@@ -13,6 +14,7 @@ import {
   SharedGlobalConfig,
 } from '../../../src/core/server';
 import { SEARCH_STRATEGY } from '../common';
+import { ConfigSchema } from '../common/config';
 import { defineRoutes } from './routes';
 import { pplSearchStrategyProvider, sqlSearchStrategyProvider } from './search';
 import {
@@ -26,7 +28,7 @@ export class QueryEnhancementsPlugin
   implements Plugin<QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart> {
   private readonly logger: Logger;
   private readonly config$: Observable<SharedGlobalConfig>;
-  constructor(initializerContext: PluginInitializerContext) {
+  constructor(private initializerContext: PluginInitializerContext) {
     this.logger = initializerContext.logger.get();
     this.config$ = initializerContext.config.legacy.globalConfig$;
   }
@@ -47,6 +49,10 @@ export class QueryEnhancementsPlugin
 
     core.http.registerRouteHandlerContext('query_assist', () => ({
       logger: this.logger,
+      configPromise: this.initializerContext.config
+        .create<ConfigSchema>()
+        .pipe(first())
+        .toPromise(),
       dataSourceEnabled: !!dataSource,
     }));
 

server/routes/query_assist/createResponse.ts

@@ -0,0 +1,23 @@
+import { QueryAssistResponse } from '../../../common/query_assist';
+import { AgentResponse } from './agents';
+import { createPPLResponseBody } from './ppl/create_response';
+
+export const createResponseBody = (
+  language: string,
+  agentResponse: AgentResponse
+): QueryAssistResponse => {
+  switch (language) {
+    case 'PPL':
+      return createPPLResponseBody(agentResponse);
+
+    default:
+      if (!agentResponse.body.inference_results[0].output[0].result)
+        throw new Error('Generated query not found.');
+      const result = JSON.parse(
+        agentResponse.body.inference_results[0].output[0].result!
+      ) as Record<string, string>;
+      const query = Object.values(result).at(0);
+      if (typeof query !== 'string') throw new Error('Generated query not found.');
+      return { query };
+  }
+};

server/routes/query_assist/index.ts

@@ -1,7 +1 @@
-import { SUPPORTED_LANGUAGES } from '../../../common/query_assist';
-
 export { registerQueryAssistRoutes } from './routes';
-
-export const AGENT_CONFIG_NAME_MAP: Record<typeof SUPPORTED_LANGUAGES[number], string> = {
-  PPL: 'os_query_assist_ppl',
-} as const;

server/routes/query_assist/routes.ts

@@ -1,35 +1,33 @@
-import { schema, Type } from '@osd/config-schema';
+import { schema } from '@osd/config-schema';
 import { IRouter } from 'opensearch-dashboards/server';
 import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors';
-import { ERROR_DETAILS, SUPPORTED_LANGUAGES } from '../../../common/query_assist';
+import { API, ERROR_DETAILS } from '../../../common';
 import { getAgentIdByConfig, requestAgentByConfig } from './agents';
-import { AGENT_CONFIG_NAME_MAP } from './index';
-import { createPPLResponseBody } from './ppl/create_response';
+import { createResponseBody } from './createResponse';
 
 export function registerQueryAssistRoutes(router: IRouter) {
-  const languageSchema = schema.oneOf(SUPPORTED_LANGUAGES.map(schema.literal) as [Type<'PPL'>]);
-
   router.get(
     {
-      path: '/api/ql/query_assist/configured_languages',
+      path: API.QUERY_ASSIST.LANGUAGES,
       validate: {
         query: schema.object({
           dataSourceId: schema.maybe(schema.string()),
         }),
       },
     },
     async (context, request, response) => {
+      const config = await context.query_assist.configPromise;
       const client =
         context.query_assist.dataSourceEnabled && request.query.dataSourceId
           ? await context.dataSource.opensearch.getClient(request.query.dataSourceId)
           : context.core.opensearch.client.asCurrentUser;
       const configuredLanguages: string[] = [];
       try {
         await Promise.allSettled(
-          SUPPORTED_LANGUAGES.map((language) =>
-            getAgentIdByConfig(client, AGENT_CONFIG_NAME_MAP[language]).then(() =>
-              // if the call does not throw any error, then the agent is properly configured
-              configuredLanguages.push(language)
+          config.queryAssist.supportedLanguages.map((languageConfig) =>
+            // if the call does not throw any error, then the agent is properly configured
+            getAgentIdByConfig(client, languageConfig.agentConfig).then(() =>
+              configuredLanguages.push(languageConfig.language)
             )
           )
         );
@@ -42,21 +40,26 @@ export function registerQueryAssistRoutes(router: IRouter) {
 
   router.post(
     {
-      path: '/api/ql/query_assist/generate',
+      path: API.QUERY_ASSIST.GENERATE,
       validate: {
         body: schema.object({
           index: schema.string(),
           question: schema.string(),
-          language: languageSchema,
+          language: schema.string(),
           dataSourceId: schema.maybe(schema.string()),
         }),
       },
     },
     async (context, request, response) => {
+      const config = await context.query_assist.configPromise;
+      const languageConfig = config.queryAssist.supportedLanguages.find(
+        (c) => c.language === request.body.language
+      );
+      if (!languageConfig) return response.badRequest({ body: 'Unsupported language' });
       try {
         const agentResponse = await requestAgentByConfig({
           context,
-          configName: AGENT_CONFIG_NAME_MAP[request.body.language],
+          configName: languageConfig.agentConfig,
           body: {
             parameters: {
               index: request.body.index,
@@ -65,7 +68,7 @@ export function registerQueryAssistRoutes(router: IRouter) {
           },
           dataSourceId: request.body.dataSourceId,
         });
-        const responseBody = createPPLResponseBody(agentResponse);
+        const responseBody = createResponseBody(languageConfig.language, agentResponse);
         return response.ok({ body: responseBody });
       } catch (error) {
         if (isResponseError(error)) {