Skip to content

Commit

Permalink
Add optional prompt template transformation in input transform (#379)
Browse files Browse the repository at this point in the history 
* Add toggle-able prompt template transformation in input transform

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

* auto toggling based on parameters empty or not; var name updates

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

* remove TODO

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

---------

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
  • Loading branch information
@ohltyler
ohltyler authored Sep 16, 2024
1 parent 4e38e8c commit 50cd402
Show file tree
Hide file tree
Showing 10 changed files with 188 additions and 52 deletions.
8 changes: 4 additions & 4 deletions public/configs/ml_processor.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,6 @@ export abstract class MLProcessor extends Processor {
},
];
this.optionalFields = [
{
id: 'description',
type: 'string',
},
{
id: 'model_config',
type: 'json',
Expand Down Expand Up @@ -62,6 +58,10 @@ export abstract class MLProcessor extends Processor {
id: 'tag',
type: 'string',
},
{
id: 'description',
type: 'string',
},
];
}
}
2 changes: 2 additions & 0 deletions public/configs/search_response_processors/collapse_processor.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
2 changes: 2 additions & 0 deletions public/configs/search_response_processors/normalization_processor.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [];
Expand Down
1 change: 1 addition & 0 deletions public/pages/workflow_detail/tools/ingest/ingest.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ export function Ingest(props: IngestProps) {
setOptions={{
fontSize: '12px',
autoScrollEditorIntoView: true,
wrap: true,
}}
tabSize={2}
/>
Expand Down
1 change: 1 addition & 0 deletions public/pages/workflow_detail/tools/query/query.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ export function Query(props: QueryProps) {
setOptions={{
fontSize: '12px',
autoScrollEditorIntoView: true,
wrap: true,
}}
tabSize={2}
/>
Expand Down
1 change: 1 addition & 0 deletions public/pages/workflow_detail/workflow_inputs/input_fields/json_field.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
175 changes: 150 additions & 25 deletions public/pages/workflow_detail/workflow_inputs/processor_inputs/input_transform_modal.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import {
EuiCodeBlock,
EuiPopoverTitle,
EuiIconTip,
EuiSwitch,
} from '@elastic/eui';
import {
IConfigField,
Expand Down Expand Up @@ -64,6 +65,7 @@ import { MapArrayField } from '../input_fields';
interface InputTransformModalProps {
uiConfig: WorkflowConfig;
config: IProcessorConfig;
baseConfigPath: string;
context: PROCESSOR_CONTEXT;
inputMapField: IConfigField;
inputMapFieldPath: string;
Expand All @@ -72,10 +74,6 @@ interface InputTransformModalProps {
onClose: () => 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
*/
Expand All @@ -84,24 +82,32 @@ 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);

// source input / transformed output state
// source input / transformed input state
const [sourceInput, setSourceInput] = useState<string>('[]');
const [transformedOutput, setTransformedOutput] = useState<string>('{}');
const [transformedInput, setTransformedInput] = useState<string>('{}');

// get the current input map
const map = getIn(values, props.inputMapFieldPath) as MapArrayFormValue;

// selected output state
const outputOptions = map.map((_, idx) => ({
// selected transform state
const transformOptions = map.map((_, idx) => ({
value: idx,
text: `Prediction ${idx + 1}`,
})) as EuiSelectOption[];
const [selectedOutputOption, setSelectedOutputOption] = useState<
const [selectedTransformOption, setSelectedTransformOption] = useState<
number | undefined
>((outputOptions[0]?.value as number) ?? undefined);
>((transformOptions[0]?.value as number) ?? undefined);

// popover state containing the model interface details, if applicable
const [popoverOpen, setPopoverOpen] = useState<boolean>(false);
Expand All @@ -115,19 +121,19 @@ export function InputTransformModal(props: InputTransformModalProps) {
if (
!isEmpty(map) &&
!isEmpty(JSON.parse(sourceInput)) &&
selectedOutputOption !== undefined
selectedTransformOption !== undefined
) {
let sampleSourceInput = {};
try {
sampleSourceInput = JSON.parse(sourceInput);
const output = generateTransform(
sampleSourceInput,
map[selectedOutputOption]
map[selectedTransformOption]
);
setTransformedOutput(customStringify(output));
setTransformedInput(customStringify(output));
} catch {}
}
}, [map, sourceInput, selectedOutputOption]);
}, [map, sourceInput, selectedTransformOption]);

// hook to re-determine validity when the generated output changes
// utilize Ajv JSON schema validator library. For more info/examples, see
Expand All @@ -140,11 +146,44 @@ export function InputTransformModal(props: InputTransformModalProps) {
const validateFn = new Ajv().compile(
props.modelInterface?.input?.properties?.parameters || {}
);
setIsValid(validateFn(JSON.parse(transformedOutput)));
setIsValid(validateFn(JSON.parse(transformedInput)));
} else {
setIsValid(undefined);
}
}, [transformedOutput]);
}, [transformedInput]);

