diff --git a/common/config.ts b/common/config.ts index b6be3f7..b9ea475 100644 --- a/common/config.ts +++ b/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; diff --git a/common/constants.ts b/common/constants.ts index aa78e9a..67fdabc 100644 --- a/common/constants.ts +++ b/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' }; diff --git a/common/query_assist/index.ts b/common/query_assist/index.ts index e36438e..375227d 100644 --- a/common/query_assist/index.ts +++ b/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'; diff --git a/common/query_assist/types.ts b/common/query_assist/types.ts new file mode 100644 index 0000000..1193d99 --- /dev/null +++ b/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; +} diff --git a/public/index.ts b/public/index.ts index d2bb583..0dc86c5 100644 --- a/public/index.ts +++ b/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'; diff --git a/public/plugin.tsx b/public/plugin.tsx index e027c21..5f45387 100644 --- a/public/plugin.tsx +++ b/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; + export class QueryEnhancementsPlugin implements Plugin { private readonly storage: IStorageWrapper; + private readonly config: PublicConfig; - constructor() { + constructor(initializerContext: PluginInitializerContext) { + this.config = initializerContext.config.get(); 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), }, }); diff --git a/public/query_assist/components/query_assist_banner.tsx b/public/query_assist/components/query_assist_banner.tsx index f1cab31..f042e8f 100644 --- a/public/query_assist/components/query_assist_banner.tsx +++ b/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 = (props) => { const storage = getStorage(); const [showCallOut, _setShowCallOut] = useState(true); const setShowCallOut: typeof _setShowCallOut = (show) => { @@ -50,7 +53,7 @@ export const QueryAssistBanner: React.FC = () => { diff --git a/public/query_assist/hooks/use_generate.ts b/public/query_assist/hooks/use_generate.ts index 5365a7c..a5e0db6 100644 --- a/public/query_assist/hooks/use_generate.ts +++ b/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( - '/api/ql/query_assist/generate', - { - body: JSON.stringify(params), - signal: abortControllerRef.current?.signal, - } - ); + const response = await services.http.post(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) }; diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx index 0fb0a49..80c9293 100644 --- a/public/query_assist/utils/create_extension.tsx +++ b/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 ( - + conf.language)} + /> ); }, diff --git a/public/query_assist/utils/errors.ts b/public/query_assist/utils/errors.ts index 630a375..4f1dbb9 100644 --- a/public/query_assist/utils/errors.ts +++ b/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) { diff --git a/server/index.ts b/server/index.ts index 45aade2..b8437d8 100644 --- a/server/index.ts +++ b/server/index.ts @@ -8,7 +8,9 @@ import { QueryEnhancementsPlugin } from './plugin'; import { configSchema, ConfigSchema } from '../common/config'; export const config: PluginConfigDescriptor = { - exposeToBrowser: {}, + exposeToBrowser: { + queryAssist: true, + }, schema: configSchema, }; diff --git a/server/plugin.ts b/server/plugin.ts index a97ddfa..e372dcf 100644 --- a/server/plugin.ts +++ b/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 { private readonly logger: Logger; private readonly config$: Observable; - 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() + .pipe(first()) + .toPromise(), dataSourceEnabled: !!dataSource, })); diff --git a/server/routes/query_assist/createResponse.ts b/server/routes/query_assist/createResponse.ts new file mode 100644 index 0000000..c72305b --- /dev/null +++ b/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; + const query = Object.values(result).at(0); + if (typeof query !== 'string') throw new Error('Generated query not found.'); + return { query }; + } +}; diff --git a/server/routes/query_assist/index.ts b/server/routes/query_assist/index.ts index 3a62d86..5f6c2c9 100644 --- a/server/routes/query_assist/index.ts +++ b/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 = { - PPL: 'os_query_assist_ppl', -} as const; diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts index c60035c..51c07d7 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -1,17 +1,14 @@ -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()), @@ -19,6 +16,7 @@ export function registerQueryAssistRoutes(router: IRouter) { }, }, 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) @@ -26,10 +24,10 @@ export function registerQueryAssistRoutes(router: IRouter) { 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)) { diff --git a/server/types.ts b/server/types.ts index 97a9379..ab85fe2 100644 --- a/server/types.ts +++ b/server/types.ts @@ -6,6 +6,7 @@ import { DataPluginSetup } from 'src/plugins/data/server/plugin'; import { Logger } from '../../../src/core/server'; import { DataSourcePluginStart } from '../../../src/plugins/data_source/server'; +import { ConfigSchema } from '../common/config'; // eslint-disable-next-line @typescript-eslint/no-empty-interface export interface QueryEnhancementsPluginSetup {} @@ -54,6 +55,7 @@ declare module '../../../src/core/server' { interface RequestHandlerContext { query_assist: { logger: Logger; + configPromise: Promise; dataSourceEnabled: boolean; }; }