From 1b8486e576e3345da4eaa8c9cba92c4b4056252a Mon Sep 17 00:00:00 2001
From: Tyler Ohlsen <ohltyler@amazon.com>
Date: Tue, 23 Jul 2024 12:30:20 -0700
Subject: [PATCH] Add guardrails to doc and query JSON inputs (#231)

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
---
 common/constants.ts                           |   6 +
 common/interfaces.ts                          |   9 +-
 .../workflow_detail/resizable_workspace.tsx   |   2 +-
 .../ingest_inputs/source_data.tsx             | 121 ++++++++++++-----
 .../input_fields/json_field.tsx               |   5 +-
 .../input_transform_modal.tsx                 |   5 +-
 .../processor_inputs/ml_processor_inputs.tsx  |   9 +-
 .../configure_search_request.tsx              | 127 ++++++++++++------
 public/pages/workflows/new_workflow/utils.ts  |   7 +-
 public/utils/config_to_form_utils.ts          |  12 +-
 public/utils/form_to_config_utils.ts          |   4 +
 11 files changed, 208 insertions(+), 99 deletions(-)

diff --git a/common/constants.ts b/common/constants.ts
index 99cbe914..e105e961 100644
--- a/common/constants.ts
+++ b/common/constants.ts
@@ -112,6 +112,12 @@ export enum COMPONENT_CLASS {
   RESULTS = 'results',
 }
 
+/**
+ * LINKS
+ */
+export const ML_INFERENCE_DOCS_LINK =
+  'https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/#configuration-parameters';
+
 /**
  * MISCELLANEOUS
  */
diff --git a/common/interfaces.ts b/common/interfaces.ts
index 205707ae..e5bc1914 100644
--- a/common/interfaces.ts
+++ b/common/interfaces.ts
@@ -29,13 +29,6 @@ export interface IConfigField {
   type: ConfigFieldType;
   id: string;
   value?: ConfigFieldValue;
-  // TODO: remove below fields out of this interface and directly into the necessary components.
-  // This is to minimize what we persist here, which is added into ui_metadata and indexed.
-  // Once the config for ML inference processors is finalized, we can migrate these out.
-  label?: string;
-  placeholder?: string;
-  helpText?: string;
-  helpLink?: string;
 }
 export interface IConfig {
   id: string;
@@ -71,7 +64,7 @@ export type IngestConfig = {
 };
 
 export type SearchConfig = {
-  request: {};
+  request: IConfigField;
   enrichRequest: ProcessorsConfig;
   enrichResponse: ProcessorsConfig;
 };
diff --git a/public/pages/workflow_detail/resizable_workspace.tsx b/public/pages/workflow_detail/resizable_workspace.tsx
index 787817ae..cabc76eb 100644
--- a/public/pages/workflow_detail/resizable_workspace.tsx
+++ b/public/pages/workflow_detail/resizable_workspace.tsx
@@ -128,7 +128,7 @@ export function ResizableWorkspace(props: ResizableWorkspaceProps) {
   // Initialize the form state based on the current UI config
   useEffect(() => {
     if (uiConfig) {
-      const initFormValues = uiConfigToFormik(uiConfig, ingestDocs, query);
+      const initFormValues = uiConfigToFormik(uiConfig, ingestDocs);
       const initFormSchema = uiConfigToSchema(uiConfig);
       setFormValues(initFormValues);
       setFormSchema(initFormSchema);
diff --git a/public/pages/workflow_detail/workflow_inputs/ingest_inputs/source_data.tsx b/public/pages/workflow_detail/workflow_inputs/ingest_inputs/source_data.tsx
index 2e8a1937..6e5b5b16 100644
--- a/public/pages/workflow_detail/workflow_inputs/ingest_inputs/source_data.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/ingest_inputs/source_data.tsx
@@ -3,12 +3,21 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useEffect } from 'react';
-import { useFormikContext } from 'formik';
+import React, { useEffect, useState } from 'react';
+import { useFormikContext, getIn } from 'formik';
 import {
+  EuiButton,
+  EuiCodeBlock,
   EuiFilePicker,
   EuiFlexGroup,
   EuiFlexItem,
+  EuiModal,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiModalHeader,
+  EuiModalHeaderTitle,
+  EuiSpacer,
+  EuiText,
   EuiTitle,
 } from '@elastic/eui';
 import { JsonField } from '../input_fields';
@@ -25,6 +34,9 @@ interface SourceDataProps {
 export function SourceData(props: SourceDataProps) {
   const { values, setFieldValue } = useFormikContext<WorkspaceFormValues>();
 
+  // edit modal state
+  const [isEditModalOpen, setIsEditModalOpen] = useState<boolean>(false);
+
   // files state. when a file is read, update the form value.
   const fileReader = new FileReader();
   fileReader.onload = (e) => {
@@ -42,36 +54,79 @@ export function SourceData(props: SourceDataProps) {
   }, [values?.ingest?.docs]);
 
   return (
-    <EuiFlexGroup direction="column">
-      <EuiFlexItem grow={false}>
-        <EuiTitle size="s">
-          <h2>Source data</h2>
-        </EuiTitle>
-      </EuiFlexItem>
-      <EuiFlexItem grow={false}>
-        <EuiFilePicker
-          accept="application/json"
-          multiple={false}
-          initialPromptText="Select a JSON file containing documents"
-          onChange={(files) => {
-            if (files && files.length > 0) {
-              fileReader.readAsText(files[0]);
-            }
-          }}
-          display="default"
-        />
-      </EuiFlexItem>
-      <EuiFlexItem grow={false}>
-        <JsonField
-          label="Upload JSON documents"
-          fieldPath={'ingest.docs'}
-          helpText="Documents should be formatted as a valid JSON array."
-          // when ingest doc values change, don't update the form
-          // since we initially only support running ingest once per configuration
-          onFormChange={() => {}}
-          editorHeight="25vh"
-        />
-      </EuiFlexItem>
-    </EuiFlexGroup>
+    <>
+      {isEditModalOpen && (
+        <EuiModal
+          onClose={() => setIsEditModalOpen(false)}
+          style={{ width: '70vw' }}
+        >
+          <EuiModalHeader>
+            <EuiModalHeaderTitle>
+              <p>{`Edit source data`}</p>
+            </EuiModalHeaderTitle>
+          </EuiModalHeader>
+          <EuiModalBody>
+            <>
+              <EuiText color="subdued">
+                Upload a JSON file or enter manually.
+              </EuiText>{' '}
+              <EuiSpacer size="s" />
+              <EuiFilePicker
+                accept="application/json"
+                multiple={false}
+                initialPromptText="Upload file"
+                onChange={(files) => {
+                  if (files && files.length > 0) {
+                    fileReader.readAsText(files[0]);
+                  }
+                }}
+                display="default"
+              />
+              <EuiSpacer size="s" />
+              <JsonField
+                fieldPath={'ingest.docs'}
+                helpText="Documents should be formatted as a valid JSON array."
+                // when ingest doc values change, don't update the form
+                // since we initially only support running ingest once per configuration
+                onFormChange={() => {}}
+                editorHeight="25vh"
+                readOnly={false}
+              />
+            </>
+          </EuiModalBody>
+          <EuiModalFooter>
+            <EuiButton
+              onClick={() => setIsEditModalOpen(false)}
+              fill={false}
+              color="primary"
+            >
+              Close
+            </EuiButton>
+          </EuiModalFooter>
+        </EuiModal>
+      )}
+      <EuiFlexGroup direction="column" gutterSize="s">
+        <EuiFlexItem grow={false}>
+          <EuiTitle size="s">
+            <h2>Source data</h2>
+          </EuiTitle>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiButton
+            fill={false}
+            style={{ width: '100px' }}
+            size="s"
+            onClick={() => setIsEditModalOpen(true)}
+          >
+            Edit
+          </EuiButton>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiCodeBlock language="json" fontSize="m" isCopyable={false}>
+            {getIn(values, 'ingest.docs')}
+          </EuiCodeBlock>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </>
   );
 }
diff --git a/public/pages/workflow_detail/workflow_inputs/input_fields/json_field.tsx b/public/pages/workflow_detail/workflow_inputs/input_fields/json_field.tsx
index 87dace8d..44584a9f 100644
--- a/public/pages/workflow_detail/workflow_inputs/input_fields/json_field.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/input_fields/json_field.tsx
@@ -11,10 +11,11 @@ import { WorkspaceFormValues } from '../../../../../common';
 interface JsonFieldProps {
   fieldPath: string; // the full path in string-form to the field (e.g., 'ingest.enrich.processors.text_embedding_processor.inputField')
   onFormChange: () => void;
-  label: string;
+  label?: string;
   helpLink?: string;
   helpText?: string;
   editorHeight?: string;
+  readOnly?: boolean;
 }
 
 /**
@@ -79,7 +80,7 @@ export function JsonField(props: JsonFieldProps) {
                   props.onFormChange();
                 }
               }}
-              readOnly={false}
+              readOnly={props.readOnly || false}
               setOptions={{
                 fontSize: '14px',
               }}
diff --git a/public/pages/workflow_detail/workflow_inputs/processor_inputs/input_transform_modal.tsx b/public/pages/workflow_detail/workflow_inputs/processor_inputs/input_transform_modal.tsx
index ba7b782a..31a903f9 100644
--- a/public/pages/workflow_detail/workflow_inputs/processor_inputs/input_transform_modal.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/processor_inputs/input_transform_modal.tsx
@@ -25,6 +25,7 @@ import {
   IProcessorConfig,
   IngestPipelineConfig,
   JSONPATH_ROOT_SELECTOR,
+  ML_INFERENCE_DOCS_LINK,
   PROCESSOR_CONTEXT,
   SimulateIngestPipelineDoc,
   SimulateIngestPipelineResponse,
@@ -135,9 +136,7 @@ export function InputTransformModal(props: InputTransformModalProps) {
                 fieldPath={props.inputMapFieldPath}
                 label="Input map"
                 helpText={`An array specifying how to map fields from the ingested document to the model’s input.`}
-                helpLink={
-                  'https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/#configuration-parameters'
-                }
+                helpLink={ML_INFERENCE_DOCS_LINK}
                 keyPlaceholder="Model input field"
                 valuePlaceholder="Document field"
                 onFormChange={props.onFormChange}
diff --git a/public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx b/public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx
index 1b93fb3b..01ab9894 100644
--- a/public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx
@@ -19,6 +19,7 @@ import {
   PROCESSOR_CONTEXT,
   WorkflowConfig,
   JSONPATH_ROOT_SELECTOR,
+  ML_INFERENCE_DOCS_LINK,
 } from '../../../../../common';
 import { MapField, ModelField } from '../input_fields';
 import { isEmpty } from 'lodash';
@@ -128,9 +129,7 @@ export function MLProcessorInputs(props: MLProcessorInputsProps) {
                 fieldPath={inputMapFieldPath}
                 label="Input map"
                 helpText={`An array specifying how to map fields from the ingested document to the model’s input.`}
-                helpLink={
-                  'https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/#configuration-parameters'
-                }
+                helpLink={ML_INFERENCE_DOCS_LINK}
                 keyPlaceholder="Model input field"
                 valuePlaceholder="Document field"
                 onFormChange={props.onFormChange}
@@ -141,9 +140,7 @@ export function MLProcessorInputs(props: MLProcessorInputsProps) {
                 fieldPath={outputMapFieldPath}
                 label="Output map"
                 helpText={`An array specifying how to map the model’s output to new fields.`}
-                helpLink={
-                  'https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/#configuration-parameters'
-                }
+                helpLink={ML_INFERENCE_DOCS_LINK}
                 keyPlaceholder="New document field"
                 valuePlaceholder="Model output field"
                 onFormChange={props.onFormChange}
diff --git a/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx b/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx
index 596abb2a..a761f2fa 100644
--- a/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx
@@ -5,18 +5,25 @@
 
 import React, { useEffect, useState } from 'react';
 import { useSelector } from 'react-redux';
+import { useFormikContext, getIn } from 'formik';
 import {
+  EuiButton,
+  EuiCodeBlock,
   EuiFieldText,
   EuiFlexGroup,
   EuiFlexItem,
   EuiFormRow,
+  EuiModal,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiModalHeader,
+  EuiModalHeaderTitle,
   EuiSuperSelect,
   EuiSuperSelectOption,
   EuiText,
   EuiTitle,
 } from '@elastic/eui';
-import { useFormikContext } from 'formik';
-import { IConfigField, WorkspaceFormValues } from '../../../../../common';
+import { WorkspaceFormValues } from '../../../../../common';
 import { JsonField } from '../input_fields';
 import { AppState, catIndices, useAppDispatch } from '../../../../store';
 
@@ -44,6 +51,9 @@ export function ConfigureSearchRequest(props: ConfigureSearchRequestProps) {
     undefined
   );
 
+  // Edit modal state
+  const [isEditModalOpen, setIsEditModalOpen] = useState<boolean>(false);
+
   // Hook to listen when the query form value changes.
   // Try to set the query request if possible
   useEffect(() => {
@@ -61,43 +71,82 @@ export function ConfigureSearchRequest(props: ConfigureSearchRequestProps) {
   }, []);
 
   return (
-    <EuiFlexGroup direction="column">
-      <EuiFlexItem grow={false}>
-        <EuiTitle size="s">
-          <h2>Configure query</h2>
-        </EuiTitle>
-      </EuiFlexItem>
-      <EuiFlexItem grow={false}>
-        <EuiFormRow label="Retrieval index">
-          {ingestEnabled ? (
-            <EuiFieldText value={indexName} readOnly={true} />
-          ) : (
-            <EuiSuperSelect
-              options={Object.values(indices).map(
-                (option) =>
-                  ({
-                    value: option.name,
-                    inputDisplay: <EuiText>{option.name}</EuiText>,
-                    disabled: false,
-                  } as EuiSuperSelectOption<string>)
-              )}
-              valueOfSelected={selectedIndex}
-              onChange={(option) => {
-                setSelectedIndex(option);
-              }}
-              isInvalid={selectedIndex !== undefined}
+    <>
+      {isEditModalOpen && (
+        <EuiModal
+          onClose={() => setIsEditModalOpen(false)}
+          style={{ width: '70vw' }}
+        >
+          <EuiModalHeader>
+            <EuiModalHeaderTitle>
+              <p>{`Edit query`}</p>
+            </EuiModalHeaderTitle>
+          </EuiModalHeader>
+          <EuiModalBody>
+            <JsonField
+              label="Query"
+              fieldPath={'search.request'}
+              onFormChange={props.onFormChange}
+              editorHeight="25vh"
+              readOnly={false}
             />
-          )}
-        </EuiFormRow>
-      </EuiFlexItem>
-      <EuiFlexItem grow={false}>
-        <JsonField
-          label="Define query"
-          fieldPath={'search.request'}
-          onFormChange={props.onFormChange}
-          editorHeight="25vh"
-        />
-      </EuiFlexItem>
-    </EuiFlexGroup>
+          </EuiModalBody>
+          <EuiModalFooter>
+            <EuiButton
+              onClick={() => setIsEditModalOpen(false)}
+              fill={false}
+              color="primary"
+            >
+              Close
+            </EuiButton>
+          </EuiModalFooter>
+        </EuiModal>
+      )}
+      <EuiFlexGroup direction="column">
+        <EuiFlexItem grow={false}>
+          <EuiTitle size="s">
+            <h2>Configure query</h2>
+          </EuiTitle>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiFormRow label="Retrieval index">
+            {ingestEnabled ? (
+              <EuiFieldText value={indexName} readOnly={true} />
+            ) : (
+              <EuiSuperSelect
+                options={Object.values(indices).map(
+                  (option) =>
+                    ({
+                      value: option.name,
+                      inputDisplay: <EuiText>{option.name}</EuiText>,
+                      disabled: false,
+                    } as EuiSuperSelectOption<string>)
+                )}
+                valueOfSelected={selectedIndex}
+                onChange={(option) => {
+                  setSelectedIndex(option);
+                }}
+                isInvalid={selectedIndex === undefined}
+              />
+            )}
+          </EuiFormRow>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiButton
+            fill={false}
+            style={{ width: '100px' }}
+            size="s"
+            onClick={() => setIsEditModalOpen(true)}
+          >
+            Edit
+          </EuiButton>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiCodeBlock language="json" fontSize="m" isCopyable={false}>
+            {getIn(values, 'search.request')}
+          </EuiCodeBlock>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </>
   );
 }
diff --git a/public/pages/workflows/new_workflow/utils.ts b/public/pages/workflows/new_workflow/utils.ts
index dcd2bddd..d45f6891 100644
--- a/public/pages/workflows/new_workflow/utils.ts
+++ b/public/pages/workflows/new_workflow/utils.ts
@@ -10,6 +10,7 @@ import {
   DEFAULT_NEW_WORKFLOW_NAME,
   UIState,
   WORKFLOW_TYPE,
+  FETCH_ALL_QUERY_BODY,
 } from '../../../../common';
 
 // Fn to produce the complete preset template with all necessary UI metadata.
@@ -64,7 +65,11 @@ function fetchEmptyMetadata(): UIState {
         },
       },
       search: {
-        request: {},
+        request: {
+          id: 'request',
+          type: 'json',
+          value: JSON.stringify(FETCH_ALL_QUERY_BODY, undefined, 2),
+        },
         enrichRequest: {
           processors: [],
         },
diff --git a/public/utils/config_to_form_utils.ts b/public/utils/config_to_form_utils.ts
index afc51414..57c1a73d 100644
--- a/public/utils/config_to_form_utils.ts
+++ b/public/utils/config_to_form_utils.ts
@@ -15,6 +15,7 @@ import {
   ConfigFieldType,
   ConfigFieldValue,
   ModelFormValue,
+  FETCH_ALL_QUERY_BODY,
 } from '../../common';
 
 /*
@@ -26,12 +27,11 @@ import {
 // and can be extremely large. so we pass that as a standalone field
 export function uiConfigToFormik(
   config: WorkflowConfig,
-  ingestDocs: string,
-  query: string
+  ingestDocs: string
 ): WorkflowFormValues {
   const formikValues = {} as WorkflowFormValues;
   formikValues['ingest'] = ingestConfigToFormik(config.ingest, ingestDocs);
-  formikValues['search'] = searchConfigToFormik(config.search, query);
+  formikValues['search'] = searchConfigToFormik(config.search);
   return formikValues;
 }
 
@@ -83,12 +83,12 @@ function indexConfigToFormik(indexConfig: IndexConfig): FormikValues {
 }
 
 function searchConfigToFormik(
-  searchConfig: SearchConfig | undefined,
-  query: string
+  searchConfig: SearchConfig | undefined
 ): FormikValues {
   let searchFormikValues = {} as FormikValues;
   if (searchConfig) {
-    searchFormikValues['request'] = query || getInitialValue('json');
+    searchFormikValues['request'] =
+      searchConfig.request.value || getInitialValue('json');
     searchFormikValues['enrichRequest'] = processorsConfigToFormik(
       searchConfig.enrichRequest
     );
diff --git a/public/utils/form_to_config_utils.ts b/public/utils/form_to_config_utils.ts
index 6ab866f1..1371be7b 100644
--- a/public/utils/form_to_config_utils.ts
+++ b/public/utils/form_to_config_utils.ts
@@ -72,6 +72,10 @@ function formikToSearchUiConfig(
 ): SearchConfig {
   return {
     ...existingConfig,
+    request: {
+      ...existingConfig.request,
+      value: searchFormValues['request'],
+    },
     enrichRequest: formikToProcessorsUiConfig(
       searchFormValues['enrichRequest'],
       existingConfig.enrichRequest