Edit viz in a dialog (#351)

* feat: supporting edit visualization with input from a dialog Signed-off-by: Yulong Ruan <ruanyl@amazon.com> * fix: save instruction input Signed-off-by: Yulong Ruan <ruanyl@amazon.com> * fix: change internal error to bad request Signed-off-by: Yulong Ruan <ruanyl@amazon.com> * fix: use sass variable for 12px font size Signed-off-by: Yulong Ruan <ruanyl@amazon.com> * add CHANGELOG Signed-off-by: Yulong Ruan <ruanyl@amazon.com> --------- Signed-off-by: Yulong Ruan <ruanyl@amazon.com>
opensearch-project · Oct 21, 2024 · 6a49d78 · 6a49d78
1 parent e3ec245
commit 6a49d78
Show file tree

Hide file tree

Showing 7 changed files with 254 additions and 58 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - feat: only display ai actions that compatible with the datasource([#350](https://github.com/opensearch-project/dashboards-assistant/pull/350))
 - feat: take index pattern and query assistant input to text2viz app([#349](https://github.com/opensearch-project/dashboards-assistant/pull/349))
 - feat: Hide incompatible index patterns ([#354] (https://github.com/opensearch-project/dashboards-assistant/pull/354))
+- feat: edit visualization with natural language in a dialog([#351](https://github.com/opensearch-project/dashboards-assistant/pull/351))
 
 
 ### 📈 Features/Enhancements

diff --git a/public/components/visualization/text2vega.ts b/public/components/visualization/text2vega.ts
@@ -13,13 +13,14 @@ import { DataSourceAttributes } from '../../../../../src/plugins/data_source/com
 const topN = (ppl: string, n: number) => `${ppl} | head ${n}`;
 
 interface Input {
-  prompt: string;
+  inputQuestion: string;
+  inputInstruction?: string;
   index: string;
   dataSourceId?: string;
 }
 
 export class Text2Vega {
-  input$ = new BehaviorSubject<Input>({ prompt: '', index: '' });
+  input$ = new BehaviorSubject<Input>({ inputQuestion: '', index: '' });
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   result$: Observable<Record<string, any> | { error: any }>;
   status$ = new BehaviorSubject<'RUNNING' | 'STOPPED'>('STOPPED');
@@ -37,7 +38,7 @@ export class Text2Vega {
     this.savedObjects = savedObjects;
     this.result$ = this.input$
       .pipe(
-        filter((v) => v.prompt.length > 0),
+        filter((v) => v.inputQuestion.length > 0),
         tap(() => this.status$.next('RUNNING')),
         debounceTime(200)
       )
@@ -46,7 +47,7 @@ export class Text2Vega {
           of(v).pipe(
             // text to ppl
             switchMap(async (value) => {
-              const pplQuestion = value.prompt.split('//')[0];
+              const pplQuestion = value.inputQuestion;
               const ppl = await this.text2ppl(pplQuestion, value.index, value.dataSourceId);
               return {
                 ...value,
@@ -71,7 +72,8 @@ export class Text2Vega {
             // call llm to generate vega
             switchMap(async (value) => {
               const result = await this.text2vega({
-                input: value.prompt,
+                inputQuestion: value.inputQuestion,
+                inputInstruction: value.inputInstruction,
                 ppl: value.ppl,
                 sampleData: JSON.stringify(value.sample.jsonData),
                 dataSchema: JSON.stringify(value.sample.schema),
@@ -96,13 +98,15 @@ export class Text2Vega {
   }
 
   async text2vega({
-    input,
+    inputQuestion,
+    inputInstruction = '',
     ppl,
     sampleData,
     dataSchema,
     dataSourceId,
   }: {
-    input: string;
+    inputQuestion: string;
+    inputInstruction?: string;
     ppl: string;
     sampleData: string;
     dataSchema: string;
@@ -122,7 +126,6 @@ export class Text2Vega {
         }
       }
     };
-    const [inputQuestion, inputInstruction = ''] = input.split('//');
     const res = await this.http.post(TEXT2VIZ_API.TEXT2VEGA, {
       body: JSON.stringify({
         input_question: inputQuestion.trim(),

diff --git a/public/components/visualization/text2viz.scss b/public/components/visualization/text2viz.scss
@@ -17,10 +17,26 @@
     padding-left: 30px;
   }
 
-  .feedback_thumbs {
+  .text2viz__actionContainer {
     position: absolute;
-    right: 16px;
     top: 4px;
+    right: 16px;
     z-index: 9999;
+
+    // No existing button from OUI with the same style, have to customize here
+    .vizStyleEditor__editButton {
+      height: 22px;
+      padding: 2px;
+      font-size: $ouiFontSizeXS;
+    }
+
+    .text2viz__feedbackContainer {
+      padding-right: $euiSizeS;
+      border-right: 1px solid $euiColorLightShade
+    }
+
+    .text2viz__vizStyleEditorContainer {
+      padding-left: $euiSizeS;
+    }
   }
 }
diff --git a/public/components/visualization/text2viz.tsx b/public/components/visualization/text2viz.tsx
@@ -51,6 +51,7 @@ import { VIS_NLQ_APP_ID, VIS_NLQ_SAVED_OBJECT } from '../../../common/constants/
 import { HeaderVariant } from '../../../../../src/core/public';
 import { TEXT2VEGA_INPUT_SIZE_LIMIT } from '../../../common/constants/llm';
 import { FeedbackThumbs } from '../feedback_thumbs';
+import { VizStyleEditor } from './viz_style_editor';
 
 export const INDEX_PATTERN_URL_SEARCH_KEY = 'indexPatternId';
 export const ASSISTANT_INPUT_URL_SEARCH_KEY = 'assistantInput';
@@ -96,7 +97,10 @@ export const Text2Viz = () => {
 
   const useUpdatedUX = uiSettings.get('home:useNewHomePage');
 
-  const [input, setInput] = useState(searchParams.get(ASSISTANT_INPUT_URL_SEARCH_KEY) ?? '');
+  const [inputQuestion, setInputQuestion] = useState(
+    searchParams.get(ASSISTANT_INPUT_URL_SEARCH_KEY) ?? ''
+  );
+  const [currentInstruction, setCurrentInstruction] = useState('');
   const [editorInput, setEditorInput] = useState('');
   const text2vegaRef = useRef(new Text2Vega(http, data.search, savedObjects));
 
@@ -174,7 +178,8 @@ export const Text2Viz = () => {
             }
           }
           if (savedVis?.uiState) {
-            setInput(JSON.parse(savedVis.uiState ?? '{}').input);
+            setInputQuestion(JSON.parse(savedVis.uiState ?? '{}').input ?? '');
+            setCurrentInstruction(JSON.parse(savedVis.uiState ?? '{}').instruction ?? '');
           }
         })
         .catch(() => {
@@ -196,42 +201,48 @@ export const Text2Viz = () => {
   /**
    * Submit user's natural language input to generate visualization
    */
-  const onSubmit = useCallback(async () => {
-    if (status === 'RUNNING' || !selectedSource) return;
-
-    const [inputQuestion = '', inputInstruction = ''] = input.split('//');
-    if (
-      inputQuestion.trim().length > TEXT2VEGA_INPUT_SIZE_LIMIT ||
-      inputInstruction.trim().length > TEXT2VEGA_INPUT_SIZE_LIMIT
-    ) {
-      notifications.toasts.addDanger({
-        title: i18n.translate('dashboardAssistant.feature.text2viz.invalidInput', {
-          defaultMessage: `Input size exceed limit: {limit}. Actual size: question({inputQuestionLength}), instruction({inputInstructionLength})`,
-          values: {
-            limit: TEXT2VEGA_INPUT_SIZE_LIMIT,
-            inputQuestionLength: inputQuestion.trim().length,
-            inputInstructionLength: inputInstruction.trim().length,
-          },
-        }),
-      });
-      return;
-    }
+  const onSubmit = useCallback(
+    async (inputInstruction: string = '') => {
+      setCurrentInstruction(inputInstruction);
+
+      if (status === 'RUNNING' || !selectedSource) return;
+
+      if (
+        inputQuestion.trim().length > TEXT2VEGA_INPUT_SIZE_LIMIT ||
+        inputInstruction.trim().length > TEXT2VEGA_INPUT_SIZE_LIMIT
+      ) {
+        notifications.toasts.addDanger({
+          title: i18n.translate('dashboardAssistant.feature.text2viz.invalidInput', {
+            defaultMessage:
+              'Input size exceed limit: {limit}. Actual size: question({inputQuestionLength}), instruction({inputInstructionLength})',
+            values: {
+              limit: TEXT2VEGA_INPUT_SIZE_LIMIT,
+              inputQuestionLength: inputQuestion.trim().length,
+              inputInstructionLength: inputInstruction.trim().length,
+            },
+          }),
+        });
+        return;
+      }
 
-    setSubmitting(true);
+      setSubmitting(true);
 
-    const indexPatterns = getIndexPatterns();
-    const indexPattern = await indexPatterns.get(selectedSource);
-    currentUsedIndexPatternRef.current = indexPattern;
+      const indexPatterns = getIndexPatterns();
+      const indexPattern = await indexPatterns.get(selectedSource);
+      currentUsedIndexPatternRef.current = indexPattern;
 
-    const text2vega = text2vegaRef.current;
-    text2vega.invoke({
-      index: indexPattern.title,
-      prompt: input,
-      dataSourceId: indexPattern.dataSourceRef?.id,
-    });
+      const text2vega = text2vegaRef.current;
+      text2vega.invoke({
+        index: indexPattern.title,
+        inputQuestion,
+        inputInstruction,
+        dataSourceId: indexPattern.dataSourceRef?.id,
+      });
 
-    setSubmitting(false);
-  }, [selectedSource, input, status, notifications.toasts]);
+      setSubmitting(false);
+    },
+    [selectedSource, inputQuestion, status, notifications.toasts]
+  );
 
   /**
    * Display the save visualization dialog to persist the current generated visualization
@@ -252,7 +263,8 @@ export const Text2Viz = () => {
         },
       });
       savedVis.uiState = JSON.stringify({
-        input,
+        input: inputQuestion,
+        instruction: currentInstruction,
       });
       savedVis.searchSourceFields = { index: indexPattern };
       savedVis.title = onSaveProps.newTitle;
@@ -311,7 +323,16 @@ export const Text2Viz = () => {
         />
       )
     );
-  }, [notifications, vegaSpec, input, overlays, selectedSource, savedObjectId, usageCollection]);
+  }, [
+    notifications,
+    vegaSpec,
+    inputQuestion,
+    overlays,
+    selectedSource,
+    savedObjectId,
+    usageCollection,
+    currentInstruction,
+  ]);
 
   const pageTitle = savedObjectId
     ? i18n.translate('dashboardAssistant.feature.text2viz.breadcrumbs.editVisualization', {
@@ -380,8 +401,8 @@ export const Text2Viz = () => {
         </EuiFlexItem>
         <EuiFlexItem grow={8}>
           <EuiFieldText
-            value={input}
-            onChange={(e) => setInput(e.target.value)}
+            value={inputQuestion}
+            onChange={(e) => setInputQuestion(e.target.value)}
             fullWidth
             compressed
             prepend={<EuiIcon type={config.branding.logo || chatIcon} />}
@@ -393,8 +414,8 @@ export const Text2Viz = () => {
         <EuiFlexItem grow={false}>
           <EuiButtonIcon
             aria-label="submit"
-            onClick={onSubmit}
-            isDisabled={loading || input.trim().length === 0 || !selectedSource}
+            onClick={() => onSubmit()}
+            isDisabled={loading || inputQuestion.trim().length === 0 || !selectedSource}
             display="base"
             size="s"
             iconType="returnKey"
@@ -453,13 +474,23 @@ export const Text2Viz = () => {
                   paddingSize="none"
                   scrollable={false}
                 >
-                  {usageCollection ? (
-                    <FeedbackThumbs
-                      usageCollection={usageCollection}
-                      appName={VIS_NLQ_APP_ID}
-                      className="feedback_thumbs"
-                    />
-                  ) : null}
+                  <EuiFlexGroup className="text2viz__actionContainer" gutterSize="none">
+                    {usageCollection ? (
+                      <EuiFlexItem className="text2viz__feedbackContainer">
+                        <FeedbackThumbs
+                          usageCollection={usageCollection}
+                          appName={VIS_NLQ_APP_ID}
+                        />
+                      </EuiFlexItem>
+                    ) : null}
+                    <EuiFlexItem className="text2viz__vizStyleEditorContainer">
+                      <VizStyleEditor
+                        iconType={config.branding.logo || chatIcon}
+                        onApply={(instruction) => onSubmit(instruction)}
+                        value={currentInstruction}
+                      />
+                    </EuiFlexItem>
+                  </EuiFlexGroup>
                   <EmbeddableRenderer factory={factory} input={visInput} />
                 </EuiResizablePanel>
                 <EuiResizableButton />

diff --git a/public/components/visualization/viz_style_editor.test.tsx b/public/components/visualization/viz_style_editor.test.tsx
@@ -0,0 +1,44 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { render, screen, fireEvent, waitFor } from '@testing-library/react';
+import { VizStyleEditor } from './viz_style_editor';
+
+describe('<VizStyleEditor />', () => {
+  test('should render visual style editor', () => {
+    const onApplyFn = jest.fn();
+    render(<VizStyleEditor onApply={onApplyFn} iconType="icon" />);
+    expect(screen.queryByText('Edit visual')).toBeInTheDocument();
+
+    // click Edit visual button to open the modal
+    expect(screen.queryByTestId('text2vizStyleEditorModal')).toBe(null);
+    fireEvent.click(screen.getByText('Edit visual'));
+    expect(screen.queryByTestId('text2vizStyleEditorModal')).toBeInTheDocument();
+
+    // Click cancel to close the modal
+    fireEvent.click(screen.getByText('Cancel'));
+    expect(screen.queryByTestId('text2vizStyleEditorModal')).toBe(null);
+
+    // Apply button is disabled
+    fireEvent.click(screen.getByText('Edit visual'));
+    expect(screen.getByTestId('text2vizStyleEditorModalApply')).toBeDisabled();
+
+    // After input text, Apply button is enabled
+    fireEvent.input(screen.getByLabelText('Input instructions to tweak the visual'), {
+      target: { value: 'test input' },
+    });
+    expect(screen.getByTestId('text2vizStyleEditorModalApply')).not.toBeDisabled();
+    fireEvent.click(screen.getByText('Apply'));
+    expect(onApplyFn).toHaveBeenCalledWith('test input');
+    expect(screen.queryByTestId('text2vizStyleEditorModal')).toBe(null);
+  });
+
+  test('should open the modal with initial value', () => {
+    render(<VizStyleEditor onApply={jest.fn()} iconType="icon" value="test input" />);
+    fireEvent.click(screen.getByText('Edit visual'));
+    expect(screen.getByDisplayValue('test input')).toBeInTheDocument();
+  });
+});