From 016e0f2f73efd8bb0649151908c67dd7ac09d174 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 4 Jun 2024 15:11:58 -0700 Subject: [PATCH] [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 --- .../common/query_assist/index.ts | 14 +++ .../opensearch_dashboards.json | 2 +- .../public/assets/query_assist_logo.svg | 18 ++++ .../public/{plugin.ts => plugin.tsx} | 25 ++++- .../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 + .../query_enhancements/public/services.ts | 6 ++ .../query_enhancements/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 +++++++++++++++ src/plugins/data/public/index.ts | 1 - .../public/ui/query_editor/query_editor.tsx | 4 +- .../ui/query_editor/query_editor_top_row.tsx | 25 ++++- .../data/public/ui/search_bar/search_bar.tsx | 31 ------ .../search_bar_extension.tsx | 2 +- src/plugins/data/public/ui/types.ts | 4 +- src/plugins/data/public/ui/ui_service.ts | 2 + 29 files changed, 631 insertions(+), 46 deletions(-) create mode 100644 plugins-extra/query_enhancements/common/query_assist/index.ts create mode 100644 plugins-extra/query_enhancements/public/assets/query_assist_logo.svg rename plugins-extra/query_enhancements/public/{plugin.ts => plugin.tsx} (78%) create mode 100644 plugins-extra/query_enhancements/public/query_assist/components/call_outs.tsx create mode 100644 plugins-extra/query_enhancements/public/query_assist/components/index.ts create mode 100644 plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx create mode 100644 plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx create mode 100644 plugins-extra/query_enhancements/public/query_assist/components/submit_button.tsx create mode 100644 plugins-extra/query_enhancements/public/query_assist/hooks/index.ts create mode 100644 plugins-extra/query_enhancements/public/query_assist/hooks/use_generate.ts create mode 100644 plugins-extra/query_enhancements/public/query_assist/index.ts create mode 100644 plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx create mode 100644 plugins-extra/query_enhancements/public/query_assist/utils/errors.ts create mode 100644 plugins-extra/query_enhancements/public/query_assist/utils/get_persisted_log.ts create mode 100644 plugins-extra/query_enhancements/public/query_assist/utils/index.ts create mode 100644 plugins-extra/query_enhancements/public/services.ts create mode 100644 plugins-extra/query_enhancements/server/routes/query_assist/agents.ts create mode 100644 plugins-extra/query_enhancements/server/routes/query_assist/index.ts create mode 100644 plugins-extra/query_enhancements/server/routes/query_assist/ppl/create_response.ts create mode 100644 plugins-extra/query_enhancements/server/routes/query_assist/routes.ts diff --git a/plugins-extra/query_enhancements/common/query_assist/index.ts b/plugins-extra/query_enhancements/common/query_assist/index.ts new file mode 100644 index 000000000000..d4173b12f8c0 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/opensearch_dashboards.json b/plugins-extra/query_enhancements/opensearch_dashboards.json index cdd81cde48c6..a13261e6f45c 100644 --- a/plugins-extra/query_enhancements/opensearch_dashboards.json +++ b/plugins-extra/query_enhancements/opensearch_dashboards.json @@ -6,5 +6,5 @@ "ui": true, "requiredPlugins": ["data"], "optionalPlugins": ["home"], - "requiredBundles": [] + "requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"] } diff --git a/plugins-extra/query_enhancements/public/assets/query_assist_logo.svg b/plugins-extra/query_enhancements/public/assets/query_assist_logo.svg new file mode 100644 index 000000000000..b744e8c35e8f --- /dev/null +++ b/plugins-extra/query_enhancements/public/assets/query_assist_logo.svg @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + + diff --git a/plugins-extra/query_enhancements/public/plugin.ts b/plugins-extra/query_enhancements/public/plugin.tsx similarity index 78% rename from plugins-extra/query_enhancements/public/plugin.ts rename to plugins-extra/query_enhancements/public/plugin.tsx index 6f926e66fec2..256142023c38 100644 --- a/plugins-extra/query_enhancements/public/plugin.ts +++ b/plugins-extra/query_enhancements/public/plugin.tsx @@ -1,20 +1,29 @@ 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 { PPLQlSearchInterceptor } from './search/ppl_search_interceptor'; +import { SQLQlSearchInterceptor } from './search/sql_search_interceptor'; +import { setData, setStorage } from './services'; import { QueryEnhancementsPluginSetup, - QueryEnhancementsPluginStart, QueryEnhancementsPluginSetupDependencies, + QueryEnhancementsPluginStart, + QueryEnhancementsPluginStartDependencies, } from './types'; -import { PPLQlSearchInterceptor } from './search/ppl_search_interceptor'; -import { SQLQlSearchInterceptor } from './search/sql_search_interceptor'; export class QueryEnhancementsPlugin implements Plugin { + private readonly storage: IStorageWrapper; + + constructor() { + this.storage = new Storage(window.localStorage); + } + public setup( core: CoreSetup, { data }: QueryEnhancementsPluginSetupDependencies ): QueryEnhancementsPluginSetup { - const pplSearchInterceptor = new PPLQlSearchInterceptor({ toasts: core.notifications.toasts, http: core.http, @@ -43,6 +52,7 @@ export class QueryEnhancementsPlugin initialTo: moment().add(2, 'days').toISOString(), }, showFilterBar: false, + extensions: [createQueryAssistExtension(core.http)], }, fields: { visualizable: false, @@ -75,7 +85,12 @@ export class QueryEnhancementsPlugin return {}; } - public start(core: CoreStart): QueryEnhancementsPluginStart { + public start( + core: CoreStart, + deps: QueryEnhancementsPluginStartDependencies + ): QueryEnhancementsPluginStart { + setStorage(this.storage); + setData(deps.data); return {}; } diff --git a/plugins-extra/query_enhancements/public/query_assist/components/call_outs.tsx b/plugins-extra/query_enhancements/public/query_assist/components/call_outs.tsx new file mode 100644 index 000000000000..f98569f42572 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/components/index.ts b/plugins-extra/query_enhancements/public/query_assist/components/index.ts new file mode 100644 index 000000000000..d6276e402dee --- /dev/null +++ b/plugins-extra/query_enhancements/public/query_assist/components/index.ts @@ -0,0 +1 @@ +export { QueryAssistBar } from './query_assist_bar'; diff --git a/plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx b/plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx new file mode 100644 index 000000000000..c80fe5d09377 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx b/plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx new file mode 100644 index 000000000000..735ff555372a --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/components/submit_button.tsx b/plugins-extra/query_enhancements/public/query_assist/components/submit_button.tsx new file mode 100644 index 000000000000..3b7a3c88fda8 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/hooks/index.ts b/plugins-extra/query_enhancements/public/query_assist/hooks/index.ts new file mode 100644 index 000000000000..59226067cb40 --- /dev/null +++ b/plugins-extra/query_enhancements/public/query_assist/hooks/index.ts @@ -0,0 +1 @@ +export * from './use_generate'; diff --git a/plugins-extra/query_enhancements/public/query_assist/hooks/use_generate.ts b/plugins-extra/query_enhancements/public/query_assist/hooks/use_generate.ts new file mode 100644 index 000000000000..0d23114fbd34 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/index.ts b/plugins-extra/query_enhancements/public/query_assist/index.ts new file mode 100644 index 000000000000..0c059cff8e8c --- /dev/null +++ b/plugins-extra/query_enhancements/public/query_assist/index.ts @@ -0,0 +1 @@ +export { createQueryAssistExtension } from './utils'; diff --git a/plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx b/plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx new file mode 100644 index 000000000000..f2f046d11f61 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/utils/errors.ts b/plugins-extra/query_enhancements/public/query_assist/utils/errors.ts new file mode 100644 index 000000000000..630a3759904c --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/utils/get_persisted_log.ts b/plugins-extra/query_enhancements/public/query_assist/utils/get_persisted_log.ts new file mode 100644 index 000000000000..185acdc18df1 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/query_assist/utils/index.ts b/plugins-extra/query_enhancements/public/query_assist/utils/index.ts new file mode 100644 index 000000000000..5384401204c8 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/public/services.ts b/plugins-extra/query_enhancements/public/services.ts new file mode 100644 index 000000000000..7132f4a05698 --- /dev/null +++ b/plugins-extra/query_enhancements/public/services.ts @@ -0,0 +1,6 @@ +import { createGetterSetter } from '../../../src/plugins/opensearch_dashboards_utils/common'; +import { IStorageWrapper } from '../../../src/plugins/opensearch_dashboards_utils/public'; +import { DataPublicPluginStart } from '../../../src/plugins/data/public'; + +export const [getStorage, setStorage] = createGetterSetter('storage'); +export const [getData, setData] = createGetterSetter('data'); diff --git a/plugins-extra/query_enhancements/server/routes/index.ts b/plugins-extra/query_enhancements/server/routes/index.ts index 2183bf38d0a9..7bb2b3ac8ce8 100644 --- a/plugins-extra/query_enhancements/server/routes/index.ts +++ b/plugins-extra/query_enhancements/server/routes/index.ts @@ -5,11 +5,12 @@ import { Logger, ResponseError, } from '../../../../src/core/server'; -import { ISearchStrategy } from '../../../../src/plugins/data/server'; import { IDataFrameResponse, IOpenSearchDashboardsSearchRequest, } from '../../../../src/plugins/data/common'; +import { ISearchStrategy } from '../../../../src/plugins/data/server'; +import { registerQueryAssistRoutes } from './query_assist'; export function defineRoutes( logger: Logger, @@ -19,6 +20,7 @@ export function defineRoutes( ISearchStrategy > ) { + registerQueryAssistRoutes(logger, router); router.post( { path: `/api/pplql/search`, diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts b/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts new file mode 100644 index 000000000000..676ea01bab5c --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/server/routes/query_assist/index.ts b/plugins-extra/query_enhancements/server/routes/query_assist/index.ts new file mode 100644 index 000000000000..0e1cdc303df5 --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/server/routes/query_assist/ppl/create_response.ts b/plugins-extra/query_enhancements/server/routes/query_assist/ppl/create_response.ts new file mode 100644 index 000000000000..ac651a0a45fc --- /dev/null +++ b/plugins-extra/query_enhancements/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/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts b/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts new file mode 100644 index 000000000000..0b8c5c893239 --- /dev/null +++ b/plugins-extra/query_enhancements/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 }); + } + } + ); +} diff --git a/src/plugins/data/public/index.ts b/src/plugins/data/public/index.ts index 3cd7c54410f4..6583fa4b26f6 100644 --- a/src/plugins/data/public/index.ts +++ b/src/plugins/data/public/index.ts @@ -528,5 +528,4 @@ export { DataSourceOption, } from './data_sources/datasource_selector'; -export { SuggestionsComponent } from './ui'; export { PersistedLog } from './query'; diff --git a/src/plugins/data/public/ui/query_editor/query_editor.tsx b/src/plugins/data/public/ui/query_editor/query_editor.tsx index 02c495c430f1..88211841f335 100644 --- a/src/plugins/data/public/ui/query_editor/query_editor.tsx +++ b/src/plugins/data/public/ui/query_editor/query_editor.tsx @@ -54,7 +54,7 @@ export interface QueryEditorProps { size?: SuggestionsListSize; className?: string; isInvalid?: boolean; - queryEditorRef: React.RefObject; + queryEditorHeaderRef: React.RefObject; } interface Props extends QueryEditorProps { @@ -519,7 +519,7 @@ export default class QueryEditorUI extends Component { -
+
; } // Needed for React.lazy @@ -66,6 +66,7 @@ export interface QueryEditorTopRowProps { export default function QueryEditorTopRow(props: QueryEditorTopRowProps) { const [isDateRangeInvalid, setIsDateRangeInvalid] = useState(false); const [isQueryEditorFocused, setIsQueryEditorFocused] = useState(false); + const queryEditorHeaderRef = useRef(null); const opensearchDashboards = useOpenSearchDashboards(); const { uiSettings, storage, appName } = opensearchDashboards.services; @@ -239,12 +240,23 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) { getQueryStringInitialValue={getQueryStringInitialValue} persistedLog={persistedLog} dataTestSubj={props.dataTestSubj} - queryEditorRef={props.queryEditorRef} + queryEditorHeaderRef={queryEditorHeaderRef} /> ); } + function renderSearchBarExtensions() { + if (!shouldRenderSearchBarExtensions() || !queryEditorHeaderRef.current) return; + return ( + + ); + } + function renderSharingMetaFields() { const { from, to } = getDateRange(); const dateRangePretty = prettyDuration( @@ -274,6 +286,12 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) { return Boolean(props.showQueryEditor && props.indexPatterns && props.query && storage); } + function shouldRenderSearchBarExtensions(): boolean { + return Boolean( + queryLanguage && props.queryEnhancements?.get(queryLanguage)?.searchBar?.extensions?.length + ); + } + function renderUpdateButton() { const button = props.customSubmitButton ? ( React.cloneElement(props.customSubmitButton, { onClick: onClickSubmitButton }) @@ -366,6 +384,7 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) { direction="column" justifyContent="flexEnd" > + {renderSearchBarExtensions()} {renderQueryEditor()} diff --git a/src/plugins/data/public/ui/search_bar/search_bar.tsx b/src/plugins/data/public/ui/search_bar/search_bar.tsx index 626578a65e8e..dd3587acfc96 100644 --- a/src/plugins/data/public/ui/search_bar/search_bar.tsx +++ b/src/plugins/data/public/ui/search_bar/search_bar.tsx @@ -45,7 +45,6 @@ import QueryEditorTopRow from '../query_editor/query_editor_top_row'; import QueryBarTopRow from '../query_string_input/query_bar_top_row'; import { SavedQueryMeta, SaveQueryForm } from '../saved_query_form'; import { SavedQueryManagementComponent } from '../saved_query_management'; -import { SearchBarExtensions } from '../search_bar_extensions'; import { QueryEnhancement, Settings } from '../types'; interface SearchBarInjectedDeps { @@ -123,12 +122,6 @@ class SearchBarUI extends Component { private services = this.props.opensearchDashboards.services; private savedQueryService = this.services.data.query.savedQueries; - /** - * queryEditorRef can't be bound to the actual editor - * https://github.com/react-monaco-editor/react-monaco-editor/blob/v0.27.0/src/editor.js#L113, - * currently it is an element above. - */ - public queryEditorRef = React.createRef(); public filterBarRef: Element | null = null; public filterBarWrapperRef: Element | null = null; @@ -243,15 +236,6 @@ class SearchBarUI extends Component { ); } - private shouldRenderExtensions() { - return ( - this.props.isEnhancementsEnabled && - (!!this.props.queryEnhancements?.get(this.state.query?.language!)?.searchBar?.extensions - ?.length ?? - false) - ); - } - /* * This Function is here to show the toggle in saved query form * in case you the date range (from/to) @@ -525,20 +509,6 @@ class SearchBarUI extends Component { filterBar={filterBar} dataTestSubj={this.props.dataTestSubj} indicateNoData={this.props.indicateNoData} - queryEditorRef={this.queryEditorRef} - /> - ); - } - - let searchBarExtensions; - if (this.shouldRenderExtensions() && this.queryEditorRef.current) { - searchBarExtensions = ( - ); } @@ -548,7 +518,6 @@ class SearchBarUI extends Component { return (
{queryBar} - {searchBarExtensions} {queryEditor} {!!!this.props.isEnhancementsEnabled && filterBar} diff --git a/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx b/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx index e22ed9220dd3..88a3fcdfbb08 100644 --- a/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx +++ b/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx @@ -18,7 +18,7 @@ export interface SearchBarExtensionDependencies { /** * Currently selected index patterns. */ - indexPatterns?: IIndexPattern[]; + indexPatterns?: Array; } export interface SearchBarExtensionConfig { diff --git a/src/plugins/data/public/ui/types.ts b/src/plugins/data/public/ui/types.ts index f03cd85169ac..c07af4db757a 100644 --- a/src/plugins/data/public/ui/types.ts +++ b/src/plugins/data/public/ui/types.ts @@ -8,7 +8,8 @@ import { SearchInterceptor } from '../search'; import { IndexPatternSelectProps } from './index_pattern_select'; import { StatefulSearchBarProps } from './search_bar'; import { Settings } from './settings'; -import { SearchBarExtensionConfig } from './search_bar_extensions/search_bar_extension'; +import { SearchBarExtensionConfig } from './search_bar_extensions'; +import { SuggestionsComponentProps } from './typeahead/suggestions_component'; export * from './settings'; @@ -64,6 +65,7 @@ export interface IUiStart { queryEnhancements: Map; IndexPatternSelect: React.ComponentType; SearchBar: React.ComponentType; + SuggestionsComponent: React.ComponentType; Settings: Settings; containerRef: HTMLDivElement | null; container$: BehaviorSubject; diff --git a/src/plugins/data/public/ui/ui_service.ts b/src/plugins/data/public/ui/ui_service.ts index 99be6b1d2862..b24fc3ee21d0 100644 --- a/src/plugins/data/public/ui/ui_service.ts +++ b/src/plugins/data/public/ui/ui_service.ts @@ -13,6 +13,7 @@ import { createSearchBar } from './search_bar/create_search_bar'; import { createSettings } from './settings'; import { DataPublicPluginStart } from '../types'; import { IStorageWrapper } from '../../../opensearch_dashboards_utils/public'; +import { SuggestionsComponent } from './typeahead'; /** @internal */ // eslint-disable-next-line @typescript-eslint/no-empty-interface @@ -71,6 +72,7 @@ export class UiService implements Plugin { queryEnhancements: this.queryEnhancements, IndexPatternSelect: createIndexPatternSelect(core.savedObjects.client), SearchBar, + SuggestionsComponent, Settings, containerRef: this.containerRef, container$: this.container$,