Skip to content

Commit

Permalink
add source selector for vega visualization
Browse files Browse the repository at this point in the history
Signed-off-by: Yulong Ruan <ruanyl@amazon.com>
  • Loading branch information
ruanyl committed Jun 6, 2024
1 parent 7a829b8 commit 7e49fa6
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -76,6 +76,9 @@ function format(
}

function VegaVisEditor({ stateParams, setValue }: VisOptionsProps<VisParams>) {
const setValueRef = useRef(setValue);
setValueRef.current = setValue;

const onChange = useCallback(
(value: string) => {
setValue('spec', value);
Expand Down Expand Up @@ -104,15 +107,15 @@ function VegaVisEditor({ stateParams, setValue }: VisOptionsProps<VisParams>) {
}),
});
} else {
setValue('spec', JSON.stringify(result, null, 4));
setValueRef.current('spec', JSON.stringify(result, null, 4));
}
}
});

return () => {
subscription.unsubscribe();
};
}, [setValue]);
}, []);

return (
<div className="vgaEditor">
Expand Down
20 changes: 10 additions & 10 deletions src/plugins/vis_type_vega/public/text_to_vega.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<Record<string, any> | { error: any }>;
status$ = new BehaviorSubject<'RUNNING' | 'STOPPED'>('STOPPED');
http: HttpSetup;
Expand All @@ -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,
};
}),
Expand All @@ -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: {
Expand All @@ -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,

Check failure on line 96 in src/plugins/vis_type_vega/public/text_to_vega.ts

View workflow job for this annotation

GitHub Actions / Build and Verify on Linux (ciGroup1)

Expected property shorthand
}),
});
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);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
import React, { useCallback, useMemo, useState, useEffect } from 'react';

Check failure on line 1 in src/plugins/visualize/public/application/components/source_selector.tsx

View workflow job for this annotation

GitHub Actions / Build and Verify on Linux (ciGroup1)

File must start with a license header
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<VisualizeServices>();
const [currentDataSources, setCurrentDataSources] = useState<DataSource[]>([]);
const [dataSourceOptions, setDataSourceOptions] = useState<DataSourceGroup[]>([]);

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]);

Check failure on line 50 in src/plugins/visualize/public/application/components/source_selector.tsx

View workflow job for this annotation

GitHub Actions / Build and Verify on Linux (ciGroup1)

React Hook useEffect has a missing dependency: 'onChange'. Either include it or remove the dependency array. If 'onChange' changes too often, find the parent component that defines it and wrap that definition in useCallback

useEffect(() => {
const subscription = dataSources.dataSourceService
.getDataSources$()
.subscribe((currentDataSources) => {

Check failure on line 55 in src/plugins/visualize/public/application/components/source_selector.tsx

View workflow job for this annotation

GitHub Actions / Build and Verify on Linux (ciGroup1)

'currentDataSources' is already declared in the upper scope
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 (
<DataSourceSelectable
dataSources={currentDataSources}
dataSourceOptionList={dataSourceOptions}
setDataSourceOptionList={setDataSourceOptions}
onDataSourceSelect={onDataSourceSelect}
selectedSources={selectedSources}
onGetDataSetError={handleGetDataSetError}
onRefresh={memorizedReload}
fullWidth
/>
);
};
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -145,6 +146,7 @@ const TopNav = ({
const [indexPatterns, setIndexPatterns] = useState<IndexPattern[]>(
vis.data.indexPattern ? [vis.data.indexPattern] : []
);
const [selectedSource, setSelectedSource] = useState<DataSourceOption>();
const showDatePicker = () => {
// tsvb loads without an indexPattern initially (TODO investigate).
// hide timefilter only if timeFieldName is explicitly undefined.
Expand Down Expand Up @@ -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 ? (
Expand All @@ -273,6 +279,12 @@ const TopNav = ({
alignItems="center"
style={{ maxHeight: '100px' }}
>
<EuiFlexItem grow={3}>
<SourceSelector
selectedSourceId={selectedSource?.value ?? ''}
onChange={(ds) => setSelectedSource(ds)}
/>
</EuiFlexItem>
<EuiFlexItem grow={9}>
<EuiInputPopover
input={
Expand Down

0 comments on commit 7e49fa6

Please sign in to comment.