From 43e0ceb342cd27ee94e7c1ac830ee642294e2f2f Mon Sep 17 00:00:00 2001
From: HugoCasa <hugo@casademont.ch>
Date: Tue, 13 Feb 2024 15:52:03 +0100
Subject: [PATCH] feat: add ai for predicates and iterator expressions (#3203)

* feat: add ai for predicates and iterator expressions

* fix: npm run check

* feat: copilot like suggestion
---
 frontend/src/lib/components/ArgInput.svelte   |   2 +
 frontend/src/lib/components/Editor.svelte     |   4 -
 .../lib/components/InputTransformForm.svelte  |  12 +-
 .../InputTransformSchemaForm.svelte           |   1 +
 .../src/lib/components/SimpleEditor.svelte    |  36 +++-
 .../copilot/FlowCopilotInputsModal.svelte     |   2 +-
 .../lib/components/copilot/IteratorGen.svelte | 181 ++++++++++++++++++
 .../components/copilot/PredicateGen.svelte    | 135 +++++++++++++
 .../components/copilot/StepInputGen.svelte    | 113 +++++++++--
 .../components/copilot/StepInputsGen.svelte   |  84 ++++++--
 .../lib/components/flows/FlowEditor.svelte    |   2 +-
 .../content/BranchPredicateEditor.svelte      |  15 +-
 .../flows/content/FlowBranchOneWrapper.svelte |   3 +-
 .../content/FlowBranchesOneWrapper.svelte     |   8 +-
 .../flows/content/FlowEditorPanel.svelte      |   2 +
 .../components/flows/content/FlowLoop.svelte  |  52 ++++-
 .../flows/content/FlowModuleWrapper.svelte    |  17 +-
 .../lib/components/flows/flowStateUtils.ts    |   2 +-
 frontend/src/lib/components/theme/nord.json   |  93 ---------
 19 files changed, 607 insertions(+), 157 deletions(-)
 create mode 100644 frontend/src/lib/components/copilot/IteratorGen.svelte
 create mode 100644 frontend/src/lib/components/copilot/PredicateGen.svelte
 delete mode 100644 frontend/src/lib/components/theme/nord.json

diff --git a/frontend/src/lib/components/ArgInput.svelte b/frontend/src/lib/components/ArgInput.svelte
index a15db45209..a3b4b801ae 100644
--- a/frontend/src/lib/components/ArgInput.svelte
+++ b/frontend/src/lib/components/ArgInput.svelte
@@ -340,6 +340,7 @@
 						<input
 							{autofocus}
 							on:focus
+							on:blur
 							{disabled}
 							type="number"
 							on:keydown={() => {
@@ -608,6 +609,7 @@
 				<input
 					{autofocus}
 					on:focus
+					on:blur
 					{disabled}
 					type="email"
 					class={valid
diff --git a/frontend/src/lib/components/Editor.svelte b/frontend/src/lib/components/Editor.svelte
index 468965eb8f..c121b00d9e 100644
--- a/frontend/src/lib/components/Editor.svelte
+++ b/frontend/src/lib/components/Editor.svelte
@@ -28,10 +28,6 @@
 	import libStdContent from '$lib/es6.d.ts.txt?raw'
 	import denoFetchContent from '$lib/deno_fetch.d.ts.txt?raw'
 
-	// import nord from '$lib/assets/nord.json'
-
-	// import nord from '$lib/assets/nord.json'
-
 	import { MonacoLanguageClient } from 'monaco-languageclient'
 
 	import { toSocket, WebSocketMessageReader, WebSocketMessageWriter } from 'vscode-ws-jsonrpc'
diff --git a/frontend/src/lib/components/InputTransformForm.svelte b/frontend/src/lib/components/InputTransformForm.svelte
index 22451ed49b..f8144470dc 100644
--- a/frontend/src/lib/components/InputTransformForm.svelte
+++ b/frontend/src/lib/components/InputTransformForm.svelte
@@ -218,13 +218,22 @@
 							bind:this={stepInputGen}
 							{focused}
 							{arg}
+							schemaProperty={schema.properties[argName]}
+							showPopup={(isStaticTemplate(inputCat) && propertyType == 'static') ||
+								propertyType === undefined ||
+								propertyType === 'static' ||
+								arg?.expr?.length > 0}
+							on:showExpr={(e) => {
+								monaco?.setSuggestion(e.detail)
+							}}
 							on:setExpr={(e) => {
 								arg = {
 									type: 'javascript',
 									expr: e.detail
 								}
 								propertyType = 'javascript'
-								monaco?.setCode(e.detail)
+								monaco?.setCode('')
+								monaco?.insertAtCursor(e.detail)
 							}}
 							{pickableProperties}
 							{argName}
@@ -395,6 +404,7 @@
 							focused = false
 						}}
 						autoHeight
+						preventTabOnEmpty={enableAi}
 					/>
 				</div>
 				<DynamicInputHelpBox />
diff --git a/frontend/src/lib/components/InputTransformSchemaForm.svelte b/frontend/src/lib/components/InputTransformSchemaForm.svelte
index deb74486aa..ae2a9dad62 100644
--- a/frontend/src/lib/components/InputTransformSchemaForm.svelte
+++ b/frontend/src/lib/components/InputTransformSchemaForm.svelte
@@ -74,6 +74,7 @@
 								(args[argName].type === 'javascript' && !args[argName].expr))
 				  )
 				: []}
+			{schema}
 		/>
 	{/if}
 	{#if keys.length > 0}
diff --git a/frontend/src/lib/components/SimpleEditor.svelte b/frontend/src/lib/components/SimpleEditor.svelte
index 3d8022afaf..e15d008439 100644
--- a/frontend/src/lib/components/SimpleEditor.svelte
+++ b/frontend/src/lib/components/SimpleEditor.svelte
@@ -13,8 +13,7 @@
 		KeyMod,
 		Uri as mUri,
 		languages,
-		type IRange,
-		type IKeyboardEvent
+		type IRange
 	} from 'monaco-editor'
 	import 'monaco-editor/esm/vs/basic-languages/sql/sql.contribution'
 	import 'monaco-editor/esm/vs/basic-languages/yaml/yaml.contribution'
