feat: text to visualization (#218)

* feat: add text to visualization

Signed-off-by: Yulong Ruan <[email protected]>

---------

Signed-off-by: Yulong Ruan <[email protected]>
(cherry picked from commit 63d9d80)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

common/constants/llm.ts

@@ -19,6 +19,11 @@ export const ASSISTANT_API = {
   ACCOUNT: `${API_BASE}/account`,
 } as const;
 
+export const TEXT2VIZ_API = {
+  TEXT2PPL: `${API_BASE}/text2ppl`,
+  TEXT2VEGA: `${API_BASE}/text2vega`,
+};
+
 export const NOTEBOOK_API = {
   CREATE_NOTEBOOK: `${NOTEBOOK_PREFIX}/note`,
   SET_PARAGRAPH: `${NOTEBOOK_PREFIX}/set_paragraphs/`,

opensearch_dashboards.json

@@ -5,16 +5,19 @@
   "server": true,
   "ui": true,
   "requiredPlugins": [
+    "data",
     "dashboard",
     "embeddable",
     "opensearchDashboardsReact",
-    "opensearchDashboardsUtils"
+    "opensearchDashboardsUtils",
+    "visualizations"
   ],
   "optionalPlugins": [
     "dataSource",
     "dataSourceManagement"
   ],
+  "requiredBundles": [],
   "configPath": [
     "assistant"
   ]
-}
+}

public/components/visualization/source_selector.tsx

@@ -0,0 +1,102 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React, { useCallback, useMemo, useState, useEffect } from 'react';
+import { i18n } from '@osd/i18n';
+
+import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
+import {
+  DataSource,
+  DataSourceGroup,
+  DataSourceSelectable,
+  DataSourceOption,
+} from '../../../../../src/plugins/data/public';
+import { StartServices } from '../../types';
+
+export const SourceSelector = ({
+  selectedSourceId,
+  onChange,
+}: {
+  selectedSourceId: string;
+  onChange: (ds: DataSourceOption) => void;
+}) => {
+  const {
+    services: {
+      data: { dataSources },
+      notifications: { toasts },
+    },
+  } = useOpenSearchDashboards<StartServices>();
+  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]);
+
+  useEffect(() => {
+    const subscription = dataSources.dataSourceService.getDataSources$().subscribe((ds) => {
+      setCurrentDataSources(Object.values(ds));
+    });
+
+    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
+    />
+  );
+};

public/components/visualization/text2vega.ts

