Support query_template optional field for ML search request process…

…or (#270) Signed-off-by: Tyler Ohlsen <[email protected]>
opensearch-project · Aug 7, 2024 · 211434e · 211434e
1 parent 1e4127c
commit 211434e
Show file tree

Hide file tree

Showing 7 changed files with 95 additions and 14 deletions.
diff --git a/common/interfaces.ts b/common/interfaces.ts
@@ -21,6 +21,7 @@ export type ConfigFieldType =
   | 'string'
   | 'json'
   | 'jsonArray'
+  | 'jsonString'
   | 'select'
   | 'model'
   | 'map'
@@ -218,6 +219,7 @@ export type MLInferenceProcessor = IngestProcessor & {
     model_id: string;
     input_map?: {};
     output_map?: {};
+    [key: string]: any;
   };
 };
 

diff --git a/public/configs/search_request_processors/ml_search_request_processor.ts b/public/configs/search_request_processors/ml_search_request_processor.ts
@@ -13,5 +13,12 @@ export class MLSearchRequestProcessor extends MLProcessor {
   constructor() {
     super();
     this.id = generateId('ml_processor_search_request');
+    this.optionalFields = [
+      {
+        id: 'query_template',
+        type: 'jsonString',
+      },
+      ...(this.optionalFields || []),
+    ];
   }
 }
diff --git a/public/pages/workflow_detail/workflow_inputs/config_field_list.tsx b/public/pages/workflow_detail/workflow_inputs/config_field_list.tsx
@@ -10,6 +10,7 @@ import {
   SelectField,
   BooleanField,
   NumberField,
+  JsonField,
 } from './input_fields';
 import { IConfigField } from '../../../../common';
 import { camelCaseToTitleString } from '../../../utils';
@@ -97,6 +98,33 @@ export function ConfigFieldList(props: ConfigFieldListProps) {
             );
             break;
           }
+          case 'json': {
+            el = (
+              <EuiFlexItem key={idx}>
+                <JsonField
+                  label={camelCaseToTitleString(field.id)}
+                  fieldPath={`${props.baseConfigPath}.${props.configId}.${field.id}`}
+                  onFormChange={props.onFormChange}
+                />
+                <EuiSpacer size={CONFIG_FIELD_SPACER_SIZE} />
+              </EuiFlexItem>
+            );
+            break;
+          }
+          case 'jsonString': {
+            el = (
+              <EuiFlexItem key={idx}>
+                <JsonField
+                  validate={false}
+                  label={camelCaseToTitleString(field.id)}
+                  fieldPath={`${props.baseConfigPath}.${props.configId}.${field.id}`}
+                  onFormChange={props.onFormChange}
+                />
+                <EuiSpacer size={CONFIG_FIELD_SPACER_SIZE} />
+              </EuiFlexItem>
+            );
+            break;
+          }
         }
         return el;
       })}

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
@@ -17,6 +17,7 @@ import { camelCaseToTitleString } from '../../../../utils';
 interface JsonFieldProps {
   fieldPath: string; // the full path in string-form to the field (e.g., 'ingest.enrich.processors.text_embedding_processor.inputField')
   onFormChange: () => void;
+  validate?: boolean;
   label?: string;
   helpLink?: string;
   helpText?: string;
@@ -29,6 +30,8 @@ interface JsonFieldProps {
  * in some custom JSON
  */
 export function JsonField(props: JsonFieldProps) {
+  const validate = props.validate !== undefined ? props.validate : true;
+
   const { errors, touched, values } = useFormikContext<WorkspaceFormValues>();
 
   // temp input state. only format when users click out of the code editor
@@ -61,8 +64,12 @@ export function JsonField(props: JsonFieldProps) {
               ) : undefined
             }
             helpText={props.helpText || undefined}
-            error={getIn(errors, field.name)}
-            isInvalid={getIn(errors, field.name) && getIn(touched, field.name)}
+            error={validate ? getIn(errors, field.name) : undefined}
+            isInvalid={
+              validate
+                ? getIn(errors, field.name) && getIn(touched, field.name)
+                : false
+            }
           >
             <EuiCodeEditor
               mode="json"
@@ -89,6 +96,7 @@ export function JsonField(props: JsonFieldProps) {
               readOnly={props.readOnly || false}
               setOptions={{
                 fontSize: '14px',
+                useWorker: validate,
               }}
               aria-label="Code Editor"
               tabSize={2}

diff --git a/public/utils/config_to_form_utils.ts b/public/utils/config_to_form_utils.ts
@@ -139,7 +139,8 @@ export function getInitialValue(fieldType: ConfigFieldType): ConfigFieldValue {
     case 'map': {
       return [];
     }
-    case 'json': {
+    case 'json':
+    case 'jsonString': {
       return '{}';
     }
     case 'jsonArray': {

diff --git a/public/utils/config_to_schema_utils.ts b/public/utils/config_to_schema_utils.ts
@@ -157,6 +157,10 @@ function getFieldSchema(
 
       break;
     }
+    case 'jsonString': {
+      baseSchema = yup.string().min(1, 'Too short');
+      break;
+    }
     case 'mapArray': {
       baseSchema = yup.array().of(
         yup.array().of(

diff --git a/public/utils/config_to_template_utils.ts b/public/utils/config_to_template_utils.ts
@@ -147,30 +147,61 @@ export function processorConfigsToTemplateProcessors(
   processorConfigs.forEach((processorConfig) => {
     switch (processorConfig.type) {
       case PROCESSOR_TYPE.ML: {
-        const { model, input_map, output_map } = processorConfigToFormik(
-          processorConfig
-        ) as {
-          model: ModelFormValue;
-          input_map: MapArrayFormValue;
-          output_map: MapArrayFormValue;
-        };
+        const {
+          model,
+          input_map,
+          output_map,
+          model_config,
+          ...formValues
+        } = processorConfigToFormik(processorConfig);
 
         let processor = {
           ml_inference: {
             model_id: model.id,
           },
         } as MLInferenceProcessor;
+
+        // process input/output maps
         if (input_map?.length > 0) {
-          processor.ml_inference.input_map = input_map.map((mapFormValue) =>
-            mergeMapIntoSingleObj(mapFormValue)
+          processor.ml_inference.input_map = input_map.map(
+            (mapFormValue: MapFormValue) => mergeMapIntoSingleObj(mapFormValue)
           );
         }
 
         if (output_map?.length > 0) {
-          processor.ml_inference.output_map = output_map.map((mapFormValue) =>
-            mergeMapIntoSingleObj(mapFormValue)
+          processor.ml_inference.output_map = output_map.map(
+            (mapFormValue: MapFormValue) => mergeMapIntoSingleObj(mapFormValue)
           );
         }
+
+        // process optional fields
+        let additionalFormValues = {} as FormikValues;
+        Object.keys(formValues).forEach((formKey: string) => {
+          const formValue = formValues[formKey];
+          additionalFormValues = optionallyAddToFinalForm(
+            additionalFormValues,
+            formKey,
+            formValue
+          );
+        });
+
+        // process model config.
+        // TODO: this special handling, plus the special handling on index settings/mappings
+        // could be improved if the 'json' obj returned {} during the conversion instead
+        // of "{}". We may have future JSON fields which right now are going to require
+        // this manual parsing before adding to the template.
+        let finalModelConfig = {};
+        try {
+          // @ts-ignore
+          finalModelConfig = JSON.parse(model_config);
+        } catch (e) {}
+
+        processor.ml_inference = {
+          ...processor.ml_inference,
+          ...additionalFormValues,
+          model_config: finalModelConfig,
+        };
+
         processorsList.push(processor);
         break;
       }