From b6d54a53c9c44d5034634f2f8b1b77895dc7d3ef Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 4 Jun 2024 15:11:58 -0700 Subject: [PATCH 01/10] [Discover-next] add query assist to query enhancements plugin (#6895) it adds query assist specific logic in query enhancements plugin to show a UI above the PPL search bar if user has configured PPL agent. Issues Resolved: #6820 * add query assist to query enhancements Signed-off-by: Joshua Li * align language to uppercase Signed-off-by: Joshua Li * pick PR 6167 Signed-off-by: Joshua Li * use useState hooks for query assist There is a bug in data explorer `AppContainer` where memorized `DiscoverCanvas` gets unmounted after `setQuery`. PR 6167 works around it by memorizing `AppContainer`. As query assist is no longer being unmounted, we can use proper hooks to persist state now. Signed-off-by: Joshua Li * Revert "pick PR 6167" This reverts commit acb0d41937e30bd76c666a225407743243692d11. Wait for official 6167 to merge to avoid conflict Signed-off-by: Joshua Li * address comments for PR 6894 Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li (cherry picked from commit 016e0f2f73efd8bb0649151908c67dd7ac09d174) --- common/query_assist/index.ts | 14 +++ opensearch_dashboards.json | 2 +- public/assets/query_assist_logo.svg | 18 ++++ public/plugin.tsx | 2 + public/query_assist/components/call_outs.tsx | 79 +++++++++++++++ public/query_assist/components/index.ts | 1 + .../components/query_assist_bar.tsx | 95 +++++++++++++++++++ .../components/query_assist_input.tsx | 75 +++++++++++++++ .../query_assist/components/submit_button.tsx | 19 ++++ public/query_assist/hooks/index.ts | 1 + public/query_assist/hooks/use_generate.ts | 36 +++++++ public/query_assist/index.ts | 1 + .../query_assist/utils/create_extension.tsx | 24 +++++ public/query_assist/utils/errors.ts | 25 +++++ .../query_assist/utils/get_persisted_log.ts | 19 ++++ public/query_assist/utils/index.ts | 3 + server/routes/index.ts | 4 +- server/routes/query_assist/agents.ts | 62 ++++++++++++ server/routes/query_assist/index.ts | 5 + .../query_assist/ppl/create_response.ts | 17 ++++ server/routes/query_assist/routes.ts | 77 +++++++++++++++ 21 files changed, 577 insertions(+), 2 deletions(-) create mode 100644 common/query_assist/index.ts create mode 100644 public/assets/query_assist_logo.svg create mode 100644 public/query_assist/components/call_outs.tsx create mode 100644 public/query_assist/components/index.ts create mode 100644 public/query_assist/components/query_assist_bar.tsx create mode 100644 public/query_assist/components/query_assist_input.tsx create mode 100644 public/query_assist/components/submit_button.tsx create mode 100644 public/query_assist/hooks/index.ts create mode 100644 public/query_assist/hooks/use_generate.ts create mode 100644 public/query_assist/index.ts create mode 100644 public/query_assist/utils/create_extension.tsx create mode 100644 public/query_assist/utils/errors.ts create mode 100644 public/query_assist/utils/get_persisted_log.ts create mode 100644 public/query_assist/utils/index.ts create mode 100644 server/routes/query_assist/agents.ts create mode 100644 server/routes/query_assist/index.ts create mode 100644 server/routes/query_assist/ppl/create_response.ts create mode 100644 server/routes/query_assist/routes.ts diff --git a/common/query_assist/index.ts b/common/query_assist/index.ts new file mode 100644 index 0000000..d4173b1 --- /dev/null +++ b/common/query_assist/index.ts @@ -0,0 +1,14 @@ +import { TimeRange } from '../../../../src/plugins/data/common'; + +export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrails triggered' }; + +export interface QueryAssistResponse { + query: string; + timeRange?: TimeRange; +} + +export interface QueryAssistParameters { + question: string; + index: string; + language: string; +} diff --git a/opensearch_dashboards.json b/opensearch_dashboards.json index 710f289..b7b3e9c 100644 --- a/opensearch_dashboards.json +++ b/opensearch_dashboards.json @@ -6,5 +6,5 @@ "ui": true, "requiredPlugins": ["data"], "optionalPlugins": ["dataSource", "dataSourceManagement"], - "requiredBundles": ["opensearchDashboardsUtils"] + "requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"] } diff --git a/public/assets/query_assist_logo.svg b/public/assets/query_assist_logo.svg new file mode 100644 index 0000000..b744e8c --- /dev/null +++ b/public/assets/query_assist_logo.svg @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + + diff --git a/public/plugin.tsx b/public/plugin.tsx index 353018a..f09dab1 100644 --- a/public/plugin.tsx +++ b/public/plugin.tsx @@ -6,6 +6,7 @@ import moment from 'moment'; import { CoreSetup, CoreStart, Plugin } from '../../../src/core/public'; import { IStorageWrapper, Storage } from '../../../src/plugins/opensearch_dashboards_utils/public'; +import { createQueryAssistExtension } from './query_assist'; import { PPLSearchInterceptor, SQLSearchInterceptor } from './search'; import { setData, setStorage } from './services'; import { @@ -54,6 +55,7 @@ export class QueryEnhancementsPlugin initialTo: moment().add(2, 'days').toISOString(), }, showFilterBar: false, + extensions: [createQueryAssistExtension(core.http)], }, fields: { filterable: false, diff --git a/public/query_assist/components/call_outs.tsx b/public/query_assist/components/call_outs.tsx new file mode 100644 index 0000000..f98569f --- /dev/null +++ b/public/query_assist/components/call_outs.tsx @@ -0,0 +1,79 @@ +import { EuiCallOut, EuiCallOutProps } from '@elastic/eui'; +import React from 'react'; + +type CalloutDismiss = Required>; +interface QueryAssistCallOutProps extends CalloutDismiss { + type: QueryAssistCallOutType; +} + +export type QueryAssistCallOutType = + | undefined + | 'invalid_query' + | 'prohibited_query' + | 'empty_query' + | 'empty_index' + | 'query_generated'; + +const EmptyIndexCallOut: React.FC = (props) => ( + +); + +const ProhibitedQueryCallOut: React.FC = (props) => ( + +); + +const EmptyQueryCallOut: React.FC = (props) => ( + +); + +const PPLGeneratedCallOut: React.FC = (props) => ( + +); + +export const QueryAssistCallOut: React.FC = (props) => { + switch (props.type) { + case 'empty_query': + return ; + case 'empty_index': + return ; + case 'invalid_query': + return ; + case 'query_generated': + return ; + default: + break; + } + return null; +}; diff --git a/public/query_assist/components/index.ts b/public/query_assist/components/index.ts new file mode 100644 index 0000000..d6276e4 --- /dev/null +++ b/public/query_assist/components/index.ts @@ -0,0 +1 @@ +export { QueryAssistBar } from './query_assist_bar'; diff --git a/public/query_assist/components/query_assist_bar.tsx b/public/query_assist/components/query_assist_bar.tsx new file mode 100644 index 0000000..c80fe5d --- /dev/null +++ b/public/query_assist/components/query_assist_bar.tsx @@ -0,0 +1,95 @@ +import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui'; +import React, { SyntheticEvent, useEffect, 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'; +import { getStorage } from '../../services'; +import { useGenerateQuery } from '../hooks'; +import { getPersistedLog, ProhibitedQueryError } from '../utils'; +import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs'; +import { QueryAssistInput } from './query_assist_input'; +import { QueryAssistSubmitButton } from './submit_button'; + +interface QueryAssistInputProps { + dependencies: SearchBarExtensionDependencies; +} + +export const QueryAssistBar: React.FC = (props) => { + const { services } = useOpenSearchDashboards(); + const inputRef = useRef(null); + const storage = getStorage(); + const persistedLog: PersistedLog = useMemo( + () => getPersistedLog(services.uiSettings, storage, 'query-assist'), + [services.uiSettings, storage] + ); + const { generateQuery, loading } = useGenerateQuery(); + const [callOutType, setCallOutType] = useState(); + const dismissCallout = () => setCallOutType(undefined); + const mounted = useRef(false); + const selectedIndex = props.dependencies.indexPatterns?.at(0)?.title; + const previousQuestionRef = useRef(); + + useEffect(() => { + mounted.current = true; + return () => { + mounted.current = false; + }; + }, []); + + const onSubmit = async (e: SyntheticEvent) => { + e.preventDefault(); + if (!inputRef.current?.value) { + setCallOutType('empty_query'); + return; + } + if (!selectedIndex) { + setCallOutType('empty_index'); + return; + } + dismissCallout(); + previousQuestionRef.current = inputRef.current.value; + persistedLog.add(inputRef.current.value); + const params = { + question: inputRef.current.value, + index: selectedIndex, + language: 'PPL', + }; + const { response, error } = await generateQuery(params); + if (!mounted.current) return; + if (error) { + if (error instanceof ProhibitedQueryError) { + setCallOutType('invalid_query'); + } else { + services.notifications.toasts.addError(error, { title: 'Failed to generate results' }); + } + } else if (response) { + services.data.query.queryString.setQuery({ + query: response.query, + language: params.language, + }); + if (response.timeRange) services.data.query.timefilter.timefilter.setTime(response.timeRange); + setCallOutType('query_generated'); + } + }; + + return ( + + + + + + + + + + + + + + ); +}; diff --git a/public/query_assist/components/query_assist_input.tsx b/public/query_assist/components/query_assist_input.tsx new file mode 100644 index 0000000..735ff55 --- /dev/null +++ b/public/query_assist/components/query_assist_input.tsx @@ -0,0 +1,75 @@ +import { EuiFieldText, EuiIcon, EuiOutsideClickDetector, EuiPortal } from '@elastic/eui'; +import React, { useMemo, useState } from 'react'; +import { PersistedLog, QuerySuggestionTypes } from '../../../../../src/plugins/data/public'; +import assistantLogo from '../../assets/query_assist_logo.svg'; +import { getData } from '../../services'; + +interface QueryAssistInputProps { + inputRef: React.RefObject; + persistedLog: PersistedLog; + initialValue?: string; + selectedIndex?: string; + previousQuestion?: string; +} + +export const QueryAssistInput: React.FC = (props) => { + const { + ui: { SuggestionsComponent }, + } = getData(); + const [isSuggestionsVisible, setIsSuggestionsVisible] = useState(false); + const [suggestionIndex, setSuggestionIndex] = useState(null); + const [value, setValue] = useState(props.initialValue ?? ''); + + const recentSearchSuggestions = useMemo(() => { + if (!props.persistedLog) return []; + return props.persistedLog + .get() + .filter((recentSearch) => recentSearch.includes(value)) + .map((recentSearch) => ({ + type: QuerySuggestionTypes.RecentSearch, + text: recentSearch, + start: 0, + end: value.length, + })); + }, [props.persistedLog, value]); + + return ( + setIsSuggestionsVisible(false)}> +
+ setIsSuggestionsVisible(true)} + onChange={(e) => setValue(e.target.value)} + onKeyDown={() => setIsSuggestionsVisible(true)} + placeholder={ + props.previousQuestion || + (props.selectedIndex + ? `Ask a natural language question about ${props.selectedIndex} to generate a query` + : 'Select an index pattern to ask a question') + } + prepend={} + fullWidth + /> + + { + if (!props.inputRef.current) return; + setValue(suggestion.text); + setIsSuggestionsVisible(false); + setSuggestionIndex(null); + props.inputRef.current.focus(); + }} + onMouseEnter={(i) => setSuggestionIndex(i)} + loadMore={() => {}} + queryBarRect={props.inputRef.current?.getBoundingClientRect()} + size="s" + /> + +
+ + ); +}; diff --git a/public/query_assist/components/submit_button.tsx b/public/query_assist/components/submit_button.tsx new file mode 100644 index 0000000..3b7a3c8 --- /dev/null +++ b/public/query_assist/components/submit_button.tsx @@ -0,0 +1,19 @@ +import React from 'react'; +import { EuiButtonIcon } from '@elastic/eui'; + +interface SubmitButtonProps { + isDisabled: boolean; +} + +export const QueryAssistSubmitButton: React.FC = (props) => { + return ( + + ); +}; diff --git a/public/query_assist/hooks/index.ts b/public/query_assist/hooks/index.ts new file mode 100644 index 0000000..5922606 --- /dev/null +++ b/public/query_assist/hooks/index.ts @@ -0,0 +1 @@ +export * from './use_generate'; diff --git a/public/query_assist/hooks/use_generate.ts b/public/query_assist/hooks/use_generate.ts new file mode 100644 index 0000000..0d23114 --- /dev/null +++ b/public/query_assist/hooks/use_generate.ts @@ -0,0 +1,36 @@ +import { useEffect, useRef, useState } from 'react'; +import { IDataPluginServices } from '../../../../../src/plugins/data/public'; +import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public'; +import { QueryAssistParameters, QueryAssistResponse } from '../../../common/query_assist'; +import { formatError } from '../utils'; + +export const useGenerateQuery = () => { + const [loading, setLoading] = useState(false); + const abortControllerRef = useRef(); + const { services } = useOpenSearchDashboards(); + + useEffect(() => () => abortControllerRef.current?.abort(), []); + + const generateQuery = async ( + params: QueryAssistParameters + ): Promise<{ response?: QueryAssistResponse; error?: Error }> => { + 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, + } + ); + return { response }; + } catch (error) { + return { error: formatError(error) }; + } finally { + setLoading(false); + } + }; + + return { generateQuery, loading, abortControllerRef }; +}; diff --git a/public/query_assist/index.ts b/public/query_assist/index.ts new file mode 100644 index 0000000..0c059cf --- /dev/null +++ b/public/query_assist/index.ts @@ -0,0 +1 @@ +export { createQueryAssistExtension } from './utils'; diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx new file mode 100644 index 0000000..f2f046d --- /dev/null +++ b/public/query_assist/utils/create_extension.tsx @@ -0,0 +1,24 @@ +import React from 'react'; +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 => { + return { + id: 'query-assist-ppl', + order: 1000, + isEnabled: (() => { + let agentConfigured: boolean; + return async () => { + if (agentConfigured === undefined) { + agentConfigured = await http + .get<{ configured: boolean }>('/api/ql/query_assist/configured/PPL') + .then((response) => response.configured) + .catch(() => false); + } + return agentConfigured; + }; + })(), + getComponent: (dependencies) => , + }; +}; diff --git a/public/query_assist/utils/errors.ts b/public/query_assist/utils/errors.ts new file mode 100644 index 0000000..630a375 --- /dev/null +++ b/public/query_assist/utils/errors.ts @@ -0,0 +1,25 @@ +import { ResponseError } from '@opensearch-project/opensearch/lib/errors'; +import { ERROR_DETAILS } from '../../../common/query_assist'; + +export class ProhibitedQueryError extends Error { + constructor(message?: string) { + super(message); + } +} + +export const formatError = (error: ResponseError | Error): Error => { + if ('body' in error) { + if (error.body.statusCode === 429) + return { + ...error.body, + message: 'Request is throttled. Try again later or contact your administrator', + } as Error; + if ( + error.body.statusCode === 400 && + error.body.message.includes(ERROR_DETAILS.GUARDRAILS_TRIGGERED) + ) + return new ProhibitedQueryError(error.body.message); + return error.body as Error; + } + return error; +}; diff --git a/public/query_assist/utils/get_persisted_log.ts b/public/query_assist/utils/get_persisted_log.ts new file mode 100644 index 0000000..185acdc --- /dev/null +++ b/public/query_assist/utils/get_persisted_log.ts @@ -0,0 +1,19 @@ +import { IUiSettingsClient } from 'opensearch-dashboards/public'; +import { UI_SETTINGS } from '../../../../../src/plugins/data/common'; +import { PersistedLog } from '../../../../../src/plugins/data/public'; +import { IStorageWrapper } from '../../../../../src/plugins/opensearch_dashboards_utils/public'; + +export function getPersistedLog( + uiSettings: IUiSettingsClient, + storage: IStorageWrapper, + language: string +) { + return new PersistedLog( + `typeahead:${language}`, + { + maxLength: uiSettings.get(UI_SETTINGS.HISTORY_LIMIT), + filterDuplicates: true, + }, + storage + ); +} diff --git a/public/query_assist/utils/index.ts b/public/query_assist/utils/index.ts new file mode 100644 index 0000000..5384401 --- /dev/null +++ b/public/query_assist/utils/index.ts @@ -0,0 +1,3 @@ +export * from './create_extension'; +export * from './errors'; +export * from './get_persisted_log'; diff --git a/server/routes/index.ts b/server/routes/index.ts index d9c0f65..49366e2 100644 --- a/server/routes/index.ts +++ b/server/routes/index.ts @@ -15,7 +15,8 @@ import { IOpenSearchDashboardsSearchRequest, } from '../../../../src/plugins/data/common'; import { ISearchStrategy } from '../../../../src/plugins/data/server'; -import { SEARCH_STRATEGY, API } from '../../common'; +import { API, SEARCH_STRATEGY } from '../../common'; +import { registerQueryAssistRoutes } from './query_assist'; function defineRoute( logger: Logger, @@ -69,4 +70,5 @@ export function defineRoutes( ) { defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.PPL); defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.SQL); + registerQueryAssistRoutes(logger, router); } diff --git a/server/routes/query_assist/agents.ts b/server/routes/query_assist/agents.ts new file mode 100644 index 0000000..676ea01 --- /dev/null +++ b/server/routes/query_assist/agents.ts @@ -0,0 +1,62 @@ +import { ApiResponse } from '@opensearch-project/opensearch'; +import { RequestBody, TransportRequestPromise } from '@opensearch-project/opensearch/lib/Transport'; +import { RequestHandlerContext } from 'src/core/server'; + +const ML_COMMONS_API_PREFIX = '/_plugins/_ml'; +const AGENT_REQUEST_OPTIONS = { + /** + * It is time-consuming for LLM to generate final answer + * Give it a large timeout window + */ + requestTimeout: 5 * 60 * 1000, + /** + * Do not retry + */ + maxRetries: 0, +}; + +export type AgentResponse = ApiResponse<{ + inference_results: Array<{ + output: Array<{ name: string; result?: string }>; + }>; +}>; + +type OpenSearchClient = RequestHandlerContext['core']['opensearch']['client']['asCurrentUser']; + +export const getAgentIdByConfig = async ( + client: OpenSearchClient, + configName: string +): Promise => { + try { + const response = (await client.transport.request({ + method: 'GET', + path: `${ML_COMMONS_API_PREFIX}/config/${configName}`, + })) as ApiResponse<{ type: string; configuration: { agent_id?: string } }>; + + if (!response || response.body.configuration.agent_id === undefined) { + throw new Error('cannot find any agent by configuration: ' + configName); + } + return response.body.configuration.agent_id; + } catch (error) { + const errorMessage = JSON.stringify(error.meta?.body) || error; + throw new Error(`Get agent '${configName}' failed, reason: ` + errorMessage); + } +}; + +export const requestAgentByConfig = async (options: { + context: RequestHandlerContext; + configName: string; + body: RequestBody; +}): Promise => { + const { context, configName, body } = options; + const client = context.core.opensearch.client.asCurrentUser; + const agentId = await getAgentIdByConfig(client, configName); + return client.transport.request( + { + method: 'POST', + path: `${ML_COMMONS_API_PREFIX}/agents/${agentId}/_execute`, + body, + }, + AGENT_REQUEST_OPTIONS + ) as TransportRequestPromise; +}; diff --git a/server/routes/query_assist/index.ts b/server/routes/query_assist/index.ts new file mode 100644 index 0000000..0e1cdc3 --- /dev/null +++ b/server/routes/query_assist/index.ts @@ -0,0 +1,5 @@ +export { registerQueryAssistRoutes } from './routes'; + +export const AGENT_CONFIG_NAME_MAP = { + PPL: 'os_query_assist_ppl', +} as const; diff --git a/server/routes/query_assist/ppl/create_response.ts b/server/routes/query_assist/ppl/create_response.ts new file mode 100644 index 0000000..ac651a0 --- /dev/null +++ b/server/routes/query_assist/ppl/create_response.ts @@ -0,0 +1,17 @@ +import { QueryAssistResponse } from '../../../../common/query_assist'; +import { AgentResponse } from '../agents'; + +export const createPPLResponseBody = (agentResponse: AgentResponse): QueryAssistResponse => { + 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 { + ppl: string; + }; + const ppl = result.ppl + .replace(/[\r\n]/g, ' ') + .trim() + .replace(/ISNOTNULL/g, 'isnotnull') // https://github.com/opensearch-project/sql/issues/2431 + .replace(/`/g, '') // https://github.com/opensearch-project/dashboards-observability/issues/509, https://github.com/opensearch-project/dashboards-observability/issues/557 + .replace(/\bSPAN\(/g, 'span('); // https://github.com/opensearch-project/dashboards-observability/issues/759 + return { query: ppl }; +}; diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts new file mode 100644 index 0000000..0b8c5c8 --- /dev/null +++ b/server/routes/query_assist/routes.ts @@ -0,0 +1,77 @@ +import { schema, Type } from '@osd/config-schema'; +import { IRouter, Logger } from 'opensearch-dashboards/server'; +import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors'; +import { ERROR_DETAILS } from '../../../common/query_assist'; +import { getAgentIdByConfig, requestAgentByConfig } from './agents'; +import { AGENT_CONFIG_NAME_MAP } from './index'; +import { createPPLResponseBody } from './ppl/create_response'; + +export function registerQueryAssistRoutes(logger: Logger, router: IRouter) { + const languageSchema = schema.oneOf( + Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'PPL'>] + ); + + router.get( + { + path: '/api/ql/query_assist/configured/{language}', + validate: { + params: schema.object({ + language: languageSchema, + }), + }, + }, + async (context, request, response) => { + const client = context.core.opensearch.client.asCurrentUser; + try { + // if the call does not throw any error, then the agent is properly configured + await getAgentIdByConfig(client, AGENT_CONFIG_NAME_MAP[request.params.language]); + return response.ok({ body: { configured: true } }); + } catch (error) { + return response.ok({ body: { configured: false, error: error.message } }); + } + } + ); + + router.post( + { + path: '/api/ql/query_assist/generate', + validate: { + body: schema.object({ + index: schema.string(), + question: schema.string(), + language: languageSchema, + }), + }, + }, + async (context, request, response) => { + try { + if (!(request.body.language in AGENT_CONFIG_NAME_MAP)) + throw new Error('Unsupported language.'); + const agentResponse = await requestAgentByConfig({ + context, + configName: AGENT_CONFIG_NAME_MAP[request.body.language], + body: { + parameters: { + index: request.body.index, + question: request.body.question, + }, + }, + }); + const responseBody = createPPLResponseBody(agentResponse); + return response.ok({ body: responseBody }); + } catch (error) { + if (isResponseError(error)) { + if (error.statusCode === 400 && error.body.includes(ERROR_DETAILS.GUARDRAILS_TRIGGERED)) + return response.badRequest({ body: ERROR_DETAILS.GUARDRAILS_TRIGGERED }); + return response.badRequest({ + body: + typeof error.meta.body === 'string' + ? error.meta.body + : JSON.stringify(error.meta.body), + }); + } + return response.custom({ statusCode: error.statusCode || 500, body: error.message }); + } + } + ); +} From b0053b9869462f4fd71f58b0c62c960d8dd64ce9 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Wed, 5 Jun 2024 19:16:29 -0700 Subject: [PATCH 02/10] [Discover-next] Address comments for search bar extensions and query assist (#6933) * pass dependencies to isEnabled func Signed-off-by: Joshua Li * add lazy and memo to search bar extensions Signed-off-by: Joshua Li * move ppl specific string out from query assist Signed-off-by: Joshua Li * prevent setstate after hook unmounts Signed-off-by: Joshua Li * add max-height to search bar extensions Signed-off-by: Joshua Li * prevent setstate after component unmounts Signed-off-by: Joshua Li * move ml-commons API to common/index.ts Signed-off-by: Joshua Li * improve i18n and accessibility usages Signed-off-by: Joshua Li * add hard-coded suggestions for sample data indices Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li (cherry picked from commit 4aade0f993559b0bae9cbcee8e889868afa88547) --- public/plugin.tsx | 2 +- public/query_assist/components/call_outs.tsx | 40 ++++++++++------- .../components/query_assist_bar.tsx | 21 ++++----- .../components/query_assist_input.tsx | 43 ++++++++++++++++--- .../query_assist/components/submit_button.tsx | 2 +- public/query_assist/hooks/use_generate.ts | 20 +++++++-- .../query_assist/utils/create_extension.tsx | 13 ++++-- server/routes/query_assist/agents.ts | 6 +-- 8 files changed, 100 insertions(+), 47 deletions(-) diff --git a/public/plugin.tsx b/public/plugin.tsx index f09dab1..e822443 100644 --- a/public/plugin.tsx +++ b/public/plugin.tsx @@ -55,7 +55,7 @@ export class QueryEnhancementsPlugin initialTo: moment().add(2, 'days').toISOString(), }, showFilterBar: false, - extensions: [createQueryAssistExtension(core.http)], + extensions: [createQueryAssistExtension(core.http, 'PPL')], }, fields: { filterable: false, diff --git a/public/query_assist/components/call_outs.tsx b/public/query_assist/components/call_outs.tsx index f98569f..fdbe901 100644 --- a/public/query_assist/components/call_outs.tsx +++ b/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>; -interface QueryAssistCallOutProps extends CalloutDismiss { +interface QueryAssistCallOutProps extends Required> { + language: string; type: QueryAssistCallOutType; } @@ -14,10 +15,12 @@ export type QueryAssistCallOutType = | 'empty_index' | 'query_generated'; -const EmptyIndexCallOut: React.FC = (props) => ( +const EmptyIndexCallOut: React.FC = (props) => ( = (props) => ( /> ); -const ProhibitedQueryCallOut: React.FC = (props) => ( +const ProhibitedQueryCallOut: React.FC = (props) => ( = (props) => ( /> ); -const EmptyQueryCallOut: React.FC = (props) => ( +const EmptyQueryCallOut: React.FC = (props) => ( = (props) => ( /> ); -const PPLGeneratedCallOut: React.FC = (props) => ( +const QueryGeneratedCallOut: React.FC = (props) => ( = (props) => ( export const QueryAssistCallOut: React.FC = (props) => { switch (props.type) { case 'empty_query': - return ; + return ; case 'empty_index': - return ; + return ; case 'invalid_query': - return ; + return ; case 'query_generated': - return ; + return ; default: break; } diff --git a/public/query_assist/components/query_assist_bar.tsx b/public/query_assist/components/query_assist_bar.tsx index c80fe5d..904382c 100644 --- a/public/query_assist/components/query_assist_bar.tsx +++ b/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 = (props) => { const { generateQuery, loading } = useGenerateQuery(); const [callOutType, setCallOutType] = useState(); 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(); - 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 = (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 = (props) => { - + ); }; diff --git a/public/query_assist/components/query_assist_input.tsx b/public/query_assist/components/query_assist_input.tsx index 735ff55..476e2d1 100644 --- a/public/query_assist/components/query_assist_input.tsx +++ b/public/query_assist/components/query_assist_input.tsx @@ -20,18 +20,49 @@ export const QueryAssistInput: React.FC = (props) => { const [suggestionIndex, setSuggestionIndex] = useState(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 ( setIsSuggestionsVisible(false)}> @@ -54,7 +85,7 @@ export const QueryAssistInput: React.FC = (props) => { { if (!props.inputRef.current) return; diff --git a/public/query_assist/components/submit_button.tsx b/public/query_assist/components/submit_button.tsx index 3b7a3c8..2896a2e 100644 --- a/public/query_assist/components/submit_button.tsx +++ b/public/query_assist/components/submit_button.tsx @@ -13,7 +13,7 @@ export const QueryAssistSubmitButton: React.FC = (props) => { isDisabled={props.isDisabled} size="s" type="submit" - aria-label="submit-question" + aria-label="Submit question to query assistant" /> ); }; diff --git a/public/query_assist/hooks/use_generate.ts b/public/query_assist/hooks/use_generate.ts index 0d23114..5365a7c 100644 --- a/public/query_assist/hooks/use_generate.ts +++ b/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(); const { services } = useOpenSearchDashboards(); - 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 }; diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx index f2f046d..ef3bcb7 100644 --- a/public/query_assist/utils/create_extension.tsx +++ b/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) => , + getComponent: (dependencies) => ( + + ), }; }; diff --git a/server/routes/query_assist/agents.ts b/server/routes/query_assist/agents.ts index 676ea01..b20ccee 100644 --- a/server/routes/query_assist/agents.ts +++ b/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 { 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: `${URI.ML}/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: `${URI.ML}/agents/${agentId}/_execute`, body, }, AGENT_REQUEST_OPTIONS From eb0e6f36d752563b168b15e79630580cfacc146f Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Fri, 7 Jun 2024 15:08:18 -0700 Subject: [PATCH 03/10] [Discover-next] Support data sources for query assist (#6972) * disable query assist for non-default datasource Signed-off-by: Joshua Li * disable query assist input when loading Signed-off-by: Joshua Li * support MDS for query assist Signed-off-by: Joshua Li * add unit tests for agents Signed-off-by: Joshua Li * Revert "add unit tests for agents" This reverts commit 983514ee11362c5efe4cdb59802b3ff402b61ef2. The test configs are not yet setup in query_enhancements plugins. Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li (cherry picked from commit 328e08e688c39de1f47fee1c357e9928c0576390) --- common/query_assist/index.ts | 2 ++ .../components/query_assist_bar.tsx | 11 ++++-- .../components/query_assist_input.tsx | 2 ++ .../query_assist/utils/create_extension.tsx | 36 +++++++++++++------ public/query_assist/utils/get_mds_id.ts | 11 ++++++ public/query_assist/utils/index.ts | 1 + server/plugin.ts | 7 +++- server/routes/index.ts | 2 +- server/routes/query_assist/agents.ts | 8 +++-- server/routes/query_assist/routes.ts | 14 ++++++-- server/types.ts | 10 ++++++ 11 files changed, 84 insertions(+), 20 deletions(-) create mode 100644 public/query_assist/utils/get_mds_id.ts diff --git a/common/query_assist/index.ts b/common/query_assist/index.ts index d4173b1..2834fbd 100644 --- a/common/query_assist/index.ts +++ b/common/query_assist/index.ts @@ -11,4 +11,6 @@ export interface QueryAssistParameters { question: string; index: string; language: string; + // for MDS + dataSourceId?: string; } diff --git a/public/query_assist/components/query_assist_bar.tsx b/public/query_assist/components/query_assist_bar.tsx index 904382c..db1a24c 100644 --- a/public/query_assist/components/query_assist_bar.tsx +++ b/public/query_assist/components/query_assist_bar.tsx @@ -3,9 +3,10 @@ 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'; +import { QueryAssistParameters } from '../../../common/query_assist'; import { getStorage } from '../../services'; import { useGenerateQuery } from '../hooks'; -import { getPersistedLog, ProhibitedQueryError } from '../utils'; +import { getMdsDataSourceId, getPersistedLog, ProhibitedQueryError } from '../utils'; import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs'; import { QueryAssistInput } from './query_assist_input'; import { QueryAssistSubmitButton } from './submit_button'; @@ -45,10 +46,15 @@ export const QueryAssistBar: React.FC = (props) => { dismissCallout(); previousQuestionRef.current = inputRef.current.value; persistedLog.add(inputRef.current.value); - const params = { + const dataSourceId = await getMdsDataSourceId( + services.data.indexPatterns, + selectedIndexPattern + ); + const params: QueryAssistParameters = { question: inputRef.current.value, index: selectedIndex, language: props.language, + dataSourceId, }; const { response, error } = await generateQuery(params); if (error) { @@ -75,6 +81,7 @@ export const QueryAssistBar: React.FC = (props) => { diff --git a/public/query_assist/components/query_assist_input.tsx b/public/query_assist/components/query_assist_input.tsx index 476e2d1..248dc54 100644 --- a/public/query_assist/components/query_assist_input.tsx +++ b/public/query_assist/components/query_assist_input.tsx @@ -7,6 +7,7 @@ import { getData } from '../../services'; interface QueryAssistInputProps { inputRef: React.RefObject; persistedLog: PersistedLog; + isDisabled: boolean; initialValue?: string; selectedIndex?: string; previousQuestion?: string; @@ -70,6 +71,7 @@ export const QueryAssistInput: React.FC = (props) => { setIsSuggestionsVisible(true)} onChange={(e) => setValue(e.target.value)} onKeyDown={() => setIsSuggestionsVisible(true)} diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx index ef3bcb7..0993d9f 100644 --- a/public/query_assist/utils/create_extension.tsx +++ b/public/query_assist/utils/create_extension.tsx @@ -1,7 +1,9 @@ -import React from 'react'; import { HttpSetup } from 'opensearch-dashboards/public'; -import { QueryAssistBar } from '../components'; +import React from 'react'; +import { getMdsDataSourceId } from '.'; import { SearchBarExtensionConfig } from '../../../../../src/plugins/data/public/ui/search_bar_extensions'; +import { getData } from '../../services'; +import { QueryAssistBar } from '../components'; export const createQueryAssistExtension = ( http: HttpSetup, @@ -11,15 +13,27 @@ export const createQueryAssistExtension = ( 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/${language}`) - .then((response) => response.configured) - .catch(() => false); - } - return agentConfigured; + const agentConfiguredMap: Map = new Map(); + return async (dependencies) => { + // currently query assist tool relies on opensearch API to get index + // mappings, other data sources are not supported + if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default') + return false; + + const dataSourceId = await getMdsDataSourceId( + getData().indexPatterns, + dependencies.indexPatterns?.at(0) + ); + const cached = agentConfiguredMap.get(dataSourceId); + if (cached !== undefined) return cached; + const configured = await http + .get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`, { + query: { dataSourceId }, + }) + .then((response) => response.configured) + .catch(() => false); + agentConfiguredMap.set(dataSourceId, configured); + return configured; }; })(), getComponent: (dependencies) => ( diff --git a/public/query_assist/utils/get_mds_id.ts b/public/query_assist/utils/get_mds_id.ts new file mode 100644 index 0000000..6dab73e --- /dev/null +++ b/public/query_assist/utils/get_mds_id.ts @@ -0,0 +1,11 @@ +import { IIndexPattern, IndexPatternsContract } from '../../../../../src/plugins/data/public'; + +export const getMdsDataSourceId = async ( + indexPatterns: IndexPatternsContract, + indexPattern: IIndexPattern | string | undefined +): Promise => { + if (!indexPattern || typeof indexPattern !== 'object' || !indexPattern.id) return undefined; + return indexPatterns + .get(indexPattern.id) + .then((indexPatternEntity) => indexPatternEntity.dataSourceRef?.id); +}; diff --git a/public/query_assist/utils/index.ts b/public/query_assist/utils/index.ts index 5384401..421893f 100644 --- a/public/query_assist/utils/index.ts +++ b/public/query_assist/utils/index.ts @@ -1,3 +1,4 @@ export * from './create_extension'; export * from './errors'; +export * from './get_mds_id'; export * from './get_persisted_log'; diff --git a/server/plugin.ts b/server/plugin.ts index a0ce775..a97ddfa 100644 --- a/server/plugin.ts +++ b/server/plugin.ts @@ -20,7 +20,7 @@ import { QueryEnhancementsPluginSetupDependencies, QueryEnhancementsPluginStart, } from './types'; -import { OpenSearchPPLPlugin, OpenSearchObservabilityPlugin } from './utils'; +import { OpenSearchObservabilityPlugin, OpenSearchPPLPlugin } from './utils'; export class QueryEnhancementsPlugin implements Plugin { @@ -45,6 +45,11 @@ export class QueryEnhancementsPlugin data.search.registerSearchStrategy(SEARCH_STRATEGY.PPL, pplSearchStrategy); data.search.registerSearchStrategy(SEARCH_STRATEGY.SQL, sqlSearchStrategy); + core.http.registerRouteHandlerContext('query_assist', () => ({ + logger: this.logger, + dataSourceEnabled: !!dataSource, + })); + defineRoutes(this.logger, router, { ppl: pplSearchStrategy, sql: sqlSearchStrategy, diff --git a/server/routes/index.ts b/server/routes/index.ts index 49366e2..467aa77 100644 --- a/server/routes/index.ts +++ b/server/routes/index.ts @@ -70,5 +70,5 @@ export function defineRoutes( ) { defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.PPL); defineRoute(logger, router, searchStrategies, SEARCH_STRATEGY.SQL); - registerQueryAssistRoutes(logger, router); + registerQueryAssistRoutes(router); } diff --git a/server/routes/query_assist/agents.ts b/server/routes/query_assist/agents.ts index b20ccee..ce93c61 100644 --- a/server/routes/query_assist/agents.ts +++ b/server/routes/query_assist/agents.ts @@ -47,9 +47,13 @@ export const requestAgentByConfig = async (options: { context: RequestHandlerContext; configName: string; body: RequestBody; + dataSourceId?: string; }): Promise => { - const { context, configName, body } = options; - const client = context.core.opensearch.client.asCurrentUser; + const { context, configName, body, dataSourceId } = options; + const client = + context.query_assist.dataSourceEnabled && dataSourceId + ? await context.dataSource.opensearch.getClient(dataSourceId) + : context.core.opensearch.client.asCurrentUser; const agentId = await getAgentIdByConfig(client, configName); return client.transport.request( { diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts index 0b8c5c8..8e9e151 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -1,12 +1,12 @@ import { schema, Type } from '@osd/config-schema'; -import { IRouter, Logger } from 'opensearch-dashboards/server'; +import { IRouter } from 'opensearch-dashboards/server'; import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors'; import { ERROR_DETAILS } from '../../../common/query_assist'; import { getAgentIdByConfig, requestAgentByConfig } from './agents'; import { AGENT_CONFIG_NAME_MAP } from './index'; import { createPPLResponseBody } from './ppl/create_response'; -export function registerQueryAssistRoutes(logger: Logger, router: IRouter) { +export function registerQueryAssistRoutes(router: IRouter) { const languageSchema = schema.oneOf( Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'PPL'>] ); @@ -18,10 +18,16 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) { params: schema.object({ language: languageSchema, }), + query: schema.object({ + dataSourceId: schema.maybe(schema.string()), + }), }, }, async (context, request, response) => { - const client = context.core.opensearch.client.asCurrentUser; + const client = + context.query_assist.dataSourceEnabled && request.query.dataSourceId + ? await context.dataSource.opensearch.getClient(request.query.dataSourceId) + : context.core.opensearch.client.asCurrentUser; try { // if the call does not throw any error, then the agent is properly configured await getAgentIdByConfig(client, AGENT_CONFIG_NAME_MAP[request.params.language]); @@ -40,6 +46,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) { index: schema.string(), question: schema.string(), language: languageSchema, + dataSourceId: schema.maybe(schema.string()), }), }, }, @@ -56,6 +63,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) { question: request.body.question, }, }, + dataSourceId: request.body.dataSourceId, }); const responseBody = createPPLResponseBody(agentResponse); return response.ok({ body: responseBody }); diff --git a/server/types.ts b/server/types.ts index 7335075..97a9379 100644 --- a/server/types.ts +++ b/server/types.ts @@ -4,6 +4,7 @@ */ import { DataPluginSetup } from 'src/plugins/data/server/plugin'; +import { Logger } from '../../../src/core/server'; import { DataSourcePluginStart } from '../../../src/plugins/data_source/server'; // eslint-disable-next-line @typescript-eslint/no-empty-interface @@ -48,3 +49,12 @@ export interface FacetRequest { format?: string; }; } + +declare module '../../../src/core/server' { + interface RequestHandlerContext { + query_assist: { + logger: Logger; + dataSourceEnabled: boolean; + }; + } +} From 56ab13b3541ddcfe3924ed3f075f842928d71d69 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Mon, 17 Jun 2024 12:33:09 -0700 Subject: [PATCH 04/10] add query assist banner Signed-off-by: Joshua Li --- common/query_assist/index.ts | 2 + public/plugin.tsx | 7 +- public/query_assist/components/call_outs.tsx | 40 +++++++---- .../components/query_assist_banner.tsx | 69 +++++++++++++++++++ .../components/query_assist_bar.tsx | 13 ++-- .../query_assist/utils/create_extension.tsx | 33 ++++++--- server/routes/query_assist/index.ts | 4 +- server/routes/query_assist/routes.ts | 8 +-- 8 files changed, 138 insertions(+), 38 deletions(-) create mode 100644 public/query_assist/components/query_assist_banner.tsx diff --git a/common/query_assist/index.ts b/common/query_assist/index.ts index 2834fbd..e36438e 100644 --- a/common/query_assist/index.ts +++ b/common/query_assist/index.ts @@ -2,6 +2,8 @@ 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; diff --git a/public/plugin.tsx b/public/plugin.tsx index e822443..e027c21 100644 --- a/public/plugin.tsx +++ b/public/plugin.tsx @@ -55,7 +55,6 @@ export class QueryEnhancementsPlugin initialTo: moment().add(2, 'days').toISOString(), }, showFilterBar: false, - extensions: [createQueryAssistExtension(core.http, 'PPL')], }, fields: { filterable: false, @@ -86,6 +85,12 @@ export class QueryEnhancementsPlugin }, }); + data.__enhance({ + ui: { + queryEditorExtension: createQueryAssistExtension(core.http), + }, + }); + return {}; } diff --git a/public/query_assist/components/call_outs.tsx b/public/query_assist/components/call_outs.tsx index fdbe901..3074ec7 100644 --- a/public/query_assist/components/call_outs.tsx +++ b/public/query_assist/components/call_outs.tsx @@ -1,5 +1,5 @@ import { EuiCallOut, EuiCallOutProps } from '@elastic/eui'; -import { i18n } from '@osd/i18n'; +import { FormattedMessage } from '@osd/i18n/react'; import React from 'react'; interface QueryAssistCallOutProps extends Required> { @@ -18,9 +18,12 @@ export type QueryAssistCallOutType = const EmptyIndexCallOut: React.FC = (props) => ( + } size="s" color="warning" iconType="iInCircle" @@ -32,9 +35,12 @@ const EmptyIndexCallOut: React.FC = (props) => ( const ProhibitedQueryCallOut: React.FC = (props) => ( + } size="s" color="danger" iconType="alert" @@ -46,10 +52,12 @@ const ProhibitedQueryCallOut: React.FC = (props) => ( const EmptyQueryCallOut: React.FC = (props) => ( + } size="s" color="warning" iconType="iInCircle" @@ -61,9 +69,13 @@ const EmptyQueryCallOut: React.FC = (props) => ( const QueryGeneratedCallOut: React.FC = (props) => ( + } size="s" color="success" iconType="check" diff --git a/public/query_assist/components/query_assist_banner.tsx b/public/query_assist/components/query_assist_banner.tsx new file mode 100644 index 0000000..1b6cfc8 --- /dev/null +++ b/public/query_assist/components/query_assist_banner.tsx @@ -0,0 +1,69 @@ +import { + EuiBadge, + EuiCallOut, + EuiFlexGroup, + EuiFlexItem, + EuiIcon, + EuiLink, + EuiTextColor, +} from '@elastic/eui'; +import { FormattedMessage } from '@osd/i18n/react'; +import React, { useState } from 'react'; +import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension'; +import { SUPPORTED_LANGUAGES } from '../../../common/query_assist'; +import assistantLogo from '../../assets/query_assist_logo.svg'; +import { getStorage } from '../../services'; + +const BANNER_STORAGE_KEY = 'queryAssist:banner:show'; + +interface QueryAssistBannerProps { + dependencies: QueryEditorExtensionDependencies; +} + +export const QueryAssistBanner: React.FC = (props) => { + const storage = getStorage(); + const [showCallOut, _setShowCallOut] = useState(true); + const setShowCallOut: typeof _setShowCallOut = (show) => { + if (!show) { + storage.set(BANNER_STORAGE_KEY, false); + } + _setShowCallOut(show); + }; + + if (!showCallOut || storage.get(BANNER_STORAGE_KEY) === false) return null; + + return ( + + + + + + + + + + + + + + + + + + + } + dismissible + onDismiss={() => setShowCallOut(false)} + /> + ); +}; diff --git a/public/query_assist/components/query_assist_bar.tsx b/public/query_assist/components/query_assist_bar.tsx index db1a24c..bd1618e 100644 --- a/public/query_assist/components/query_assist_bar.tsx +++ b/public/query_assist/components/query_assist_bar.tsx @@ -1,7 +1,7 @@ import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui'; 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 { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension'; import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public'; import { QueryAssistParameters } from '../../../common/query_assist'; import { getStorage } from '../../services'; @@ -12,8 +12,7 @@ import { QueryAssistInput } from './query_assist_input'; import { QueryAssistSubmitButton } from './submit_button'; interface QueryAssistInputProps { - language: string; - dependencies: SearchBarExtensionDependencies; + dependencies: QueryEditorExtensionDependencies; } export const QueryAssistBar: React.FC = (props) => { @@ -53,7 +52,7 @@ export const QueryAssistBar: React.FC = (props) => { const params: QueryAssistParameters = { question: inputRef.current.value, index: selectedIndex, - language: props.language, + language: props.dependencies.language, dataSourceId, }; const { response, error } = await generateQuery(params); @@ -91,7 +90,11 @@ export const QueryAssistBar: React.FC = (props) => { - + ); }; diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx index 0993d9f..5c56aa9 100644 --- a/public/query_assist/utils/create_extension.tsx +++ b/public/query_assist/utils/create_extension.tsx @@ -1,14 +1,13 @@ import { HttpSetup } from 'opensearch-dashboards/public'; import React from 'react'; import { getMdsDataSourceId } from '.'; -import { SearchBarExtensionConfig } from '../../../../../src/plugins/data/public/ui/search_bar_extensions'; +import { QueryEditorExtensionConfig } from '../../../../../src/plugins/data/public/ui/query_editor'; +import { SUPPORTED_LANGUAGES } from '../../../common/query_assist'; import { getData } from '../../services'; import { QueryAssistBar } from '../components'; +import { QueryAssistBanner } from '../components/query_assist_banner'; -export const createQueryAssistExtension = ( - http: HttpSetup, - language: string -): SearchBarExtensionConfig => { +export const createQueryAssistExtension = (http: HttpSetup): QueryEditorExtensionConfig => { return { id: 'query-assist', order: 1000, @@ -27,17 +26,29 @@ export const createQueryAssistExtension = ( const cached = agentConfiguredMap.get(dataSourceId); if (cached !== undefined) return cached; const configured = await http - .get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`, { - query: { dataSourceId }, - }) + .get<{ configured: boolean }>( + `/api/ql/query_assist/configured/${dependencies.language}`, + { + query: { dataSourceId }, + } + ) .then((response) => response.configured) .catch(() => false); agentConfiguredMap.set(dataSourceId, configured); return configured; }; })(), - getComponent: (dependencies) => ( - - ), + getComponent: (dependencies) => { + // only show the component if user is on a supported language. + // @ts-expect-error language can be an arbitrary string and fail the check + if (!SUPPORTED_LANGUAGES.includes(dependencies.language)) return null; + return ; + }, + getBanner: (dependencies) => { + // advertise query assist if user is not on a supported language. + // @ts-expect-error language can be an arbitrary string and fail the check + if (SUPPORTED_LANGUAGES.includes(dependencies.language)) return null; + return ; + }, }; }; diff --git a/server/routes/query_assist/index.ts b/server/routes/query_assist/index.ts index 0e1cdc3..3a62d86 100644 --- a/server/routes/query_assist/index.ts +++ b/server/routes/query_assist/index.ts @@ -1,5 +1,7 @@ +import { SUPPORTED_LANGUAGES } from '../../../common/query_assist'; + export { registerQueryAssistRoutes } from './routes'; -export const AGENT_CONFIG_NAME_MAP = { +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 8e9e151..677ee64 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -1,15 +1,13 @@ import { schema, Type } from '@osd/config-schema'; import { IRouter } from 'opensearch-dashboards/server'; import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors'; -import { ERROR_DETAILS } from '../../../common/query_assist'; +import { ERROR_DETAILS, SUPPORTED_LANGUAGES } from '../../../common/query_assist'; import { getAgentIdByConfig, requestAgentByConfig } from './agents'; import { AGENT_CONFIG_NAME_MAP } from './index'; import { createPPLResponseBody } from './ppl/create_response'; export function registerQueryAssistRoutes(router: IRouter) { - const languageSchema = schema.oneOf( - Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'PPL'>] - ); + const languageSchema = schema.oneOf(SUPPORTED_LANGUAGES.map(schema.literal) as [Type<'PPL'>]); router.get( { @@ -52,8 +50,6 @@ export function registerQueryAssistRoutes(router: IRouter) { }, async (context, request, response) => { try { - if (!(request.body.language in AGENT_CONFIG_NAME_MAP)) - throw new Error('Unsupported language.'); const agentResponse = await requestAgentByConfig({ context, configName: AGENT_CONFIG_NAME_MAP[request.body.language], From 43d1d433ca4ec3df081b95cb3c55bdb7c6185f8e Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 18 Jun 2024 18:24:31 +0000 Subject: [PATCH 05/10] implement banner rendering logic Signed-off-by: Joshua Li --- .../components/query_assist_banner.tsx | 6 +- .../query_assist/utils/create_extension.tsx | 108 ++++++++++++------ server/routes/query_assist/routes.ts | 20 ++-- 3 files changed, 86 insertions(+), 48 deletions(-) diff --git a/public/query_assist/components/query_assist_banner.tsx b/public/query_assist/components/query_assist_banner.tsx index 1b6cfc8..3611e7a 100644 --- a/public/query_assist/components/query_assist_banner.tsx +++ b/public/query_assist/components/query_assist_banner.tsx @@ -16,11 +16,7 @@ import { getStorage } from '../../services'; const BANNER_STORAGE_KEY = 'queryAssist:banner:show'; -interface QueryAssistBannerProps { - dependencies: QueryEditorExtensionDependencies; -} - -export const QueryAssistBanner: React.FC = (props) => { +export const QueryAssistBanner: React.FC = () => { const storage = getStorage(); const [showCallOut, _setShowCallOut] = useState(true); const setShowCallOut: typeof _setShowCallOut = (show) => { diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx index 5c56aa9..7a64afd 100644 --- a/public/query_assist/utils/create_extension.tsx +++ b/public/query_assist/utils/create_extension.tsx @@ -1,54 +1,92 @@ import { HttpSetup } from 'opensearch-dashboards/public'; -import React from 'react'; +import React, { useEffect, useState } from 'react'; import { getMdsDataSourceId } from '.'; import { QueryEditorExtensionConfig } from '../../../../../src/plugins/data/public/ui/query_editor'; -import { SUPPORTED_LANGUAGES } from '../../../common/query_assist'; +import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension'; import { getData } from '../../services'; import { QueryAssistBar } from '../components'; import { QueryAssistBanner } from '../components/query_assist_banner'; +let availableLanguagesByDataSource: Map; + +/** + * @param dependencies - QueryEditorExtensionDependencies. + * @param http - HttpSetup. + * @returns list of query assist agents configured languages in the data source + * associated with the currently selected index pattern. + */ +const getAvailableLanguages = async ( + dependencies: QueryEditorExtensionDependencies, + http: HttpSetup +) => { + if (!availableLanguagesByDataSource) availableLanguagesByDataSource = new Map(); + + const dataSourceId = await getMdsDataSourceId( + getData().indexPatterns, + dependencies.indexPatterns?.at(0) + ); + const cached = availableLanguagesByDataSource.get(dataSourceId); + if (cached !== undefined) return cached; + + const languages = await http + .get<{ configuredLanguages: string[] }>('/api/ql/query_assist/configured_languages', { + query: { dataSourceId }, + }) + .then((response) => response.configuredLanguages) + .catch(() => []); + availableLanguagesByDataSource.set(dataSourceId, languages); + return languages; +}; + export const createQueryAssistExtension = (http: HttpSetup): QueryEditorExtensionConfig => { return { id: 'query-assist', order: 1000, - isEnabled: (() => { - const agentConfiguredMap: Map = new Map(); - return async (dependencies) => { - // currently query assist tool relies on opensearch API to get index - // mappings, other data sources are not supported - if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default') - return false; - - const dataSourceId = await getMdsDataSourceId( - getData().indexPatterns, - dependencies.indexPatterns?.at(0) - ); - const cached = agentConfiguredMap.get(dataSourceId); - if (cached !== undefined) return cached; - const configured = await http - .get<{ configured: boolean }>( - `/api/ql/query_assist/configured/${dependencies.language}`, - { - query: { dataSourceId }, - } - ) - .then((response) => response.configured) - .catch(() => false); - agentConfiguredMap.set(dataSourceId, configured); - return configured; - }; - })(), + isEnabled: async (dependencies) => { + // currently query assist tool relies on opensearch API to get index + // mappings, non-default data source types are not supported + if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default') return false; + + const languages = await getAvailableLanguages(dependencies, http); + return languages.length > 0; + }, getComponent: (dependencies) => { // only show the component if user is on a supported language. - // @ts-expect-error language can be an arbitrary string and fail the check - if (!SUPPORTED_LANGUAGES.includes(dependencies.language)) return null; - return ; + return ( + + + + ); }, getBanner: (dependencies) => { // advertise query assist if user is not on a supported language. - // @ts-expect-error language can be an arbitrary string and fail the check - if (SUPPORTED_LANGUAGES.includes(dependencies.language)) return null; - return ; + return ( + + + + ); }, }; }; + +interface QueryAssistWrapperProps { + dependencies: QueryEditorExtensionDependencies; + http: HttpSetup; + invert?: boolean; +} + +const QueryAssistWrapper: React.FC = (props) => { + const [visible, setVisible] = useState(false); + useEffect(() => { + const checkAvailability = async () => { + const available = (await getAvailableLanguages(props.dependencies, props.http)).includes( + props.dependencies.language + ); + setVisible(props.invert ? !available : available); + }; + checkAvailability(); + }, [props]); + + if (!visible) return null; + return <>{props.children}; +}; diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts index 677ee64..c60035c 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -11,11 +11,8 @@ export function registerQueryAssistRoutes(router: IRouter) { router.get( { - path: '/api/ql/query_assist/configured/{language}', + path: '/api/ql/query_assist/configured_languages', validate: { - params: schema.object({ - language: languageSchema, - }), query: schema.object({ dataSourceId: schema.maybe(schema.string()), }), @@ -26,12 +23,19 @@ export function registerQueryAssistRoutes(router: IRouter) { context.query_assist.dataSourceEnabled && request.query.dataSourceId ? await context.dataSource.opensearch.getClient(request.query.dataSourceId) : context.core.opensearch.client.asCurrentUser; + const configuredLanguages: string[] = []; try { - // if the call does not throw any error, then the agent is properly configured - await getAgentIdByConfig(client, AGENT_CONFIG_NAME_MAP[request.params.language]); - return response.ok({ body: { configured: true } }); + 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) + ) + ) + ); + return response.ok({ body: { configuredLanguages } }); } catch (error) { - return response.ok({ body: { configured: false, error: error.message } }); + return response.ok({ body: { configuredLanguages, error: error.message } }); } } ); From c97cd2de85ec64184dbac22ff818655efd6e8df1 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 18 Jun 2024 18:42:00 +0000 Subject: [PATCH 06/10] pick opensearch-project/dashboards-observability/pull/1896 Signed-off-by: Joshua Li --- server/routes/query_assist/ppl/create_response.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/server/routes/query_assist/ppl/create_response.ts b/server/routes/query_assist/ppl/create_response.ts index ac651a0..5d541a0 100644 --- a/server/routes/query_assist/ppl/create_response.ts +++ b/server/routes/query_assist/ppl/create_response.ts @@ -11,7 +11,6 @@ export const createPPLResponseBody = (agentResponse: AgentResponse): QueryAssist .replace(/[\r\n]/g, ' ') .trim() .replace(/ISNOTNULL/g, 'isnotnull') // https://github.com/opensearch-project/sql/issues/2431 - .replace(/`/g, '') // https://github.com/opensearch-project/dashboards-observability/issues/509, https://github.com/opensearch-project/dashboards-observability/issues/557 .replace(/\bSPAN\(/g, 'span('); // https://github.com/opensearch-project/dashboards-observability/issues/759 return { query: ppl }; }; From 780fbd5dc72b001701322ed9ad9714badbb82d23 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Thu, 20 Jun 2024 20:47:51 +0000 Subject: [PATCH 07/10] add basic tests Signed-off-by: Joshua Li --- .gitignore | 3 +- babel.config.js | 25 +++ package.json | 1 + .../__snapshots__/call_outs.test.tsx.snap | 163 ++++++++++++++++++ .../components/call_outs.test.tsx | 46 +++++ .../components/query_assist_banner.tsx | 1 - .../components/query_assist_input.tsx | 2 +- .../components/submit_button.test.tsx | 29 ++++ .../query_assist/components/submit_button.tsx | 3 +- .../query_assist/utils/create_extension.tsx | 12 +- server/routes/query_assist/agents.test.ts | 102 +++++++++++ test/__mocks__/fileMock.js | 6 + test/__mocks__/styleMock.js | 6 + test/jest.config.js | 31 ++++ test/setup.jest.ts | 48 ++++++ test/setupTests.ts | 7 + 16 files changed, 478 insertions(+), 7 deletions(-) create mode 100644 babel.config.js create mode 100644 public/query_assist/components/__snapshots__/call_outs.test.tsx.snap create mode 100644 public/query_assist/components/call_outs.test.tsx create mode 100644 public/query_assist/components/submit_button.test.tsx create mode 100644 server/routes/query_assist/agents.test.ts create mode 100644 test/__mocks__/fileMock.js create mode 100644 test/__mocks__/styleMock.js create mode 100644 test/jest.config.js create mode 100644 test/setup.jest.ts create mode 100644 test/setupTests.ts diff --git a/.gitignore b/.gitignore index 289ed31..8a3d5ba 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /build -/target \ No newline at end of file +/target +/coverage/ diff --git a/babel.config.js b/babel.config.js new file mode 100644 index 0000000..16a64bf --- /dev/null +++ b/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 {}; +}; diff --git a/package.json b/package.json index c920fb9..6c23b9c 100644 --- a/package.json +++ b/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": {}, diff --git a/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap b/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap new file mode 100644 index 0000000..e7cccd4 --- /dev/null +++ b/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`] = ` +
+
+
+ + + + Select a data source or index to ask a question. + + +
+
+
+`; + +exports[`CallOuts spec should display empty_query call out 1`] = ` +
+
+
+ + + + Enter a natural language question to automatically generate a query to view results. + + +
+
+
+`; + +exports[`CallOuts spec should display invalid_query call out 1`] = ` +
+
+
+ + + + I am unable to respond to this query. Try another question. + + +
+
+
+`; + +exports[`CallOuts spec should display query_generated call out 1`] = ` +
+
+
+ + + + 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. + + +
+ +
+
+`; diff --git a/public/query_assist/components/call_outs.test.tsx b/public/query_assist/components/call_outs.test.tsx new file mode 100644 index 0000000..ddace95 --- /dev/null +++ b/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; + +const renderCallOut = (overrideProps: Partial = {}) => { + const props: Props = Object.assign>( + { + type: 'empty_query', + language: 'test lang', + onDismiss: jest.fn(), + }, + overrideProps + ); + const component = render(); + return { component, props: props as jest.MockedObjectDeep }; +}; + +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(); + }); +}); diff --git a/public/query_assist/components/query_assist_banner.tsx b/public/query_assist/components/query_assist_banner.tsx index 3611e7a..de054ed 100644 --- a/public/query_assist/components/query_assist_banner.tsx +++ b/public/query_assist/components/query_assist_banner.tsx @@ -9,7 +9,6 @@ import { } from '@elastic/eui'; import { FormattedMessage } from '@osd/i18n/react'; import React, { useState } from 'react'; -import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension'; import { SUPPORTED_LANGUAGES } from '../../../common/query_assist'; import assistantLogo from '../../assets/query_assist_logo.svg'; import { getStorage } from '../../services'; diff --git a/public/query_assist/components/query_assist_input.tsx b/public/query_assist/components/query_assist_input.tsx index 248dc54..7a96f4f 100644 --- a/public/query_assist/components/query_assist_input.tsx +++ b/public/query_assist/components/query_assist_input.tsx @@ -79,7 +79,7 @@ export const QueryAssistInput: React.FC = (props) => { props.previousQuestion || (props.selectedIndex ? `Ask a natural language question about ${props.selectedIndex} to generate a query` - : 'Select an index pattern to ask a question') + : 'Select an index to ask a question') } prepend={} fullWidth diff --git a/public/query_assist/components/submit_button.test.tsx b/public/query_assist/components/submit_button.test.tsx new file mode 100644 index 0000000..a3f4726 --- /dev/null +++ b/public/query_assist/components/submit_button.test.tsx @@ -0,0 +1,29 @@ +import { fireEvent, render } from '@testing-library/react'; +import React, { ComponentProps } from 'react'; +import { QueryAssistSubmitButton } from './submit_button'; + +type SubmitButtonProps = ComponentProps; + +const renderSubmitButton = (overrideProps: Partial = {}) => { + const props: SubmitButtonProps = Object.assign>( + { + isDisabled: false, + }, + overrideProps + ); + const onSubmit = jest.fn((e) => e.preventDefault()); + const component = render( +
+ + + ); + return { component, onSubmit, props: props as jest.MockedObjectDeep }; +}; + +describe(' spec', () => { + it('should trigger submit form', () => { + const { component, onSubmit } = renderSubmitButton(); + fireEvent.click(component.getByTestId('query-assist-submit-button')); + expect(onSubmit).toBeCalled(); + }); +}); diff --git a/public/query_assist/components/submit_button.tsx b/public/query_assist/components/submit_button.tsx index 2896a2e..ddae082 100644 --- a/public/query_assist/components/submit_button.tsx +++ b/public/query_assist/components/submit_button.tsx @@ -1,5 +1,5 @@ -import React from 'react'; import { EuiButtonIcon } from '@elastic/eui'; +import React from 'react'; interface SubmitButtonProps { isDisabled: boolean; @@ -14,6 +14,7 @@ export const QueryAssistSubmitButton: React.FC = (props) => { size="s" type="submit" aria-label="Submit question to query assistant" + data-test-subj="query-assist-submit-button" /> ); }; diff --git a/public/query_assist/utils/create_extension.tsx b/public/query_assist/utils/create_extension.tsx index 7a64afd..0fb0a49 100644 --- a/public/query_assist/utils/create_extension.tsx +++ b/public/query_assist/utils/create_extension.tsx @@ -77,14 +77,20 @@ interface QueryAssistWrapperProps { const QueryAssistWrapper: React.FC = (props) => { const [visible, setVisible] = useState(false); + useEffect(() => { - const checkAvailability = async () => { + let mounted = true; + + (async () => { const available = (await getAvailableLanguages(props.dependencies, props.http)).includes( props.dependencies.language ); - setVisible(props.invert ? !available : available); + if (mounted) setVisible(props.invert ? !available : available); + })(); + + return () => { + mounted = false; }; - checkAvailability(); }, [props]); if (!visible) return null; diff --git a/server/routes/query_assist/agents.test.ts b/server/routes/query_assist/agents.test.ts new file mode 100644 index 0000000..f3b4cfa --- /dev/null +++ b/server/routes/query_assist/agents.test.ts @@ -0,0 +1,102 @@ +import { ApiResponse } from '@opensearch-project/opensearch'; +import { ResponseError } from '@opensearch-project/opensearch/lib/errors'; +import { RequestHandlerContext } from 'src/core/server'; +import { CoreRouteHandlerContext } from '../../../../../src/core/server/core_route_handler_context'; +import { coreMock, httpServerMock } from '../../../../../src/core/server/mocks'; +import { loggerMock } from '../../../../../src/core/server/logging/logger.mock'; +import { getAgentIdByConfig, requestAgentByConfig } from './agents'; + +describe('Agents helper functions', () => { + const coreContext = new CoreRouteHandlerContext( + coreMock.createInternalStart(), + httpServerMock.createOpenSearchDashboardsRequest() + ); + const client = coreContext.opensearch.client.asCurrentUser; + const mockedTransport = client.transport.request as jest.Mock; + const context: RequestHandlerContext = { + core: coreContext, + dataSource: jest.fn(), + query_assist: { dataSourceEnabled: false, logger: loggerMock.create() }, + }; + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('searches agent id by name', async () => { + mockedTransport.mockResolvedValueOnce({ + body: { + type: 'agent', + configuration: { agent_id: 'agentId' }, + }, + }); + const id = await getAgentIdByConfig(client, 'test_agent'); + expect(id).toEqual('agentId'); + expect(mockedTransport.mock.calls[0]).toMatchInlineSnapshot(` + Array [ + Object { + "method": "GET", + "path": "/_plugins/_ml/config/test_agent", + }, + ] + `); + }); + + it('handles not found errors', async () => { + mockedTransport.mockRejectedValueOnce( + new ResponseError(({ + body: { + error: { + root_cause: [ + { + type: 'status_exception', + reason: 'Failed to find config with the provided config id: test_agent', + }, + ], + type: 'status_exception', + reason: 'Failed to find config with the provided config id: test_agent', + }, + status: 404, + }, + statusCode: 404, + } as unknown) as ApiResponse) + ); + await expect( + getAgentIdByConfig(client, 'test agent') + ).rejects.toThrowErrorMatchingInlineSnapshot( + `"Get agent 'test agent' failed, reason: {\\"error\\":{\\"root_cause\\":[{\\"type\\":\\"status_exception\\",\\"reason\\":\\"Failed to find config with the provided config id: test_agent\\"}],\\"type\\":\\"status_exception\\",\\"reason\\":\\"Failed to find config with the provided config id: test_agent\\"},\\"status\\":404}"` + ); + }); + + it('handles search errors', async () => { + mockedTransport.mockRejectedValueOnce('request failed'); + await expect( + getAgentIdByConfig(client, 'test agent') + ).rejects.toThrowErrorMatchingInlineSnapshot( + `"Get agent 'test agent' failed, reason: request failed"` + ); + }); + + it('searches for agent id and sends request', async () => { + mockedTransport + .mockResolvedValueOnce({ + body: { + type: 'agent', + configuration: { agent_id: 'new-id' }, + }, + }) + .mockResolvedValueOnce({ + body: { inference_results: [{ output: [{ result: 'test response' }] }] }, + }); + const response = await requestAgentByConfig({ + context, + configName: 'new_agent', + body: { parameters: { param1: 'value1' } }, + }); + expect(mockedTransport).toBeCalledWith( + expect.objectContaining({ path: '/_plugins/_ml/agents/new-id/_execute' }), + expect.anything() + ); + expect(response.body.inference_results[0].output[0].result).toEqual('test response'); + }); +}); diff --git a/test/__mocks__/fileMock.js b/test/__mocks__/fileMock.js new file mode 100644 index 0000000..cac247b --- /dev/null +++ b/test/__mocks__/fileMock.js @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +module.exports = 'file-stub'; diff --git a/test/__mocks__/styleMock.js b/test/__mocks__/styleMock.js new file mode 100644 index 0000000..28de3c8 --- /dev/null +++ b/test/__mocks__/styleMock.js @@ -0,0 +1,6 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +module.exports = {}; diff --git a/test/jest.config.js b/test/jest.config.js new file mode 100644 index 0000000..f5a1419 --- /dev/null +++ b/test/jest.config.js @@ -0,0 +1,31 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +process.env.TZ = 'UTC'; + +module.exports = { + rootDir: '../', + setupFiles: ['/test/setupTests.ts'], + setupFilesAfterEnv: ['/test/setup.jest.ts'], + roots: [''], + testMatch: ['**/*.test.js', '**/*.test.jsx', '**/*.test.ts', '**/*.test.tsx'], + clearMocks: true, + modulePathIgnorePatterns: ['/offline-module-cache/'], + testPathIgnorePatterns: ['/build/', '/node_modules/', '/__utils__/'], + snapshotSerializers: ['enzyme-to-json/serializer'], + coveragePathIgnorePatterns: [ + '/build/', + '/node_modules/', + '/test/', + '/public/requests/', + '/__utils__/', + ], + moduleNameMapper: { + '\\.(css|less|sass|scss)$': '/test/__mocks__/styleMock.js', + '\\.(gif|ttf|eot|svg|png)$': '/test/__mocks__/fileMock.js', + '^!!raw-loader!.*': 'jest-raw-loader', + }, + testEnvironment: 'jsdom', +}; diff --git a/test/setup.jest.ts b/test/setup.jest.ts new file mode 100644 index 0000000..9d19503 --- /dev/null +++ b/test/setup.jest.ts @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { configure } from '@testing-library/react'; +import { TextDecoder, TextEncoder } from 'util'; +import '@testing-library/jest-dom'; + +configure({ testIdAttribute: 'data-test-subj' }); + +// https://github.com/inrupt/solid-client-authn-js/issues/1676#issuecomment-917016646 +global.TextEncoder = TextEncoder; +global.TextDecoder = TextDecoder as typeof global.TextDecoder; + +window.URL.createObjectURL = () => ''; +// eslint-disable-next-line @typescript-eslint/no-explicit-any +HTMLCanvasElement.prototype.getContext = () => '' as any; +Element.prototype.scrollIntoView = jest.fn(); +window.IntersectionObserver = (class IntersectionObserver { + constructor() {} + + disconnect() { + return null; + } + + observe() { + return null; + } + + takeRecords() { + return null; + } + + unobserve() { + return null; + } +} as unknown) as typeof window.IntersectionObserver; + +jest.mock('@elastic/eui/lib/components/form/form_row/make_id', () => () => 'random-id'); + +jest.mock('@elastic/eui/lib/services/accessibility/html_id_generator', () => ({ + htmlIdGenerator: () => { + return () => 'random_html_id'; + }, +})); + +jest.setTimeout(30000); diff --git a/test/setupTests.ts b/test/setupTests.ts new file mode 100644 index 0000000..5a996f6 --- /dev/null +++ b/test/setupTests.ts @@ -0,0 +1,7 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +require('babel-polyfill'); +require('core-js/stable'); From 9b811a5459b63a3c8fc29ee25003df47df68c5e6 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Fri, 21 Jun 2024 19:59:01 +0000 Subject: [PATCH 08/10] add index selector 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 --- .../components/index_selector.tsx | 42 ++++++++ .../components/query_assist_bar.tsx | 29 ++++-- public/query_assist/hooks/use_indices.ts | 97 +++++++++++++++++++ 3 files changed, 158 insertions(+), 10 deletions(-) create mode 100644 public/query_assist/components/index_selector.tsx create mode 100644 public/query_assist/hooks/use_indices.ts diff --git a/public/query_assist/components/index_selector.tsx b/public/query_assist/components/index_selector.tsx new file mode 100644 index 0000000..68aa418 --- /dev/null +++ b/public/query_assist/components/index_selector.tsx @@ -0,0 +1,42 @@ +import { EuiComboBox, EuiComboBoxOptionOption, EuiText } from '@elastic/eui'; +import React from 'react'; +import { useIndexPatterns, useIndices } from '../hooks/use_indices'; + +interface IndexSelectorProps { + dataSourceId?: string; + selectedIndex?: string; + setSelectedIndex: React.Dispatch>; +} + +// TODO this is a temporary solution, there will be a dataset selector from discover +export const IndexSelector: React.FC = (props) => { + const { data: indices, loading: indicesLoading } = useIndices(props.dataSourceId); + const { data: indexPatterns, loading: indexPatternsLoading } = useIndexPatterns(); + const loading = indicesLoading || indexPatternsLoading; + const indicesAndIndexPatterns = + indexPatterns && indices + ? [...indexPatterns, ...indices].filter( + (v1, index, array) => array.findIndex((v2) => v1 === v2) === index + ) + : []; + const options: EuiComboBoxOptionOption[] = indicesAndIndexPatterns.map((index) => ({ + label: index, + })); + const selectedOptions = props.selectedIndex ? [{ label: props.selectedIndex }] : undefined; + + return ( + Index} + singleSelection={{ asPlainText: true }} + isLoading={loading} + options={options} + selectedOptions={selectedOptions} + onChange={(index) => { + props.setSelectedIndex(index[0].label); + }} + /> + ); +}; diff --git a/public/query_assist/components/query_assist_bar.tsx b/public/query_assist/components/query_assist_bar.tsx index bd1618e..98c6caf 100644 --- a/public/query_assist/components/query_assist_bar.tsx +++ b/public/query_assist/components/query_assist_bar.tsx @@ -1,5 +1,5 @@ import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui'; -import React, { SyntheticEvent, useMemo, useRef, useState } from 'react'; +import React, { SyntheticEvent, useEffect, useMemo, useRef, useState } from 'react'; import { IDataPluginServices, PersistedLog } from '../../../../../src/plugins/data/public'; import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension'; import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public'; @@ -8,6 +8,7 @@ import { getStorage } from '../../services'; import { useGenerateQuery } from '../hooks'; import { getMdsDataSourceId, getPersistedLog, ProhibitedQueryError } from '../utils'; import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs'; +import { IndexSelector } from './index_selector'; import { QueryAssistInput } from './query_assist_input'; import { QueryAssistSubmitButton } from './submit_button'; @@ -26,12 +27,17 @@ export const QueryAssistBar: React.FC = (props) => { const { generateQuery, loading } = useGenerateQuery(); const [callOutType, setCallOutType] = useState(); const dismissCallout = () => setCallOutType(undefined); - const selectedIndexPattern = props.dependencies.indexPatterns?.at(0); - const selectedIndex = - selectedIndexPattern && - (typeof selectedIndexPattern === 'string' ? selectedIndexPattern : selectedIndexPattern.title); + const [selectedIndex, setSelectedIndex] = useState(''); + const dataSourceIdRef = useRef(); const previousQuestionRef = useRef(); + useEffect(() => { + // TODO need proper way to get dataSourceId when discover index pattern selector is removed + getMdsDataSourceId(services.data.indexPatterns, props.dependencies.indexPatterns?.at(0)).then( + (id) => (dataSourceIdRef.current = id) + ); + }, [props.dependencies.indexPatterns, services.data.indexPatterns]); + const onSubmit = async (e: SyntheticEvent) => { e.preventDefault(); if (!inputRef.current?.value) { @@ -45,15 +51,11 @@ export const QueryAssistBar: React.FC = (props) => { dismissCallout(); previousQuestionRef.current = inputRef.current.value; persistedLog.add(inputRef.current.value); - const dataSourceId = await getMdsDataSourceId( - services.data.indexPatterns, - selectedIndexPattern - ); const params: QueryAssistParameters = { question: inputRef.current.value, index: selectedIndex, language: props.dependencies.language, - dataSourceId, + dataSourceId: dataSourceIdRef.current, }; const { response, error } = await generateQuery(params); if (error) { @@ -76,6 +78,13 @@ export const QueryAssistBar: React.FC = (props) => { + + + { + data?: T; + loading: boolean; + error?: Error; +} + +type Action = + | { type: 'request' } + | { type: 'success'; payload: State['data'] } + | { type: 'failure'; error: NonNullable['error']> }; + +// TODO use instantiation expressions when typescript is upgraded to >= 4.7 +type GenericReducer = Reducer, Action>; +export const genericReducer: GenericReducer = (state, action) => { + switch (action.type) { + case 'request': + return { data: state.data, loading: true }; + case 'success': + return { loading: false, data: action.payload }; + case 'failure': + return { loading: false, error: action.error }; + default: + return state; + } +}; + +export const useIndices = (dataSourceId: string | undefined) => { + const reducer: GenericReducer = genericReducer; + const [state, dispatch] = useReducer(reducer, { loading: false }); + const [refresh, setRefresh] = useState({}); + const { services } = useOpenSearchDashboards(); + + useEffect(() => { + const abortController = new AbortController(); + dispatch({ type: 'request' }); + services.http + .post('/api/console/proxy', { + query: { path: '_cat/indices?format=json', method: 'GET', dataSourceId }, + signal: abortController.signal, + }) + .then((payload: CatIndicesResponse) => + dispatch({ + type: 'success', + payload: payload + .filter((meta) => meta.index && !meta.index.startsWith('.')) + .map((meta) => meta.index!), + }) + ) + .catch((error) => dispatch({ type: 'failure', error })); + + return () => abortController.abort(); + }, [refresh, services.http, dataSourceId]); + + return { ...state, refresh: () => setRefresh({}) }; +}; + +export const useIndexPatterns = () => { + const reducer: GenericReducer = genericReducer; + const [state, dispatch] = useReducer(reducer, { loading: false }); + const [refresh, setRefresh] = useState({}); + const { services } = useOpenSearchDashboards(); + + useEffect(() => { + let abort = false; + dispatch({ type: 'request' }); + + services.data.indexPatterns + .getTitles() + .then((payload) => { + if (!abort) + dispatch({ + type: 'success', + // temporary solution does not support index patterns from other data sources + payload: payload.filter((title) => !title.includes('::')), + }); + }) + .catch((error) => { + if (!abort) dispatch({ type: 'failure', error }); + }); + + return () => { + abort = true; + }; + }, [refresh, services.data.indexPatterns]); + + return { ...state, refresh: () => setRefresh({}) }; +}; From fd651e7a3e0ef3b6c8a00e9a76cc3eadfc99dad4 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 25 Jun 2024 20:24:22 +0000 Subject: [PATCH 09/10] rename logo to mark Signed-off-by: Joshua Li --- .../assets/{query_assist_logo.svg => query_assist_mark.svg} | 0 public/query_assist/components/query_assist_banner.tsx | 4 ++-- public/query_assist/components/query_assist_input.tsx | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) rename public/assets/{query_assist_logo.svg => query_assist_mark.svg} (100%) diff --git a/public/assets/query_assist_logo.svg b/public/assets/query_assist_mark.svg similarity index 100% rename from public/assets/query_assist_logo.svg rename to public/assets/query_assist_mark.svg diff --git a/public/query_assist/components/query_assist_banner.tsx b/public/query_assist/components/query_assist_banner.tsx index de054ed..f1cab31 100644 --- a/public/query_assist/components/query_assist_banner.tsx +++ b/public/query_assist/components/query_assist_banner.tsx @@ -10,7 +10,7 @@ import { import { FormattedMessage } from '@osd/i18n/react'; import React, { useState } from 'react'; import { SUPPORTED_LANGUAGES } from '../../../common/query_assist'; -import assistantLogo from '../../assets/query_assist_logo.svg'; +import assistantMark from '../../assets/query_assist_mark.svg'; import { getStorage } from '../../services'; const BANNER_STORAGE_KEY = 'queryAssist:banner:show'; @@ -33,7 +33,7 @@ export const QueryAssistBanner: React.FC = () => { title={ - + diff --git a/public/query_assist/components/query_assist_input.tsx b/public/query_assist/components/query_assist_input.tsx index 7a96f4f..22f8382 100644 --- a/public/query_assist/components/query_assist_input.tsx +++ b/public/query_assist/components/query_assist_input.tsx @@ -1,7 +1,7 @@ import { EuiFieldText, EuiIcon, EuiOutsideClickDetector, EuiPortal } from '@elastic/eui'; import React, { useMemo, useState } from 'react'; import { PersistedLog, QuerySuggestionTypes } from '../../../../../src/plugins/data/public'; -import assistantLogo from '../../assets/query_assist_logo.svg'; +import assistantMark from '../../assets/query_assist_mark.svg'; import { getData } from '../../services'; interface QueryAssistInputProps { @@ -81,7 +81,7 @@ export const QueryAssistInput: React.FC = (props) => { ? `Ask a natural language question about ${props.selectedIndex} to generate a query` : 'Select an index to ask a question') } - prepend={} + prepend={} fullWidth /> From 8ad8600b03ee22976cf2c146e64bdefec0356ea4 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 25 Jun 2024 20:38:17 +0000 Subject: [PATCH 10/10] extract supported languages to config Signed-off-by: Joshua Li --- common/config.ts | 11 +++++++ common/constants.ts | 6 ++++ common/query_assist/index.ts | 19 +---------- common/query_assist/types.ts | 14 ++++++++ public/index.ts | 6 ++-- public/plugin.tsx | 11 +++++-- .../components/query_assist_banner.tsx | 9 +++-- public/query_assist/hooks/use_generate.ts | 12 +++---- .../query_assist/utils/create_extension.tsx | 13 ++++++-- public/query_assist/utils/errors.ts | 2 +- server/index.ts | 4 ++- server/plugin.ts | 8 ++++- server/routes/query_assist/createResponse.ts | 23 +++++++++++++ server/routes/query_assist/index.ts | 6 ---- server/routes/query_assist/routes.ts | 33 ++++++++++--------- server/types.ts | 2 ++ 16 files changed, 118 insertions(+), 61 deletions(-) create mode 100644 common/query_assist/types.ts create mode 100644 server/routes/query_assist/createResponse.ts 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; }; }