@@ -0,0 +1,163 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BehaviorSubject, Observable, of } from 'rxjs';
+import { debounceTime, switchMap, tap, filter, catchError } from 'rxjs/operators';
+import { TEXT2VIZ_API } from '.../../../common/constants/llm';
+import { HttpSetup } from '../../../../../src/core/public';
+import { DataPublicPluginStart } from '../../../../../src/plugins/data/public';
+
+const DATA_SOURCE_DELIMITER = '::';
+
+const topN = (ppl: string, n: number) => `${ppl} | head ${n}`;
+
+const getDataSourceAndIndexFromLabel = (label: string) => {
+  if (label.includes(DATA_SOURCE_DELIMITER)) {
+    return [
+      label.slice(0, label.indexOf(DATA_SOURCE_DELIMITER)),
+      label.slice(label.indexOf(DATA_SOURCE_DELIMITER) + DATA_SOURCE_DELIMITER.length),
+    ] as const;
+  }
+  return [, label] as const;
+};
+
+interface Input {
+  prompt: string;
+  index: string;
+  dataSourceId?: string;
+}
+
+export class Text2Vega {
+  input$ = new BehaviorSubject<Input>({ prompt: '', index: '' });
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  result$: Observable<Record<string, any> | { error: any }>;
+  status$ = new BehaviorSubject<'RUNNING' | 'STOPPED'>('STOPPED');
+  http: HttpSetup;
+  searchClient: DataPublicPluginStart['search'];
+
+  constructor(http: HttpSetup, searchClient: DataPublicPluginStart['search']) {
+    this.http = http;
+    this.searchClient = searchClient;
+    this.result$ = this.input$
+      .pipe(
+        filter((v) => v.prompt.length > 0),
+        debounceTime(200),
+        tap(() => this.status$.next('RUNNING'))
+      )
+      .pipe(
+        switchMap((v) =>
+          of(v).pipe(
+            // text to ppl
+            switchMap(async (value) => {
+              const [, indexName] = getDataSourceAndIndexFromLabel(value.index);
+              const pplQuestion = value.prompt.split('//')[0];
+              const ppl = await this.text2ppl(pplQuestion, indexName, value.dataSourceId);
+              return {
+                ...value,
+                ppl,
+              };
+            }),
+            // query sample data with ppl
+            switchMap(async (value) => {
+              const ppl = topN(value.ppl, 2);
+              const res = await this.searchClient
+                .search(
+                  { params: { body: { query: ppl } }, dataSourceId: value.dataSourceId },
+                  { strategy: 'pplraw' }
+                )
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                .toPromise<any>();
+              return { ...value, sample: res.rawResponse };
+            }),
+            // call llm to generate vega
+            switchMap(async (value) => {
+              const result = await this.text2vega({
+                input: value.prompt,
+                ppl: value.ppl,
+                sampleData: JSON.stringify(value.sample.jsonData),
+                dataSchema: JSON.stringify(value.sample.schema),
+                dataSourceId: value.dataSourceId,
+              });
+              const [dataSourceName] = getDataSourceAndIndexFromLabel(value.index);
+              result.data = {
+                url: {
+                  '%type%': 'ppl',
+                  body: { query: value.ppl },
+                  data_source_name: dataSourceName,
+                },
+              };
+              return result;
+            }),
+            catchError((e) => of({ error: e }))
+          )
+        )
+      )
+      .pipe(tap(() => this.status$.next('STOPPED')));
+  }
+
+  async text2vega({
+    input,
+    ppl,
+    sampleData,
+    dataSchema,
+    dataSourceId,
+  }: {
+    input: string;
+    ppl: string;
+    sampleData: string;
+    dataSchema: string;
+    dataSourceId?: string;
+  }) {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const escapeField = (json: any, field: string) => {
+      if (json[field]) {
+        if (typeof json[field] === 'string') {
+          json[field] = json[field].replace(/\./g, '\\.');
+        }
+        if (typeof json[field] === 'object') {
+          Object.keys(json[field]).forEach((p) => {
+            escapeField(json[field], p);
+          });
+        }
+      }
+    };
+    const res = await this.http.post(TEXT2VIZ_API.TEXT2VEGA, {
+      body: JSON.stringify({
+        input,
+        ppl,
+        sampleData: JSON.stringify(sampleData),
+        dataSchema: JSON.stringify(dataSchema),
+      }),
+      query: { dataSourceId },
+    });
+
+    // need to escape field: geo.city -> field: geo\\.city
+    escapeField(res, 'encoding');
+    return res;
+  }
+
+  async text2ppl(query: string, index: string, dataSourceId?: string) {
+    const pplResponse = await this.http.post(TEXT2VIZ_API.TEXT2PPL, {
+      body: JSON.stringify({
+        question: query,
+        index,
+      }),
+      query: { dataSourceId },
+    });
+    return pplResponse.ppl;
+  }
+
+  invoke(value: Input) {
+    this.input$.next(value);
+  }
+
+  getStatus$() {
+    return this.status$;
+  }
+
+  getResult$() {
+    return this.result$;
+  }
+}

public/components/visualization/text2viz.scss

@@ -0,0 +1,10 @@
+.text2viz__page {
+  .visualize {
+    height: 400px;
+  }
+
+  .text2viz__right {
+    padding-top: 15px;
+    padding-left: 30px;
+  }
+}