Support multiple input/output maps for ML processors (#244)

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>

Author: ohltyler

common/interfaces.ts

@@ -23,7 +23,8 @@ export type ConfigFieldType =
   | 'jsonArray'
   | 'select'
   | 'model'
-  | 'map';
+  | 'map'
+  | 'mapArray';
 export type ConfigFieldValue = string | {};
 export interface IConfigField {
   type: ConfigFieldType;
@@ -81,6 +82,8 @@ export type MapEntry = {
 
 export type MapFormValue = MapEntry[];
 
+export type MapArrayFormValue = MapFormValue[];
+
 export type WorkflowFormValues = {
   ingest: FormikValues;
   search: FormikValues;

public/configs/ml_processor.ts

@@ -22,11 +22,11 @@ export abstract class MLProcessor extends Processor {
       },
       {
         id: 'inputMap',
-        type: 'map',
+        type: 'mapArray',
       },
       {
         id: 'outputMap',
-        type: 'map',
+        type: 'mapArray',
       },
     ];
   }

public/pages/workflow_detail/workflow_inputs/input_fields/index.ts

@@ -7,4 +7,5 @@ export { TextField } from './text_field';
 export { JsonField } from './json_field';
 export { ModelField } from './model_field';
 export { MapField } from './map_field';
+export { MapArrayField } from './map_array_field';
 export { BooleanField } from './boolean_field';

public/pages/workflow_detail/workflow_inputs/input_fields/map_array_field.tsx

@@ -0,0 +1,149 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import {
+  EuiAccordion,
+  EuiButton,
+  EuiButtonIcon,
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiFormRow,
+  EuiLink,
+  EuiPanel,
+  EuiText,
+} from '@elastic/eui';
+import { Field, FieldProps, getIn, useFormikContext } from 'formik';
+import {
+  IConfigField,
+  MapArrayFormValue,
+  MapEntry,
+  WorkflowFormValues,
+} from '../../../../../common';
+import { MapField } from './map_field';
+
+interface MapArrayFieldProps {
+  field: IConfigField;
+  fieldPath: string; // the full path in string-form to the field (e.g., 'ingest.enrich.processors.text_embedding_processor.inputField')
+  label?: string;
+  helpLink?: string;
+  helpText?: string;
+  keyPlaceholder?: string;
+  valuePlaceholder?: string;
+  onFormChange: () => void;
+  onMapAdd?: (curArray: MapArrayFormValue) => void;
+  onMapDelete?: (idxToDelete: number) => void;
+}
+
+/**
+ * Input component for configuring an array of field mappings
+ */
+export function MapArrayField(props: MapArrayFieldProps) {
+  const { setFieldValue, setFieldTouched, errors, touched } = useFormikContext<
+    WorkflowFormValues
+  >();
+
+  // Adding a map to the end of the existing arr
+  function addMap(curMapArray: MapArrayFormValue): void {
+    setFieldValue(props.fieldPath, [...curMapArray, []]);
+    setFieldTouched(props.fieldPath, true);
+    props.onFormChange();
+    if (props.onMapAdd) {
+      props.onMapAdd(curMapArray);
+    }
+  }
+
+  // Deleting a map
+  function deleteMap(
+    curMapArray: MapArrayFormValue,
+    entryIndexToDelete: number
+  ): void {
+    const updatedMapArray = [...curMapArray];
+    updatedMapArray.splice(entryIndexToDelete, 1);
+    setFieldValue(props.fieldPath, updatedMapArray);
+    setFieldTouched(props.fieldPath, true);
+    props.onFormChange();
+    if (props.onMapDelete) {
+      props.onMapDelete(entryIndexToDelete);
+    }
+  }
+
+  return (
+    <Field name={props.fieldPath}>
+      {({ field, form }: FieldProps) => {
+        return (
+          <EuiFormRow
+            fullWidth={true}
+            key={props.fieldPath}
+            label={props.label}
+            labelAppend={
+              props.helpLink ? (
+                <EuiText size="xs">
+                  <EuiLink href={props.helpLink} target="_blank">
+                    Learn more
+                  </EuiLink>
+                </EuiText>
+              ) : undefined
+            }
+            helpText={props.helpText || undefined}
+            isInvalid={
+              getIn(errors, field.name) !== undefined &&
+              getIn(errors, field.name).length > 0 &&
+              getIn(touched, field.name) !== undefined &&
+              getIn(touched, field.name).length > 0
+            }
+          >
+            <EuiFlexGroup direction="column">
+              {field.value?.map((mapping: MapEntry, idx: number) => {
+                return (
+                  <EuiFlexItem key={idx}>
+                    <EuiAccordion
+                      key={idx}
+                      id={`accordion${idx}`}
+                      buttonContent={`Prediction ${idx + 1}`}
+                      paddingSize="m"
+                      extraAction={
+                        <EuiButtonIcon
+                          style={{ marginTop: '8px' }}
+                          iconType={'trash'}
+                          color="danger"
+                          aria-label="Delete"
+                          onClick={() => {
+                            deleteMap(field.value, idx);
+                          }}
+                        />
+                      }
+                    >
+                      <EuiPanel grow={true}>
+                        <MapField
+                          fieldPath={`${props.fieldPath}.${idx}`}
+                          keyPlaceholder={props.keyPlaceholder}
+                          valuePlaceholder={props.valuePlaceholder}
+                          onFormChange={props.onFormChange}
+                        />
+                      </EuiPanel>
+                    </EuiAccordion>
+                  </EuiFlexItem>
+                );
+              })}
+              <EuiFlexItem grow={false}>
+                <div>
+                  <EuiButton
+                    size="s"
+                    onClick={() => {
+                      addMap(field.value);
+                    }}
+                  >
+                    {field.value?.length > 0 ? 'Add another map' : 'Add map'}
+                  </EuiButton>
+                </div>
+              </EuiFlexItem>
+            </EuiFlexGroup>
+          </EuiFormRow>
+        );
+      }}
+    </Field>
+  );
+}

public/pages/workflow_detail/workflow_inputs/input_fields/map_field.tsx

@@ -16,16 +16,14 @@ import {
 } from '@elastic/eui';
 import { Field, FieldProps, getIn, useFormikContext } from 'formik';
 import {
-  IConfigField,
   MapEntry,
   MapFormValue,
   WorkflowFormValues,
 } from '../../../../../common';
 
 interface MapFieldProps {
-  field: IConfigField;
   fieldPath: string; // the full path in string-form to the field (e.g., 'ingest.enrich.processors.text_embedding_processor.inputField')
-  label: string;
+  label?: string;
   helpLink?: string;
   helpText?: string;
   keyPlaceholder?: string;
@@ -62,10 +60,11 @@ export function MapField(props: MapFieldProps) {
   }
 
   return (
-    <Field name={props.fieldPath}>
+    <Field name={props.fieldPath} key={props.fieldPath}>
       {({ field, form }: FieldProps) => {
         return (
           <EuiFormRow
+            fullWidth={true}
             key={props.fieldPath}
             label={props.label}
             labelAppend={
@@ -95,9 +94,10 @@ export function MapField(props: MapFieldProps) {
               {field.value?.map((mapping: MapEntry, idx: number) => {
                 return (
                   <EuiFlexItem key={idx}>
-                    <EuiFlexGroup direction="row" justifyContent="spaceBetween">
-                      <EuiFlexItem grow={false}>
+                    <EuiFlexGroup direction="row">
+                      <EuiFlexItem grow={true}>
                         <EuiFormControlLayoutDelimited
+                          fullWidth={true}
                           startControl={
                             <input
                               type="string"

public/pages/workflow_detail/workflow_inputs/processor_inputs/input_transform_modal.tsx

@@ -16,6 +16,8 @@ import {
   EuiModalFooter,
   EuiModalHeader,
   EuiModalHeaderTitle,
+  EuiSelect,
+  EuiSelectOption,
   EuiSpacer,
   EuiText,
 } from '@elastic/eui';
@@ -25,6 +27,7 @@ import {
   IngestPipelineConfig,
   JSONPATH_ROOT_SELECTOR,
   ML_INFERENCE_DOCS_LINK,
+  MapArrayFormValue,
   PROCESSOR_CONTEXT,
   SimulateIngestPipelineResponse,
   WorkflowConfig,
@@ -38,7 +41,7 @@ import {
 } from '../../../../utils';
 import { simulatePipeline, useAppDispatch } from '../../../../store';
 import { getCore } from '../../../../services';
-import { MapField } from '../input_fields';
+import { MapArrayField } from '../input_fields';
 
 interface InputTransformModalProps {
   uiConfig: WorkflowConfig;
@@ -50,6 +53,10 @@ interface InputTransformModalProps {
   onFormChange: () => void;
 }
 
+// TODO: InputTransformModal and OutputTransformModal are very similar, and can
+// likely be refactored and have more reusable components. Leave as-is until the
+// UI is more finalized.
+
 /**
  * A modal to configure advanced JSON-to-JSON transforms into a model's expected input
  */
@@ -59,10 +66,19 @@ export function InputTransformModal(props: InputTransformModalProps) {
 
   // source input / transformed output state
   const [sourceInput, setSourceInput] = useState<string>('[]');
-  const [transformedOutput, setTransformedOutput] = useState<string>('[]');
+  const [transformedOutput, setTransformedOutput] = useState<string>('{}');
 
   // get the current input map
-  const map = getIn(values, `ingest.enrich.${props.config.id}.inputMap`);
+  const map = getIn(values, props.inputMapFieldPath) as MapArrayFormValue;
+
+  // selected output state
+  const outputOptions = map.map((_, idx) => ({
+    value: idx,
+    text: `Prediction ${idx + 1}`,
+  })) as EuiSelectOption[];
+  const [selectedOutputOption, setSelectedOutputOption] = useState<
+    number | undefined
+  >((outputOptions[0]?.value as number) ?? undefined);
 
   return (
     <EuiModal onClose={props.onClose} style={{ width: '70vw' }}>
@@ -149,34 +165,62 @@ export function InputTransformModal(props: InputTransformModalProps) {
                 root object selector "${JSONPATH_ROOT_SELECTOR}"`}
               </EuiText>
               <EuiSpacer size="s" />
-              <MapField
+              <MapArrayField
                 field={props.inputMapField}
                 fieldPath={props.inputMapFieldPath}
-                label="Input map"
+                label="Input Map"
                 helpText={`An array specifying how to map fields from the ingested document to the model’s input.`}
                 helpLink={ML_INFERENCE_DOCS_LINK}
                 keyPlaceholder="Model input field"
                 valuePlaceholder="Document field"
                 onFormChange={props.onFormChange}
+                // If the map we are adding is the first one, populate the selected option to index 0
+                onMapAdd={(curArray) => {
+                  if (isEmpty(curArray)) {
+                    setSelectedOutputOption(0);
+                  }
+                }}
+                // If the map we are deleting is the one we last used to test, reset the state and
+                // default to the first map in the list.
+                onMapDelete={(idxToDelete) => {
+                  if (selectedOutputOption === idxToDelete) {
+                    setSelectedOutputOption(0);
+                    setTransformedOutput('{}');
+                  }
+                }}
               />
             </>
           </EuiFlexItem>
           <EuiFlexItem>
             <>
-              <EuiText>Expected output</EuiText>
+              <EuiSelect
+                prepend={<EuiText>Expected output for</EuiText>}
+                compressed={true}
+                options={outputOptions}
+                value={selectedOutputOption}
+                onChange={(e) => {
+                  setSelectedOutputOption(Number(e.target.value));
+                  setTransformedOutput('{}');
+                }}
+              />
+              <EuiSpacer size="s" />
               <EuiButton
                 style={{ width: '100px' }}
                 disabled={isEmpty(map) || isEmpty(JSON.parse(sourceInput))}
                 onClick={async () => {
                   switch (props.context) {
                     case PROCESSOR_CONTEXT.INGEST: {
-                      if (!isEmpty(map) && !isEmpty(JSON.parse(sourceInput))) {
+                      if (
+                        !isEmpty(map) &&
+                        !isEmpty(JSON.parse(sourceInput)) &&
+                        selectedOutputOption !== undefined
+                      ) {
                         let sampleSourceInput = {};
                         try {
                           sampleSourceInput = JSON.parse(sourceInput)[0];
                           const output = generateTransform(
                             sampleSourceInput,
-                            map
+                            map[selectedOutputOption]
                           );
                           setTransformedOutput(
                             JSON.stringify(output, undefined, 2)