From b46c79e08a1bbd2592e18af96dfec12d5f0cafb5 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Mon, 21 Oct 2024 07:01:10 +0000
Subject: [PATCH] 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>
(cherry picked from commit 6a49d78b697f9fc3fdbeb2ee3cae765c2aa8d51b)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md
---
 public/components/visualization/text2vega.ts  |  19 +--
 public/components/visualization/text2viz.scss |  20 ++-
 public/components/visualization/text2viz.tsx  | 125 +++++++++++-------
 .../visualization/viz_style_editor.test.tsx   |  44 ++++++
 .../visualization/viz_style_editor.tsx        | 101 ++++++++++++++
 server/routes/agent_routes.ts                 |   2 +-
 6 files changed, 253 insertions(+), 58 deletions(-)
 create mode 100644 public/components/visualization/viz_style_editor.test.tsx
 create mode 100644 public/components/visualization/viz_style_editor.tsx

diff --git a/public/components/visualization/text2vega.ts b/public/components/visualization/text2vega.ts
index ea8b0936..10108aa5 100644
--- 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
index c917471e..bc27122a 100644
--- 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
index bdc4f8ab..c4a76dec 100644
--- 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
new file mode 100644
index 00000000..7c2c6a61
--- /dev/null
+++ 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();
+  });
+});
diff --git a/public/components/visualization/viz_style_editor.tsx b/public/components/visualization/viz_style_editor.tsx
new file mode 100644
index 00000000..0d35796e
--- /dev/null
+++ b/public/components/visualization/viz_style_editor.tsx
@@ -0,0 +1,101 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React, { useCallback, useState } from 'react';
+import {
+  EuiButton,
+  EuiButtonEmpty,
+  EuiModal,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiModalHeader,
+  EuiModalHeaderTitle,
+  EuiSpacer,
+  EuiTextArea,
+} from '@elastic/eui';
+import { i18n } from '@osd/i18n';
+
+interface Props {
+  onApply: (input: string) => void;
+  value?: string;
+  iconType: string;
+  className?: string;
+}
+
+export const VizStyleEditor = ({ onApply, className, iconType, value }: Props) => {
+  const [modalVisible, setModalVisible] = useState(false);
+  const [inputValue, setInputValue] = useState('');
+
+  const onApplyClick = useCallback(() => {
+    onApply(inputValue.trim());
+    setModalVisible(false);
+  }, [inputValue, onApply]);
+
+  const openModal = useCallback(() => {
+    if (value) {
+      setInputValue(value);
+    }
+    setModalVisible(true);
+  }, [value]);
+
+  return (
+    <div className={className}>
+      <EuiButton
+        className="vizStyleEditor__editButton"
+        size="s"
+        iconType={iconType}
+        onClick={openModal}
+      >
+        {i18n.translate('dashboardAssistant.feature.text2viz.editVisualButton.label', {
+          defaultMessage: 'Edit visual',
+        })}
+      </EuiButton>
+      {modalVisible && (
+        <EuiModal data-test-subj="text2vizStyleEditorModal" onClose={() => setModalVisible(false)}>
+          <EuiModalHeader>
+            <EuiModalHeaderTitle>
+              <h1>
+                {i18n.translate('dashboardAssistant.feature.text2viz.editVisualModal.title', {
+                  defaultMessage: 'Edit visual',
+                })}
+              </h1>
+            </EuiModalHeaderTitle>
+          </EuiModalHeader>
+          <EuiModalBody>
+            {i18n.translate('dashboardAssistant.feature.text2viz.editVisualModal.body', {
+              defaultMessage: 'How would you like to edit the visual?',
+            })}
+            <EuiSpacer size="s" />
+            <EuiTextArea
+              compressed
+              autoFocus
+              aria-label="Input instructions to tweak the visual"
+              value={inputValue}
+              onChange={(e) => setInputValue(e.target.value)}
+            />
+          </EuiModalBody>
+          <EuiModalFooter>
+            <EuiButtonEmpty size="s" onClick={() => setModalVisible(false)}>
+              {i18n.translate('dashboardAssistant.feature.text2viz.editVisualModal.cancel', {
+                defaultMessage: 'Cancel',
+              })}
+            </EuiButtonEmpty>
+            <EuiButton
+              data-test-subj="text2vizStyleEditorModalApply"
+              size="s"
+              fill
+              onClick={onApplyClick}
+              disabled={inputValue.trim().length === 0}
+            >
+              {i18n.translate('dashboardAssistant.feature.text2viz.editVisualModal.apply', {
+                defaultMessage: 'Apply',
+              })}
+            </EuiButton>
+          </EuiModalFooter>
+        </EuiModal>
+      )}
+    </div>
+  );
+};
diff --git a/server/routes/agent_routes.ts b/server/routes/agent_routes.ts
index 227143a7..57e85e66 100644
--- a/server/routes/agent_routes.ts
+++ b/server/routes/agent_routes.ts
@@ -39,7 +39,7 @@ export function registerAgentRoutes(router: IRouter, assistantService: Assistant
         );
         return res.ok({ body: response });
       } catch (e) {
-        return res.internalError();
+        return res.badRequest();
       }
     })
   );