// 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(() => {
const transformedInputObj = JSON.parse(transformedInput);
if (!isEmpty(originalPrompt) && !isEmpty(transformedInputObj)) {
setTransformedPrompt(
injectValuesIntoPrompt(originalPrompt, transformedInputObj)
);
setViewPromptDetails(true);
setViewTransformedPrompt(true);
} else {
setViewPromptDetails(false);
setViewTransformedPrompt(false);
setTransformedPrompt(originalPrompt);
}
}, [originalPrompt, transformedInput]);

return (
<EuiModal onClose={props.onClose} style={{ width: '70vw' }}>
Expand Down Expand Up @@ -303,6 +342,7 @@ export function InputTransformModal(props: InputTransformModalProps) {
showLineNumbers: false,
showGutter: false,
showPrintMargin: false,
wrap: true,
}}
tabSize={2}
/>
Expand Down Expand Up @@ -330,15 +370,15 @@ export function InputTransformModal(props: InputTransformModalProps) {
// If the map we are adding is the first one, populate the selected option to index 0
onMapAdd={(curArray) => {
if (isEmpty(curArray)) {
setSelectedOutputOption(0);
setSelectedTransformOption(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('{}');
if (selectedTransformOption === idxToDelete) {
setSelectedTransformOption(0);
setTransformedInput('{}');
}
}}
/>
Expand Down Expand Up @@ -369,15 +409,15 @@ export function InputTransformModal(props: InputTransformModalProps) {
</EuiFlexItem>
)}
<EuiFlexItem grow={true}>
{outputOptions.length <= 1 ? (
{transformOptions.length <= 1 ? (
<EuiText>Transformed input</EuiText>
) : (
<EuiCompressedSelect
prepend={<EuiText>Transformed input for</EuiText>}
options={outputOptions}
value={selectedOutputOption}
options={transformOptions}
value={selectedTransformOption}
onChange={(e) => {
setSelectedOutputOption(Number(e.target.value));
setSelectedTransformOption(Number(e.target.value));
}}
/>
)}
Expand Down Expand Up @@ -417,19 +457,80 @@ export function InputTransformModal(props: InputTransformModalProps) {
theme="textmate"
width="100%"
height="15vh"
value={transformedOutput}
value={transformedInput}
readOnly={true}
setOptions={{
fontSize: '12px',
autoScrollEditorIntoView: true,
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)}
disabled={isEmpty(JSON.parse(transformedInput))}
/>
</EuiFlexItem>
{isEmpty(JSON.parse(transformedInput)) && (
<EuiFlexItem grow={false} style={{ marginTop: '16px' }}>
<EuiText size="s" color="subdued">
Transformed input is empty
</EuiText>
</EuiFlexItem>
)}
</EuiFlexGroup>
{viewPromptDetails && (
<>
<EuiSpacer size="s" />
<EuiSwitch
label="With transformed 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>
Expand All @@ -440,3 +541,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;
}
1 change: 1 addition & 0 deletions public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
Loading

0 comments on commit 50cd402

Please sign in to comment.