@@ -57,6 +56,7 @@
 	export let fixedOverflowWidgets = true
 	export let small = false
 	export let domLib = false
+	export let autofocus = false
 
 	const dispatch = createEventDispatcher()
 
@@ -117,12 +117,6 @@
 		}
 	}
 
-	export function onKeyUp(f: (e: IKeyboardEvent) => void) {
-		if (editor) {
-			return editor.onKeyUp(f)
-		}
-	}
-
 	export function show(): void {
 		divEl?.classList.remove('hidden')
 	}
@@ -131,9 +125,17 @@
 		divEl?.classList.add('hidden')
 	}
 
+	let suggestion = ''
+	export function setSuggestion(value: string): void {
+		suggestion = value
+	}
+
 	let width = 0
 	let initialized = false
 
+	let disableTabCond: meditor.IContextKey<boolean> | undefined
+	$: disableTabCond?.set(!code && !!suggestion)
+
 	async function loadMonaco() {
 		await initializeVscode()
 		initialized = true
@@ -184,6 +186,7 @@
 
 		let timeoutModel: NodeJS.Timeout | undefined = undefined
 		editor.onDidChangeModelContent((event) => {
+			suggestion = ''
 			timeoutModel && clearTimeout(timeoutModel)
 			timeoutModel = setTimeout(() => {
 				code = getCode()
@@ -198,6 +201,9 @@
 				code = getCode()
 				shouldBindKey && format && format()
 			})
+
+			disableTabCond = editor.createContextKey('disableTabCond', !code)
+			editor.addCommand(KeyCode.Tab, function () {}, 'disableTabCond')
 		})
 
 		if (autoHeight) {
@@ -332,6 +338,11 @@
 		if (BROWSER) {
 			mounted = true
 			await loadMonaco()
+			if (autofocus) {
+				setTimeout(() => {
+					focus()
+				}, 0)
+			}
 		}
 	})
 
@@ -347,6 +358,15 @@
 
 <EditorTheme />
 
