From 7e49fa61164d0dd3e080d88d479387f9f4790650 Mon Sep 17 00:00:00 2001 From: Yulong Ruan Date: Thu, 6 Jun 2024 16:01:16 +0800 Subject: [PATCH] add source selector for vega visualization Signed-off-by: Yulong Ruan --- .../public/components/vega_vis_editor.tsx | 9 +- .../vis_type_vega/public/text_to_vega.ts | 20 ++-- .../components/source_selector.tsx | 99 +++++++++++++++++++ .../components/visualize_top_nav.tsx | 16 ++- 4 files changed, 129 insertions(+), 15 deletions(-) create mode 100644 src/plugins/visualize/public/application/components/source_selector.tsx diff --git a/src/plugins/vis_type_vega/public/components/vega_vis_editor.tsx b/src/plugins/vis_type_vega/public/components/vega_vis_editor.tsx index 73abdd4788bc..324e7aacd902 100644 --- a/src/plugins/vis_type_vega/public/components/vega_vis_editor.tsx +++ b/src/plugins/vis_type_vega/public/components/vega_vis_editor.tsx @@ -28,7 +28,7 @@ * under the License. */ -import React, { useCallback, useEffect } from 'react'; +import React, { useCallback, useEffect, useRef } from 'react'; import { EuiCodeEditor } from '@elastic/eui'; import compactStringify from 'json-stringify-pretty-compact'; import hjson from 'hjson'; @@ -76,6 +76,9 @@ function format( } function VegaVisEditor({ stateParams, setValue }: VisOptionsProps) { + const setValueRef = useRef(setValue); + setValueRef.current = setValue; + const onChange = useCallback( (value: string) => { setValue('spec', value); @@ -104,7 +107,7 @@ function VegaVisEditor({ stateParams, setValue }: VisOptionsProps) { }), }); } else { - setValue('spec', JSON.stringify(result, null, 4)); + setValueRef.current('spec', JSON.stringify(result, null, 4)); } } }); @@ -112,7 +115,7 @@ function VegaVisEditor({ stateParams, setValue }: VisOptionsProps) { return () => { subscription.unsubscribe(); }; - }, [setValue]); + }, []); return (
diff --git a/src/plugins/vis_type_vega/public/text_to_vega.ts b/src/plugins/vis_type_vega/public/text_to_vega.ts index 6718a283f521..4944af04555f 100644 --- a/src/plugins/vis_type_vega/public/text_to_vega.ts +++ b/src/plugins/vis_type_vega/public/text_to_vega.ts @@ -29,7 +29,7 @@ Just reply with the json based Vega-Lite object, do not include any other conten }; export class Text2Vega { - input$ = new BehaviorSubject({ input: '' }); + input$ = new BehaviorSubject({ prompt: '', index: '' }); result$: Observable | { error: any }>; status$ = new BehaviorSubject<'RUNNING' | 'STOPPED'>('STOPPED'); http: HttpSetup; @@ -38,19 +38,19 @@ export class Text2Vega { this.http = http; this.result$ = this.input$ .pipe( - filter((v) => v.input.length > 0), + filter((v) => v.prompt.length > 0), debounceTime(200), tap(() => this.status$.next('RUNNING')) ) .pipe( switchMap((v) => - of(v.input).pipe( + of(v).pipe( // text to ppl switchMap(async (value) => { - const pplQuestion = value.split('//')[0]; - const ppl = await this.text2ppl(pplQuestion); + const pplQuestion = value.prompt.split('//')[0]; + const ppl = await this.text2ppl(pplQuestion, value.index); return { - input: value, + ...value, ppl, }; }), @@ -64,7 +64,7 @@ export class Text2Vega { }), // call llm to generate vega switchMap(async (value) => { - const prompt = createPrompt(value.input, value.ppl, value.sample); + const prompt = createPrompt(value.prompt, value.ppl, value.sample); const result = await this.text2vega(prompt); result.data = { url: { @@ -89,18 +89,18 @@ export class Text2Vega { return result; } - async text2ppl(query: string) { + async text2ppl(query: string, index: string) { const pplResponse = await this.http.post('/api/llm/text2ppl', { body: JSON.stringify({ question: query, - index: 'opensearch_dashboards_sample_data_logs', + index: index, }), }); const result = JSON.parse(pplResponse.body.inference_results[0].output[0].result); return result.ppl; } - invoke(value: { input: string }) { + invoke(value: { prompt: string; index: string }) { this.input$.next(value); } diff --git a/src/plugins/visualize/public/application/components/source_selector.tsx b/src/plugins/visualize/public/application/components/source_selector.tsx new file mode 100644 index 000000000000..673c804283a7 --- /dev/null +++ b/src/plugins/visualize/public/application/components/source_selector.tsx @@ -0,0 +1,99 @@ +import React, { useCallback, useMemo, useState, useEffect } from 'react'; +import { i18n } from '@osd/i18n'; + +import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public'; +import { + DataSource, + DataSourceGroup, + DataSourceSelectable, + DataSourceOption, +} from '../../../../data/public'; +import { VisualizeServices } from '../types'; + +export const SourceSelector = ({ + selectedSourceId, + onChange, +}: { + selectedSourceId: string; + onChange: (ds: DataSourceOption) => void; +}) => { + const { + services: { + data: { dataSources }, + notifications: { toasts }, + }, + } = useOpenSearchDashboards(); + const [currentDataSources, setCurrentDataSources] = useState([]); + const [dataSourceOptions, setDataSourceOptions] = useState([]); + + const selectedSources = useMemo(() => { + if (selectedSourceId) { + for (const group of dataSourceOptions) { + for (const item of group.options) { + if (item.value === selectedSourceId) { + return [item]; + } + } + } + } + return []; + }, [selectedSourceId, dataSourceOptions]); + + useEffect(() => { + if ( + !selectedSourceId && + dataSourceOptions.length > 0 && + dataSourceOptions[0].options.length > 0 + ) { + onChange(dataSourceOptions[0].options[0]); + } + }, [selectedSourceId, dataSourceOptions]); + + useEffect(() => { + const subscription = dataSources.dataSourceService + .getDataSources$() + .subscribe((currentDataSources) => { + setCurrentDataSources(Object.values(currentDataSources)); + }); + + return () => { + subscription.unsubscribe(); + }; + }, [dataSources]); + + const onDataSourceSelect = useCallback( + (selectedDataSources: DataSourceOption[]) => { + onChange(selectedDataSources[0]); + }, + [onChange] + ); + + const handleGetDataSetError = useCallback( + () => (error: Error) => { + toasts.addError(error, { + title: + i18n.translate('visualize.vega.failedToGetDataSetErrorDescription', { + defaultMessage: 'Failed to get data set: ', + }) + (error.message || error.name), + }); + }, + [toasts] + ); + + const memorizedReload = useCallback(() => { + dataSources.dataSourceService.reload(); + }, [dataSources.dataSourceService]); + + return ( + + ); +}; diff --git a/src/plugins/visualize/public/application/components/visualize_top_nav.tsx b/src/plugins/visualize/public/application/components/visualize_top_nav.tsx index 9a0611c6f1a9..25d940e1b848 100644 --- a/src/plugins/visualize/public/application/components/visualize_top_nav.tsx +++ b/src/plugins/visualize/public/application/components/visualize_top_nav.tsx @@ -51,8 +51,9 @@ import { } from '../types'; import { APP_NAME } from '../visualize_constants'; import { getTopNavConfig } from '../utils'; -import type { IndexPattern } from '../../../../data/public'; +import type { IndexPattern, DataSourceOption } from '../../../../data/public'; import chatLogo from './query_assistant_logo.svg'; +import { SourceSelector } from './source_selector'; interface VisualizeTopNavProps { currentAppState: VisualizeAppState; @@ -145,6 +146,7 @@ const TopNav = ({ const [indexPatterns, setIndexPatterns] = useState( vis.data.indexPattern ? [vis.data.indexPattern] : [] ); + const [selectedSource, setSelectedSource] = useState(); const showDatePicker = () => { // tsvb loads without an indexPattern initially (TODO investigate). // hide timefilter only if timeFieldName is explicitly undefined. @@ -254,7 +256,11 @@ const TopNav = ({ const indexName = 'opensearch_dashboards_sample_data_logs'; const onGenerate = async () => { - (window as any).llm.text2vega.invoke({ input: value }); + if (selectedSource) { + (window as any).llm.text2vega.invoke({ prompt: value, index: selectedSource.label }); + } else { + services.notifications.toasts.addWarning('Please select a index'); + } }; return isChromeVisible ? ( @@ -273,6 +279,12 @@ const TopNav = ({ alignItems="center" style={{ maxHeight: '100px' }} > + + setSelectedSource(ds)} + /> +