Add toggle-able prompt template transformation in input transform

Signed-off-by: Tyler Ohlsen <[email protected]>
opensearch-project · Sep 13, 2024 · ed3ec58 · ed3ec58
1 parent 4e38e8c
commit ed3ec58
Show file tree

Hide file tree

Showing 10 changed files with 129 additions and 4 deletions.
diff --git a/public/configs/ml_processor.ts b/public/configs/ml_processor.ts
@@ -30,10 +30,6 @@ export abstract class MLProcessor extends Processor {
       },
     ];
     this.optionalFields = [
-      {
-        id: 'description',
-        type: 'string',
-      },
       {
         id: 'model_config',
         type: 'json',
@@ -62,6 +58,10 @@ export abstract class MLProcessor extends Processor {
         id: 'tag',
         type: 'string',
       },
+      {
+        id: 'description',
+        type: 'string',
+      },
     ];
   }
 }
diff --git a/public/configs/search_response_processors/collapse_processor.ts b/public/configs/search_response_processors/collapse_processor.ts
@@ -5,13 +5,15 @@
 
 import { PROCESSOR_TYPE } from '../../../common';
 import { Processor } from '../processor';
+import { generateId } from '../../utils';
 
 /**
  * The collapse processor config. Used in search flows.
  */
 export class CollapseProcessor extends Processor {
   constructor() {
     super();
+    this.id = generateId('collapse_processor');
     this.type = PROCESSOR_TYPE.COLLAPSE;
     this.name = 'Collapse Processor';
     this.fields = [

diff --git a/public/configs/search_response_processors/normalization_processor.ts b/public/configs/search_response_processors/normalization_processor.ts
@@ -5,13 +5,15 @@
 
 import { PROCESSOR_TYPE } from '../../../common';
 import { Processor } from '../processor';
+import { generateId } from '../../utils';
 
 /**
  * The normalization processor config. Used in search flows.
  */
 export class NormalizationProcessor extends Processor {
   constructor() {
     super();
+    this.id = generateId('normalization_processor');
     this.type = PROCESSOR_TYPE.NORMALIZATION;
     this.name = 'Normalization Processor';
     this.fields = [];

diff --git a/public/pages/workflow_detail/tools/ingest/ingest.tsx b/public/pages/workflow_detail/tools/ingest/ingest.tsx
@@ -29,6 +29,7 @@ export function Ingest(props: IngestProps) {
       setOptions={{
         fontSize: '12px',
         autoScrollEditorIntoView: true,
+        wrap: true,
       }}
       tabSize={2}
     />

diff --git a/public/pages/workflow_detail/tools/query/query.tsx b/public/pages/workflow_detail/tools/query/query.tsx
@@ -29,6 +29,7 @@ export function Query(props: QueryProps) {
       setOptions={{
         fontSize: '12px',
         autoScrollEditorIntoView: true,
+        wrap: true,
       }}
       tabSize={2}
     />

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
@@ -99,6 +99,7 @@ export function JsonField(props: JsonFieldProps) {
                 highlightActiveLine: !props.readOnly,
                 highlightSelectedWord: !props.readOnly,
                 highlightGutterLine: !props.readOnly,
+                wrap: true,
               }}
               aria-label="Code Editor"
               tabSize={2}

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
@@ -26,6 +26,7 @@ import {
   EuiCodeBlock,
   EuiPopoverTitle,
   EuiIconTip,
+  EuiSwitch,
 } from '@elastic/eui';
 import {
   IConfigField,
@@ -64,6 +65,7 @@ import { MapArrayField } from '../input_fields';
 interface InputTransformModalProps {
   uiConfig: WorkflowConfig;
   config: IProcessorConfig;
+  baseConfigPath: string;
   context: PROCESSOR_CONTEXT;
   inputMapField: IConfigField;
   inputMapFieldPath: string;
@@ -84,6 +86,14 @@ export function InputTransformModal(props: InputTransformModalProps) {
   const dataSourceId = getDataSourceId();
   const { values } = useFormikContext<WorkflowFormValues>();
 
+  // various prompt states
+  const [viewPromptDetails, setViewPromptDetails] = useState<boolean>(false);
+  const [viewTransformedPrompt, setViewTransformedPrompt] = useState<boolean>(
+    false
+  );
+  const [originalPrompt, setOriginalPrompt] = useState<string>('');
+  const [transformedPrompt, setTransformedPrompt] = useState<string>('');
+
   // fetching input data state
   const [isFetching, setIsFetching] = useState<boolean>(false);
 
@@ -146,6 +156,32 @@ export function InputTransformModal(props: InputTransformModalProps) {
     }
   }, [transformedOutput]);
 
+  // hook to set the prompt if found in the model config
+  useEffect(() => {
+    const modelConfigString = getIn(
+      values,
+      `${props.baseConfigPath}.${props.config.id}.model_config`
+    );
+    try {
+      const prompt = JSON.parse(modelConfigString)?.prompt;
+      if (!isEmpty(prompt)) {
+        setOriginalPrompt(prompt);
+      }
+    } catch {}
+  }, [
+    getIn(values, `${props.baseConfigPath}.${props.config.id}.model_config`),
+  ]);
+
+  // hook to set the transformed prompt, if a valid prompt found, and
+  // valid parameters set
+  useEffect(() => {
+    if (!isEmpty(originalPrompt)) {
+      setTransformedPrompt(
+        injectValuesIntoPrompt(originalPrompt, JSON.parse(transformedOutput))
+      );
+    }
+  }, [originalPrompt, transformedOutput]);
+
   return (
     <EuiModal onClose={props.onClose} style={{ width: '70vw' }}>
       <EuiModalHeader>
@@ -303,6 +339,7 @@ export function InputTransformModal(props: InputTransformModalProps) {
                   showLineNumbers: false,
                   showGutter: false,
                   showPrintMargin: false,
+                  wrap: true,
                 }}
                 tabSize={2}
               />
@@ -425,11 +462,64 @@ export function InputTransformModal(props: InputTransformModalProps) {
                   showLineNumbers: false,
                   showGutter: false,
                   showPrintMargin: false,
+                  wrap: true,
                 }}
                 tabSize={2}
               />
             </>
           </EuiFlexItem>
+          {!isEmpty(originalPrompt) && (
+            <EuiFlexItem>
+              <>
+                <EuiFlexGroup direction="row">
+                  <EuiFlexItem grow={false}>
+                    <EuiText>Transformed prompt</EuiText>
+                  </EuiFlexItem>
+                  <EuiFlexItem grow={false} style={{ marginTop: '16px' }}>
+                    <EuiSwitch
+                      label="Show"
+                      checked={viewPromptDetails}
+                      onChange={() => setViewPromptDetails(!viewPromptDetails)}
+                    />
+                  </EuiFlexItem>
+                </EuiFlexGroup>
+                {viewPromptDetails && (
+                  <>
+                    <EuiSpacer size="s" />
+                    <EuiSwitch
+                      label="With injected inputs"
+                      checked={viewTransformedPrompt}
+                      onChange={() =>
+                        setViewTransformedPrompt(!viewTransformedPrompt)
+                      }
+                    />
+                    <EuiSpacer size="m" />
+                    <EuiCodeEditor
+                      mode="json"
+                      theme="textmate"
+                      width="100%"
+                      height="15vh"
+                      value={
+                        viewTransformedPrompt
+                          ? transformedPrompt
+                          : originalPrompt
+                      }
+                      readOnly={true}
+                      setOptions={{
+                        fontSize: '12px',
+                        autoScrollEditorIntoView: true,
+                        showLineNumbers: false,
+                        showGutter: false,
+                        showPrintMargin: false,
+                        wrap: true,
+                      }}
+                      tabSize={2}
+                    />
+                  </>
+                )}
+              </>
+            </EuiFlexItem>
+          )}
         </EuiFlexGroup>
       </EuiModalBody>
       <EuiModalFooter>
@@ -440,3 +530,27 @@ export function InputTransformModal(props: InputTransformModalProps) {
     </EuiModal>
   );
 }
+
+function injectValuesIntoPrompt(
+  promptString: string,
+  parameters: { [key: string]: string }
+): string {
+  let finalPromptString = promptString;
+  // replace any parameter placeholders in the prompt with any values found in the
+  // parameters obj.
+  // we do 2 checks - one for the regular prompt, and one with "toString()" appended.
+  // this is required for parameters that have values as a list, for example.
+  Object.keys(parameters).forEach((parameterKey) => {
+    const parameterValue = parameters[parameterKey];
+    const regex = new RegExp(`\\$\\{parameters.${parameterKey}\\}`, 'g');
+    const regexWithToString = new RegExp(
+      `\\$\\{parameters.${parameterKey}.toString\\(\\)\\}`,
+      'g'
+    );
+    finalPromptString = finalPromptString
+      .replace(regex, parameterValue)
+      .replace(regexWithToString, parameterValue);
+  });
+
+  return finalPromptString;
+}
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
@@ -209,6 +209,7 @@ export function MLProcessorInputs(props: MLProcessorInputsProps) {
         <InputTransformModal
           uiConfig={props.uiConfig}
           config={props.config}
+          baseConfigPath={props.baseConfigPath}
           context={props.context}
           inputMapField={inputMapField}
           inputMapFieldPath={inputMapFieldPath}

diff --git a/public/pages/workflow_detail/workflow_inputs/processor_inputs/output_transform_modal.tsx b/public/pages/workflow_detail/workflow_inputs/processor_inputs/output_transform_modal.tsx
@@ -258,6 +258,7 @@ export function OutputTransformModal(props: OutputTransformModalProps) {
                   showLineNumbers: false,
                   showGutter: false,
                   showPrintMargin: false,
+                  wrap: true,
                 }}
                 tabSize={2}
               />
@@ -322,6 +323,7 @@ export function OutputTransformModal(props: OutputTransformModalProps) {
                   showLineNumbers: false,
                   showGutter: false,
                   showPrintMargin: false,
+                  wrap: true,
                 }}
                 tabSize={2}
               />

diff --git a/public/pages/workflow_detail/workspace/workspace.tsx b/public/pages/workflow_detail/workspace/workspace.tsx
@@ -214,6 +214,7 @@ export function Workspace(props: WorkspaceProps) {
                 setOptions={{
                   fontSize: '12px',
                   autoScrollEditorIntoView: true,
+                  wrap: true,
                 }}
                 tabSize={2}
               />