+{#if editor && suggestion && code.length === 0}
+	<div
+		class="absolute top-[0.05rem] left-[2.05rem] z-10 text-sm text-[#0007] italic font-mono dark:text-[#ffffff56] text-ellipsis overflow-hidden whitespace-nowrap"
+		style={`max-width: calc(${width}px - 2.05rem)`}
+	>
+		{suggestion}
+	</div>
+{/if}
+
 <div bind:this={divEl} class="{$$props.class ?? ''} editor" bind:clientWidth={width} />
 
 <style lang="postcss">
diff --git a/frontend/src/lib/components/copilot/FlowCopilotInputsModal.svelte b/frontend/src/lib/components/copilot/FlowCopilotInputsModal.svelte
index edb2be8db4..35553c70f4 100644
--- a/frontend/src/lib/components/copilot/FlowCopilotInputsModal.svelte
+++ b/frontend/src/lib/components/copilot/FlowCopilotInputsModal.svelte
@@ -19,7 +19,7 @@
 	on:canceled
 	title="Windmill AI wants to add the following inputs to the flow:"
 >
-	<ul class="text-lg list-disc pl-5">
+	<ul class=" list-disc pl-5">
 		{#each inputs as input}
 			<li>{input}</li>
 		{/each}
diff --git a/frontend/src/lib/components/copilot/IteratorGen.svelte b/frontend/src/lib/components/copilot/IteratorGen.svelte
new file mode 100644
index 0000000000..140672620e
--- /dev/null
+++ b/frontend/src/lib/components/copilot/IteratorGen.svelte
@@ -0,0 +1,181 @@
+<script lang="ts">
+	import { Check, Loader2, Wand2 } from 'lucide-svelte'
+	import Button from '../common/button/Button.svelte'
+	import { getNonStreamingCompletion } from './lib'
+	import { sendUserToast } from '$lib/toast'
+	import type { InputTransform } from '$lib/gen'
+	import ManualPopover from '../ManualPopover.svelte'
+	import { createEventDispatcher, getContext } from 'svelte'
+	import type { FlowEditorContext } from '../flows/types'
+	import type { PickableProperties } from '../flows/previousResults'
+	import YAML from 'yaml'
+	import { sliceModules } from '../flows/flowStateUtils'
+	import { dfs } from '../flows/dfs'
+	import { yamlStringifyExceptKeys } from './utils'
+	import { copilotInfo, stepInputCompletionEnabled } from '$lib/stores'
+
+	let generatedContent = ''
+	let loading = false
+	export let focused = false
+	export let arg: InputTransform | any
+	export let pickableProperties: PickableProperties | undefined = undefined
+
+	let btnFocused = false
+	let empty = false
+	$: empty =
+		Object.keys(arg ?? {}).length === 0 ||
+		(arg.type === 'static' && !arg.value) ||
+		(arg.type === 'javascript' && !arg.expr)
+
+	let abortController = new AbortController()
+	const { flowStore, selectedId } = getContext<FlowEditorContext>('FlowEditorContext')
+
+	async function generateIteratorExpr() {
+		if (generatedContent.length > 0 || loading) {
+			return
+		}
+		abortController = new AbortController()
+		loading = true
+		const idOrders = dfs($flowStore.value.modules, (x) => x.id)
+		const upToIndex = idOrders.indexOf($selectedId)
+		if (upToIndex === -1) {
+			throw new Error('Could not find the selected id in the flow')
+		}
+
+		const flowDetails =
+			'Take into account the following information for never tested results:\n<flowDetails>\n' +
+			yamlStringifyExceptKeys(sliceModules($flowStore.value.modules, upToIndex, idOrders), [
+				'lock'
+			]) +
+			'</flowDetails>'
+		try {
+			const availableData = {
+				results: pickableProperties?.priorIds,
+				flow_input: pickableProperties?.flow_input
+			}
+			const user = `I'm building a workflow which is a DAG of script steps.
+The current step is ${selectedId} and represents a for-loop. You can find the details of all the steps below:
+${flowDetails}
+Determine the iterator expression to pass either from the previous results or the flow inputs. Here's a summary of the available data:
+<available>
+${YAML.stringify(availableData)}</available>
+Reply with the most probable answer, do not explain or discuss.
+Use javascript object dot notation to access the properties.
+Only output the expression, do not explain or discuss.`
+
+			generatedContent = await getNonStreamingCompletion(
+				[
+					{
+						role: 'user',
+						content: user
+					}
+				],
+				abortController
+			)
+		} catch (err) {
+			if (!abortController.signal.aborted) {
+				sendUserToast('Could not generate summary: ' + err, true)
+			}
+		} finally {
+			loading = false
+		}
+	}
+
+	export function onKeyUp(event: KeyboardEvent) {
+		if (!$copilotInfo.exists_openai_resource_path || !$stepInputCompletionEnabled) {
+			return
+		}
+		if (event.key === 'Tab') {
+			if (!loading && generatedContent) {
+				event.preventDefault()
+				dispatch('setExpr', generatedContent)
+				generatedContent = ''
+			}
+		} else {
+			cancel()
+		}
+	}
+
+	const dispatch = createEventDispatcher()
+
+	function automaticGeneration() {
+		if (empty) {
+			generateIteratorExpr()
+		}
+	}
+
+	function cancelOnOutOfFocus() {
+		setTimeout(() => {
+			if (!focused && !btnFocused) {
+				// only cancel if out of focus is not due to click on btn
+				cancel()
+			}
+		}, 150)
+	}
+
+	$: if (!focused) {
+		cancelOnOutOfFocus()
+	}
+
+	$: if ($copilotInfo.exists_openai_resource_path && $stepInputCompletionEnabled && focused) {
+		automaticGeneration()
+	}
+
+	function cancel() {
+		abortController.abort()
+		generatedContent = ''
+	}
+
+	$: dispatch('showExpr', generatedContent)
+
+	let out = true // hack to prevent regenerating answer when accepting the answer due to mouseenter on new icon
+</script>
+
+{#if $copilotInfo.exists_openai_resource_path && $stepInputCompletionEnabled}
+	<ManualPopover showTooltip={!empty && generatedContent.length > 0} placement="bottom" class="p-2">
+		<Button
+			size="xs"
+			color="light"
+			btnClasses="text-violet-800 dark:text-violet-400 bg-violet-100 dark:bg-gray-700 dark:hover:bg-surface-hover"
+			on:click={() => {
+				if (!loading && generatedContent.length > 0) {
+					dispatch('setExpr', generatedContent)
+					generatedContent = ''
+				}
+			}}
+			on:focus={() => {
+				btnFocused = true
+			}}
+			on:blur={() => {
+				btnFocused = false
+			}}
+			on:mouseenter={(ev) => {
+				if (out) {
+					out = false
+					generateIteratorExpr()
+				}
+			}}
+			on:mouseleave={() => {
+				out = true
+				cancel()
+			}}
+			endIcon={{
+				icon: loading ? Loader2 : generatedContent.length > 0 ? Check : Wand2,
+				classes: loading ? 'animate-spin' : ''
+			}}
+		>
+			{#if focused}
+				{#if loading}
+					ESC
+				{:else if generatedContent.length > 0}
+					TAB
+				{/if}
+			{/if}
+		</Button>
+		<svelte:fragment slot="content">
+			<div class="text-sm text-tertiary">
+				{generatedContent}
+			</div>
+		</svelte:fragment>
+	</ManualPopover>
+{/if}
diff --git a/frontend/src/lib/components/copilot/PredicateGen.svelte b/frontend/src/lib/components/copilot/PredicateGen.svelte
new file mode 100644
index 0000000000..7177898d37
--- /dev/null
+++ b/frontend/src/lib/components/copilot/PredicateGen.svelte
@@ -0,0 +1,135 @@
+<script lang="ts">
+	import { Wand2 } from 'lucide-svelte'
+	import Button from '../common/button/Button.svelte'
+	import { getNonStreamingCompletion } from './lib'
+	import { sendUserToast } from '$lib/toast'
+	import { createEventDispatcher, getContext } from 'svelte'
+	import type { FlowEditorContext } from '../flows/types'
+	import type { PickableProperties } from '../flows/previousResults'
+	import YAML from 'yaml'
+	import { sliceModules } from '../flows/flowStateUtils'
+	import { dfs } from '../flows/dfs'
+	import { yamlStringifyExceptKeys } from './utils'
+	import { copilotInfo, stepInputCompletionEnabled } from '$lib/stores'
+	import Popup from '../common/popup/Popup.svelte'
+
+	let loading = false
+	export let pickableProperties: PickableProperties | undefined = undefined
+
+	let instructions = ''
+	let instructionsField: HTMLInputElement | undefined = undefined
+	$: instructionsField && setTimeout(() => instructionsField?.focus(), 100)
+
+	let abortController = new AbortController()
+	const { flowStore, selectedId } = getContext<FlowEditorContext>('FlowEditorContext')
+
+	const dispatch = createEventDispatcher()
+
+	async function generatePredicate() {
+		abortController = new AbortController()
+		loading = true
+		const idOrders = dfs($flowStore.value.modules, (x) => x.id)
+		const upToIndex = idOrders.indexOf($selectedId)
+		if (upToIndex === -1) {
+			throw new Error('Could not find the selected id in the flow')
+		}
+
+		const flowDetails =
+			'Take into account the following information for never tested results:\n<flowDetails>\n' +
+			yamlStringifyExceptKeys(sliceModules($flowStore.value.modules, upToIndex, idOrders), [
+				'lock'
+			]) +
+			'</flowDetails>'
+		try {
+			const availableData = {
+				results: pickableProperties?.priorIds,
+				flow_input: pickableProperties?.flow_input
+			}
+			const user = `I'm building a workflow which is a DAG of script steps.
+The current step is ${selectedId} and is a branching step (if-else). 
+The user wants to generate a predicate for the branching condition.
+Here's the user's request: ${instructions}
+You can find the details of all the steps below:
+${flowDetails}
+
+Determine for the user the javascript expression for the branching condition composed of the previous results or the flow inputs.
+All inputs start with either results. or flow_input. and are followed by the key of the input.
+Here's a summary of the available data:
+<available>
+${YAML.stringify(availableData)}</available>
+If the branching is made inside a for-loop, the iterator value is accessible as flow_input.iter.value
+Only return the expression without any wrapper. Do not explain or discuss.`
+
+			const result = await getNonStreamingCompletion(
+				[
+					{
+						role: 'user',
+						content: user
+					}
+				],
+				abortController
+			)
+
+			dispatch('setExpr', result)
+		} catch (err) {
+			if (!abortController.signal.aborted) {
+				sendUserToast('Could not generate summary: ' + err, true)
+			}
+		} finally {
+			loading = false
+		}
+	}
+</script>
+
+{#if $copilotInfo.exists_openai_resource_path && $stepInputCompletionEnabled}
+	<Popup
+		floatingConfig={{ strategy: 'absolute', placement: 'bottom-end' }}
+		containerClasses="border rounded-lg shadow-lg p-4 bg-surface"
+		let:close
+	>
+		<svelte:fragment slot="button">
+			<Button
+				color={loading ? 'red' : 'light'}
+				size="xs"
+				nonCaptureEvent={!loading}
+				startIcon={{ icon: Wand2 }}
+				iconOnly
+				title="AI Assistant"
+				btnClasses="min-h-[30px] text-violet-800 dark:text-violet-400 bg-violet-100 dark:bg-gray-700"
+				{loading}
+				clickableWhileLoading
+				on:click={loading ? () => abortController?.abort() : undefined}
+			/>
+		</svelte:fragment>
+		<div class="flex w-96">
+			<input
+				bind:this={instructionsField}
+				type="text"
+				placeholder="Predicate description"
+				bind:value={instructions}
+				on:keypress={({ key }) => {
+					if (key === 'Enter' && instructions.length > 0) {
+						close(instructionsField || null)
+						generatePredicate()
+					}
+				}}
+			/>
+			<Button
+				size="xs"
+				color="light"
+				variant="contained"
+				buttonType="button"
+				btnClasses="!p-1 !w-[38px] !ml-2 text-violet-800 dark:text-violet-400 bg-violet-100 dark:bg-gray-700"
+				title="Generate predicate from prompt"
+				aria-label="Generate"
+				iconOnly
+				on:click={() => {
+					close(instructionsField || null)
+					generatePredicate()
+				}}
+				disabled={instructions.length == 0}
+				startIcon={{ icon: Wand2 }}
+			/>
+		</div>
+	</Popup>
+{/if}
diff --git a/frontend/src/lib/components/copilot/StepInputGen.svelte b/frontend/src/lib/components/copilot/StepInputGen.svelte
index 706c17e3bd..b72dbe2b30 100644
--- a/frontend/src/lib/components/copilot/StepInputGen.svelte
+++ b/frontend/src/lib/components/copilot/StepInputGen.svelte
@@ -14,24 +14,57 @@
 	import { yamlStringifyExceptKeys } from './utils'
 	import type { FlowCopilotContext } from './flow'
 	import { copilotInfo, stepInputCompletionEnabled } from '$lib/stores'
+	import type { SchemaProperty } from '$lib/common'
+	import FlowCopilotInputsModal from './FlowCopilotInputsModal.svelte'
 
 	let generatedContent = ''
 	let loading = false
 	export let focused = false
 	export let arg: InputTransform | any
+	export let schemaProperty: SchemaProperty
 	export let pickableProperties: PickableProperties | undefined = undefined
 	export let argName: string
+	export let showPopup: boolean
 
 	let empty = false
 	$: empty =
-		!arg || (arg.type === 'static' && !arg.value) || (arg.type === 'javascript' && !arg.expr)
+		Object.keys(arg ?? {}).length === 0 ||
+		(arg.type === 'static' && !arg.value) ||
+		(arg.type === 'javascript' && !arg.expr)
+
+	let btnFocused = false
 
 	let abortController = new AbortController()
+	let newFlowInput = ''
+
 	const { flowStore, selectedId } = getContext<FlowEditorContext>('FlowEditorContext')
 	const { stepInputsLoading, generatedExprs } =
 		getContext<FlowCopilotContext | undefined>('FlowCopilotContext') || {}
 
+	function createFlowInput() {
+		if (!newFlowInput) {
+			return
+		}
+		const properties = {
+			...($flowStore.schema?.properties as Record<string, SchemaProperty> | undefined),
+			[newFlowInput]: schemaProperty
+		}
+		const required = [
+			...(($flowStore.schema?.required as string[] | undefined) ?? []),
+			newFlowInput
+		]
+		$flowStore.schema = {
+			$schema: 'https://json-schema.org/draft/2020-12/schema',
+			properties,
+			required,
+			type: 'object'
+		}
+	}
+
 	async function generateStepInput() {
+		if (generatedContent.length > 0 || loading) {
+			return
+		}
 		abortController = new AbortController()
 		loading = true
 		const idOrders = dfs($flowStore.value.modules, (x) => x.id)
@@ -43,8 +76,7 @@
 		const flowDetails =
 			'Take into account the following information for never tested results:\n<flowDetails>\n' +
 			yamlStringifyExceptKeys(sliceModules($flowStore.value.modules, upToIndex, idOrders), [
-				'lock',
-				'input_transforms'
+				'lock'
 			]) +
 			'</flowDetails>'
 		try {
@@ -55,13 +87,15 @@
 			const user = `I'm building a workflow which is a DAG of script steps.
 The current step is ${selectedId}, you can find the details for the step and previous ones below:
 ${flowDetails}
-Determine for the input "${argName}", what to pass either from the previous results of the flow inputs. Here's a summary of the available data:
+Determine for the input "${argName}", what to pass either from the previous results or the flow inputs. 
+All possibles inputs either start with results. or flow_input. and are follow by the key of the input.
+Here's a summary of the available data:
 <available>
 ${YAML.stringify(availableData)}</available>
 If none of the available results are appropriate, are already used or are more appropriate for other inputs, you can also imagine new flow_input properties which we will create programmatically based on what you provide.
 Reply with the most probable answer, do not explain or discuss.
 Use javascript object dot notation to access the properties.
-Return the input element directly: e.g. flow_input.property, results.a, results.a.property, flow_input.iter.value`
+Only return the expression without any wrapper.`
 
 			generatedContent = await getNonStreamingCompletion(
 				[
@@ -72,6 +106,17 @@ Return the input element directly: e.g. flow_input.property, results.a, results.
 				],
 				abortController
 			)
+
+			if (
+				pickableProperties &&
+				generatedContent.startsWith('flow_input.') &&
+				generatedContent.split('.')[1] &&
+				!(generatedContent.split('.')[1] in pickableProperties.flow_input)
+			) {
+				newFlowInput = generatedContent.split('.')[1]
+			} else {
+				newFlowInput = ''
+			}
 		} catch (err) {
 			if (!abortController.signal.aborted) {
 				sendUserToast('Could not generate summary: ' + err, true)
@@ -89,6 +134,9 @@ Return the input element directly: e.g. flow_input.property, results.a, results.
 			if (!loading && generatedContent) {
 				event.preventDefault()
 				dispatch('setExpr', generatedContent)
+				if (newFlowInput) {
+					openInputsModal = true
+				}
 				generatedContent = ''
 			}
 		} else {
@@ -98,30 +146,52 @@ Return the input element directly: e.g. flow_input.property, results.a, results.
 
 	const dispatch = createEventDispatcher()
 
+	function cancel() {
+		abortController.abort()
+		generatedContent = ''
+	}
+
 	function automaticGeneration() {
 		if (empty) {
 			generateStepInput()
 		}
 	}
 
+	function cancelOnOutOfFocus() {
+		setTimeout(() => {
+			if (!focused && !btnFocused) {
+				// only cancel if out of focus is not due to click on btn
+				cancel()
+			}
+		}, 150)
+	}
+
+	$: if (!focused) {
+		cancelOnOutOfFocus()
+	}
+
 	$: if ($copilotInfo.exists_openai_resource_path && $stepInputCompletionEnabled && focused) {
 		automaticGeneration()
 	}
 
-	function cancel() {
-		abortController.abort()
-		generatedContent = ''
-	}
-	$: if (!focused) {
-		cancel()
-	}
+	$: dispatch('showExpr', generatedContent)
+
+	$: dispatch('showExpr', $generatedExprs?.[argName] || '')
 
 	let out = true // hack to prevent regenerating answer when accepting the answer due to mouseenter on new icon
+	let openInputsModal = false
 </script>
 
 {#if $copilotInfo.exists_openai_resource_path && $stepInputCompletionEnabled}
+	<FlowCopilotInputsModal
+		on:confirmed={async () => {
+			createFlowInput()
+		}}
+		bind:open={openInputsModal}
+		inputs={[newFlowInput]}
+	/>
 	<ManualPopover
-		showTooltip={generatedContent.length > 0 || !!$generatedExprs?.[argName]}
+		showTooltip={showPopup && (generatedContent.length > 0 || !!$generatedExprs?.[argName])}
 		placement="bottom"
 		class="p-2"
 	>
@@ -130,17 +200,16 @@ Return the input element directly: e.g. flow_input.property, results.a, results.
 			color="light"
 			btnClasses="text-violet-800 dark:text-violet-400 bg-violet-100 dark:bg-gray-700 dark:hover:bg-surface-hover"
 			on:click={() => {
-				if (loading) {
-					cancel()
-				} else if (generatedContent.length > 0) {
+				if (!loading && generatedContent.length > 0) {
 					dispatch('setExpr', generatedContent)
+					if (newFlowInput) {
+						openInputsModal = true
+					}
 					generatedContent = ''
-				} else {
-					generateStepInput()
 				}
 			}}
 			on:mouseenter={(ev) => {
-				if (!generatedContent && !loading && out) {
+				if (out) {
 					out = false
 					generateStepInput()
 				}
@@ -158,6 +227,12 @@ Return the input element directly: e.g. flow_input.property, results.a, results.
 						: Wand2,
 				classes: loading || ($stepInputsLoading && empty) ? 'animate-spin' : ''
 			}}
+			on:focus={() => {
+				btnFocused = true
+			}}
+			on:blur={() => {
+				btnFocused = false
+			}}
 		>
 			{#if focused}
 				{#if loading}
diff --git a/frontend/src/lib/components/copilot/StepInputsGen.svelte b/frontend/src/lib/components/copilot/StepInputsGen.svelte
index d7866859d9..309540eab2 100644
--- a/frontend/src/lib/components/copilot/StepInputsGen.svelte
+++ b/frontend/src/lib/components/copilot/StepInputsGen.svelte
@@ -13,10 +13,13 @@
 	import { Check, ExternalLink, Loader2, Wand2 } from 'lucide-svelte'
 	import { copilotInfo, stepInputCompletionEnabled } from '$lib/stores'
 	import { Popup } from '../common'
+	import type { SchemaProperty, Schema } from '$lib/common'
+	import FlowCopilotInputsModal from './FlowCopilotInputsModal.svelte'
 
 	let loading = false
 	export let pickableProperties: PickableProperties | undefined = undefined
 	export let argNames: string[] = []
+	export let schema: Schema
 
 	const { flowStore, selectedId } = getContext<FlowEditorContext>('FlowEditorContext')
 
@@ -25,9 +28,13 @@
 
 	let generatedContent = ''
 	let parsedInputs: string[][] = []
+	let newFlowInputs: string[] = []
 
 	let abortController = new AbortController()
 	async function generateStepInputs() {
+		if (Object.keys($generatedExprs || {}).length > 0 || loading) {
+			return
+		}
 		abortController = new AbortController()
 		loading = true
 		stepInputsLoading?.set(true)
@@ -39,8 +46,7 @@
 		const flowDetails =
 			'Take into account the following information for never tested results:\n<flowDetails>\n' +
 			yamlStringifyExceptKeys(sliceModules($flowStore.value.modules, upToIndex, idOrders), [
-				'lock',
-				'input_transforms'
+				'lock'
 			]) +
 			'</flowDetails>'
 
@@ -53,21 +59,22 @@
 The current step is ${selectedId}, you can find the details for the step and previous ones below:
 ${flowDetails}
 
-Determine for the inputs "${argNames.join(
+Determine for all the inputs "${argNames.join(
 				'", "'
-			)}", what to pass either from the previous results of the flow inputs. Here's a summary of the available data:
+			)}", what to pass either from the previous results of the flow inputs. 
+All possibles inputs either start with results. or flow_input. and are follow by the key of the input.			
+Here's a summary of the available data:
 <available>
 ${YAML.stringify(availableData)}</available>
 If none of the available results are appropriate, are already used or are more appropriate for other inputs, you can also imagine new flow_input properties which we will create programmatically based on what you provide.
 
 Reply with the most probable answer, do not explain or discuss.
 Use javascript object dot notation to access the properties.
-Return the input element directly: e.g. flow_input.property, results.a, results.a.property, flow_input.iter.value
 
 Your answer has to be in the following format (one line per input):
-input_name: expr`
-
-			console.log(user)
+{input_name1}: {expression1}
+{input_name2}: {expression2}
+...`
 
 			generatedContent = await getNonStreamingCompletion(
 				[
@@ -82,9 +89,18 @@ input_name: expr`
 			parsedInputs = generatedContent.split('\n').map((x) => x.split(': '))
 
 			const exprs = {}
+			newFlowInputs = []
 			for (const [key, value] of parsedInputs) {
 				if (argNames.includes(key)) {
 					exprs[key] = value
+					if (
+						pickableProperties &&
+						value.startsWith('flow_input.') &&
+						value.split('.')[1] &&
+						!(value.split('.')[1] in pickableProperties.flow_input)
+					) {
+						newFlowInputs.push(value.split('.')[1])
+					}
 				}
 			}
 			generatedExprs?.set(exprs)
@@ -99,6 +115,29 @@ input_name: expr`
 		}
 	}
 
+	function createFlowInputs() {
+		if (!newFlowInputs.length) {
+			return
+		}
+		const properties = {
+			...($flowStore.schema?.properties as Record<string, SchemaProperty> | undefined),
+			...newFlowInputs.reduce((acc, x) => {
+				acc[x] = (schema.properties ?? {})[x]
+				return acc
+			}, {})
+		}
+		const required = [
+			...(($flowStore.schema?.required as string[] | undefined) ?? []),
+			...newFlowInputs
+		]
+		$flowStore.schema = {
+			$schema: 'https://json-schema.org/draft/2020-12/schema',
+			properties,
+			required,
+			type: 'object'
+		}
+	}
+
 	function applyExprs() {
 		const argsUpdate = {}
 		for (const [key, value] of parsedInputs) {
@@ -111,19 +150,30 @@ input_name: expr`
 		}
 		exprsToSet?.set(argsUpdate)
 		generatedExprs?.set({})
+		if (newFlowInputs.length) {
+			openInputsModal = true
+		}
 	}
 
 	let out = true // hack to prevent regenerating answer when accepting the answer due to mouseenter on new icon
+	let openInputsModal = false
 </script>
 
 <div class="flex flex-row justify-end">
 	{#if $copilotInfo.exists_openai_resource_path && $stepInputCompletionEnabled}
+		<FlowCopilotInputsModal
+			on:confirmed={async () => {
+				createFlowInputs()
+			}}
+			bind:open={openInputsModal}
+			inputs={newFlowInputs}
+		/>
 		<Button
 			size="xs"
 			color="light"
 			btnClasses="text-violet-800 dark:text-violet-400"
 			on:mouseenter={(ev) => {
-				if (Object.keys($generatedExprs || {}).length === 0 && !loading && out) {
+				if (out) {
 					out = false
 					generateStepInputs()
 				}
@@ -134,12 +184,8 @@ input_name: expr`
 				generatedExprs?.set({})
 			}}
 			on:click={() => {
-				if (loading) {
-					abortController.abort()
-				} else if (Object.keys($generatedExprs || {}).length > 0) {
+				if (!loading && Object.keys($generatedExprs || {}).length > 0) {
 					applyExprs()
-				} else if (!loading) {
-					generateStepInputs()
 				}
 			}}
 			startIcon={{
@@ -188,9 +234,13 @@ input_name: expr`
 					</a>
 				{:else}
 					Enable step input completion in the{' '}
-					<a href="#user-settings" class="inline-flex flex-row items-center gap-1" on:click={() => {
-						close(null)
-					}}>
+					<a
+						href="#user-settings"
+						class="inline-flex flex-row items-center gap-1"
+						on:click={() => {
+							close(null)
+						}}
+					>
 						user settings
 					</a>
 				{/if}
diff --git a/frontend/src/lib/components/flows/FlowEditor.svelte b/frontend/src/lib/components/flows/FlowEditor.svelte
index f0aec6e3b5..471e5db665 100644
--- a/frontend/src/lib/components/flows/FlowEditor.svelte
+++ b/frontend/src/lib/components/flows/FlowEditor.svelte
@@ -47,7 +47,7 @@
 					</div>
 				</div>
 			{:else}
-				<FlowEditorPanel />
+				<FlowEditorPanel enableAi />
 			{/if}
 		</Pane>
 	</Splitpanes>
diff --git a/frontend/src/lib/components/flows/content/BranchPredicateEditor.svelte b/frontend/src/lib/components/flows/content/BranchPredicateEditor.svelte
index 1eeb5a5594..4e948b4c7a 100644
--- a/frontend/src/lib/components/flows/content/BranchPredicateEditor.svelte
+++ b/frontend/src/lib/components/flows/content/BranchPredicateEditor.svelte
@@ -8,6 +8,7 @@
 	import type { FlowEditorContext } from '../types'
 	import Button from '$lib/components/common/button/Button.svelte'
 	import { Pen } from 'lucide-svelte'
+	import PredicateGen from '$lib/components/copilot/PredicateGen.svelte'
 
 	export let branch: {
 		summary?: string
@@ -16,6 +17,7 @@
 	}
 	export let parentModule: FlowModule
 	export let previousModule: FlowModule | undefined
+	export let enableAi = false
 
 	const { previewArgs, flowStateStore, flowStore } =
 		getContext<FlowEditorContext>('FlowEditorContext')
@@ -55,8 +57,19 @@
 	</PropPickerWrapper>
 {:else}
 	<div class="flex justify-between gap-4 p-2">
-		<div><pre class="text-sm">{branch.expr}</pre></div><div>
+		<div><pre class="text-sm">{branch.expr}</pre></div><div
+			class="flex flex-row gap-2 items-center"
+		>
+			{#if enableAi}
+				<PredicateGen
+					on:setExpr={(e) => {
+						branch.expr = e.detail
+					}}
+					pickableProperties={stepPropPicker.pickableProperties}
+				/>
+			{/if}
 			<Button
+				size="xs"
 				startIcon={{ icon: Pen }}
 				variant="border"
 				on:click={() => (open = !open)}
diff --git a/frontend/src/lib/components/flows/content/FlowBranchOneWrapper.svelte b/frontend/src/lib/components/flows/content/FlowBranchOneWrapper.svelte
index d4b6d4bcf4..a4c842cc71 100644
--- a/frontend/src/lib/components/flows/content/FlowBranchOneWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowBranchOneWrapper.svelte
@@ -12,6 +12,7 @@
 	export let parentModule: FlowModule
 	export let previousModule: FlowModule | undefined
 	export let noEditor: boolean
+	export let enableAi = false
 </script>
 
 <div class="h-full flex flex-col">
@@ -21,7 +22,7 @@
 		</div>
 		<div class="overflow-hidden flex-grow">
 			<h3 class="p-2">Predicate expression</h3>
-			<BranchPredicateEditor {branch} {parentModule} {previousModule} />
+			<BranchPredicateEditor {branch} {parentModule} {previousModule} {enableAi} />
 		</div>
 	</FlowCard>
 </div>
diff --git a/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte b/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte
index dd15640dd4..21a2c124c4 100644
--- a/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte
@@ -19,6 +19,7 @@
 	export let flowModule: FlowModule
 	export let previousModule: FlowModule | undefined
 	export let noEditor: boolean
+	export let enableAi = false
 
 	let value = flowModule.value as BranchOne
 	$: value = flowModule.value as BranchOne
@@ -65,7 +66,12 @@
 										/>
 									</div>
 									<div class="w-full border">
-										<BranchPredicateEditor {branch} parentModule={flowModule} {previousModule} />
+										<BranchPredicateEditor
+											{branch}
+											parentModule={flowModule}
+											{previousModule}
+											{enableAi}
+										/>
 									</div>
 								</div>
 							{/each}
diff --git a/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte b/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
index 37b2a7113c..5cb06246d7 100644
--- a/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
+++ b/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
@@ -10,6 +10,7 @@
 	import type { FlowModule } from '$lib/gen'
 
 	export let noEditor = false
+	export let enableAi = false
 
 	const { selectedId, flowStore } = getContext<FlowEditorContext>('FlowEditorContext')
 
@@ -46,6 +47,7 @@
 					{noEditor}
 					bind:flowModule
 					previousModule={$flowStore.value.modules[index - 1]}
+					{enableAi}
 				/>
 			{/each}
 		{/key}
diff --git a/frontend/src/lib/components/flows/content/FlowLoop.svelte b/frontend/src/lib/components/flows/content/FlowLoop.svelte
index 821036d2cf..ddda211ab9 100644
--- a/frontend/src/lib/components/flows/content/FlowLoop.svelte
+++ b/frontend/src/lib/components/flows/content/FlowLoop.svelte
@@ -21,6 +21,7 @@
 	import type { Job } from '$lib/gen'
 	import FlowLoopIterationPreview from '$lib/components/FlowLoopIterationPreview.svelte'
 	import FlowModuleDeleteAfterUse from './FlowModuleDeleteAfterUse.svelte'
+	import IteratorGen from '$lib/components/copilot/IteratorGen.svelte'
 
 	const { previewArgs, flowStateStore, flowStore } =
 		getContext<FlowEditorContext>('FlowEditorContext')
@@ -29,6 +30,7 @@
 	export let parentModule: FlowModule | undefined
 	export let previousModule: FlowModule | undefined
 	export let noEditor: boolean
+	export let enableAi = false
 
 	let editor: SimpleEditor | undefined = undefined
 	let selected: string = 'early-stop'
@@ -47,6 +49,9 @@
 	let jobId: string | undefined = undefined
 	let job: Job | undefined = undefined
 
+	let iteratorFieldFocused = false
+	let iteratorGen: IteratorGen | undefined = undefined
+
 	$: previewIterationArgs = $flowStateStore[mod.id]?.previewArgs ?? {}
 </script>
 
@@ -121,14 +126,42 @@
 							/>
 						</div>
 					</div>
-					<div class="my-2 text-sm font-bold">
-						Iterator expression
-						<Tooltip documentationLink="https://www.windmill.dev/docs/flows/flow_loops">
-							List to iterate over.
-						</Tooltip>
+					<div class="my-2 flex flex-row gap-2 items-center">
+						<div class="text-sm font-bold">
+							Iterator expression
+							<Tooltip documentationLink="https://www.windmill.dev/docs/flows/flow_loops">
+								List to iterate over.
+							</Tooltip>
+						</div>
+						{#if enableAi}
+							<IteratorGen
+								bind:this={iteratorGen}
+								focused={iteratorFieldFocused}
+								arg={mod.value.iterator}
+								on:showExpr={(e) => {
+									editor?.setSuggestion(e.detail)
+								}}
+								on:setExpr={(e) => {
+									if (mod.value.type === 'forloopflow') {
+										mod.value.iterator = {
+											type: 'javascript',
+											expr: e.detail
+										}
+									}
+									editor?.setCode('')
+									editor?.insertAtCursor(e.detail)
+								}}
+								pickableProperties={stepPropPicker.pickableProperties}
+							/>
+						{/if}
 					</div>
 					{#if mod.value.iterator.type == 'javascript'}
-						<div class="border w-full" id="flow-editor-iterator-expression">
+						<!-- svelte-ignore a11y-no-static-element-interactions -->
+						<div
+							class="border w-full"
+							id="flow-editor-iterator-expression"
+							on:keyup={iteratorGen?.onKeyUp}
+						>
 							<PropPickerWrapper
 								notSelectable
 								pickableProperties={stepPropPicker.pickableProperties}
@@ -140,6 +173,13 @@
 							>
 								<SimpleEditor
 									bind:this={editor}
+									on:focus={() => {
+										iteratorFieldFocused = true
+									}}
+									on:blur={() => {
+										iteratorFieldFocused = false
+									}}
+									autofocus
 									lang="javascript"
 									bind:code={mod.value.iterator.expr}
 									class="small-editor"
diff --git a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
index 5c1fac8595..d66caa4f4c 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
@@ -20,6 +20,7 @@
 
 	export let flowModule: FlowModule
 	export let noEditor: boolean = false
+	export let enableAi = false
 
 	const { selectedId, schedule, flowStateStore } =
 		getContext<FlowEditorContext>('FlowEditorContext')
@@ -64,9 +65,9 @@
 
 {#if flowModule.id === $selectedId}
 	{#if flowModule.value.type === 'forloopflow'}
-		<FlowLoop {noEditor} bind:mod={flowModule} {parentModule} {previousModule} />
+		<FlowLoop {noEditor} bind:mod={flowModule} {parentModule} {previousModule} {enableAi} />
 	{:else if flowModule.value.type === 'branchone'}
-		<FlowBranchesOneWrapper {noEditor} {previousModule} bind:flowModule />
+		<FlowBranchesOneWrapper {noEditor} {previousModule} bind:flowModule {enableAi} />
 	{:else if flowModule.value.type === 'branchall'}
 		<FlowBranchesAllWrapper {noEditor} {previousModule} bind:flowModule />
 	{:else if flowModule.value.type === 'identity'}
@@ -151,6 +152,7 @@
 			bind:flowModule={submodule}
 			bind:parentModule={flowModule}
 			previousModule={flowModule.value.modules[index - 1]}
+			{enableAi}
 		/>
 	{/each}
 {:else if flowModule.value.type === 'branchone'}
@@ -165,18 +167,26 @@
 				bind:flowModule={submodule}
 				bind:parentModule={flowModule}
 				previousModule={flowModule.value.default[index - 1]}
+				{enableAi}
 			/>
 		{/each}
 	{/if}
 	{#each flowModule.value.branches as branch, branchIndex (branchIndex)}
 		{#if $selectedId === `${flowModule?.id}-branch-${branchIndex}`}
-			<FlowBranchOneWrapper {noEditor} bind:branch parentModule={flowModule} {previousModule} />
+			<FlowBranchOneWrapper
+				{noEditor}
+				bind:branch
+				parentModule={flowModule}
+				{previousModule}
+				{enableAi}
+			/>
 		{:else}
 			{#each branch.modules as submodule, index}
 				<svelte:self
 					bind:flowModule={submodule}
 					bind:parentModule={flowModule}
 					previousModule={flowModule.value.branches[branchIndex].modules[index - 1]}
+					{enableAi}
 				/>
 			{/each}
 		{/if}
@@ -191,6 +201,7 @@
 					bind:flowModule={submodule}
 					bind:parentModule={flowModule}
 					previousModule={flowModule.value.branches[branchIndex].modules[index - 1]}
+					{enableAi}
 				/>
 			{/each}
 		{/if}
diff --git a/frontend/src/lib/components/flows/flowStateUtils.ts b/frontend/src/lib/components/flows/flowStateUtils.ts
index 6904f0930a..e96a60a961 100644
--- a/frontend/src/lib/components/flows/flowStateUtils.ts
+++ b/frontend/src/lib/components/flows/flowStateUtils.ts
@@ -90,7 +90,7 @@ export async function createLoop(id: string): Promise<[FlowModule, FlowModuleSta
 		value: {
 			type: 'forloopflow',
 			modules: [],
-			iterator: { type: 'javascript', expr: '["dynamic or static array"]' },
+			iterator: { type: 'javascript', expr: '' },
 			skip_failures: true
 		}
 	}
diff --git a/frontend/src/lib/components/theme/nord.json b/frontend/src/lib/components/theme/nord.json
deleted file mode 100644
index ea03bbb84c..0000000000
--- a/frontend/src/lib/components/theme/nord.json
+++ /dev/null
@@ -1,93 +0,0 @@
-{
-  "base": "vs-dark",
-  "inherit": true,
-  "rules": [
-    {
-      "background": "2E3440",
-      "token": ""
-    },
-    {
-      "foreground": "616e88",
-      "token": "comment"
-    },
-    {
-      "foreground": "a3be8c",
-      "token": "string"
-    },
-    {
-      "foreground": "b48ead",
-      "token": "constant.numeric"
-    },
-    {
-      "foreground": "81a1c1",
-      "token": "constant.language"
-    },
-    {
-      "foreground": "81a1c1",
-      "token": "keyword"
-    },
-    {
-      "foreground": "81a1c1",
-      "token": "storage"
-    },
-    {
-      "foreground": "81a1c1",
-      "token": "storage.type"
-    },
-    {
-      "foreground": "8fbcbb",
-      "token": "entity.name.class"
-    },
-    {
-      "foreground": "8fbcbb",
-      "fontStyle": "  bold",
-      "token": "entity.other.inherited-class"
-    },
-    {
-      "foreground": "88c0d0",
-      "token": "entity.name.function"
-    },
-    {
-      "foreground": "81a1c1",
-      "token": "entity.name.tag"
-    },
-    {
-      "foreground": "8fbcbb",
-      "token": "entity.other.attribute-name"
-    },
-    {
-      "foreground": "88c0d0",
-      "token": "support.function"
-    },
-    {
-      "foreground": "f8f8f0",
-      "background": "f92672",
-      "token": "invalid"
-    },
-    {
-      "foreground": "f8f8f0",
-      "background": "ae81ff",
-      "token": "invalid.deprecated"
-    },
-    {
-      "foreground": "b48ead",
-      "token": "constant.color.other.rgb-value"
-    },
-    {
-      "foreground": "ebcb8b",
-      "token": "constant.character.escape"
-    },
-    {
-      "foreground": "8fbcbb",
-      "token": "variable.other.constant"
-    }
-  ],
-  "colors": {
-    "editor.foreground": "#D8DEE9",
-    "editor.background": "#2E3440",
-    "editor.selectionBackground": "#434C5ECC",
-    "editor.lineHighlightBackground": "#3B4252",
-    "editorCursor.foreground": "#D8DEE9",
-    "editorWhitespace.foreground": "#434C5ECC"
-  }
-}
\ No newline at end of file